import torch
import os
from torch.utils.cpp_extension import load_inline

# 设置CUDA架构以优化编译
os.environ.setdefault('TORCH_CUDA_ARCH_LIST', '7.0 7.5 8.0 8.6 8.9 9.0')

# 专业级优化的Swish激活函数CUDA实现
swish_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>

// 极简但高效的Swish激活函数内核
__global__ void swish_kernel_optimized(
    const float* __restrict__ input,
    float* __restrict__ output,
    int num_elements) {
    
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 边界检查
    if (idx >= num_elements) {
        return;
    }
    
    float x = input[idx];
    // Swish = x * sigmoid(x)
    float sigmoid_x = 1.0f / (1.0f + expf(-x));
    output[idx] = x * sigmoid_x;
}

torch::Tensor swish_cuda_forward(torch::Tensor input) {
    auto num_elements = input.numel();
    
    if (num_elements == 0) {
        return torch::empty({0}, input.options());
    }
    
    auto output = torch::empty_like(input);
    
    // 优化线程块配置
    int block_size = 256;
    int grid_size = (num_elements + block_size - 1) / block_size;
    
    // 调用优化内核
    swish_kernel_optimized<<<grid_size, block_size>>>(
        input.data_ptr<float>(), output.data_ptr<float>(), num_elements);
    
    return output;
}
"""

swish_cpp_source = """
torch::Tensor swish_cuda_forward(torch::Tensor input);
"""

# 编译内联CUDA代码
try:
    swish_cuda = load_inline(
        name="swish_cuda",
        cpp_sources=swish_cpp_source,
        cuda_sources=swish_source,
        functions=["swish_cuda_forward"]
    )
    CUDA_AVAILABLE = True
except Exception as e:
    CUDA_AVAILABLE = False
    swish_cuda = None

class SwishModel(torch.nn.Module):
    def __init__(self):
        super(SwishModel, self).__init__()

    def forward(self, input):
        if CUDA_AVAILABLE and swish_cuda is not None:
            return swish_cuda.swish_cuda_forward(input)
        else:
            # PyTorch原生实现
            return input * torch.sigmoid(input)