import torch
from torch.utils.cpp_extension import load_inline

sigmoid_derivative_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

// 向量化版本：处理4个元素
__global__ void sigmoid_derivative_vectorized4_kernel(
    const float* __restrict__ x,
    float* __restrict__ y,
    int size
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    
    // 每次处理4个元素
    int vec_size = size / 4;
    
    for (int i = tid; i < vec_size; i += stride) {
        int base_idx = i * 4;
        
        // 加载4个元素
        float x0 = x[base_idx + 0];
        float x1 = x[base_idx + 1];
        float x2 = x[base_idx + 2];
        float x3 = x[base_idx + 3];
        
        // 计算4个sigmoid derivative
        float sigmoid0, sigmoid1, sigmoid2, sigmoid3;
        
        // 批量计算sigmoid
        if (x0 >= 0.0f) {
            float exp_neg_x = expf(-x0);
            sigmoid0 = 1.0f / (1.0f + exp_neg_x);
        } else {
            float exp_x = expf(x0);
            sigmoid0 = exp_x / (1.0f + exp_x);
        }
        
        if (x1 >= 0.0f) {
            float exp_neg_x = expf(-x1);
            sigmoid1 = 1.0f / (1.0f + exp_neg_x);
        } else {
            float exp_x = expf(x1);
            sigmoid1 = exp_x / (1.0f + exp_x);
        }
        
        if (x2 >= 0.0f) {
            float exp_neg_x = expf(-x2);
            sigmoid2 = 1.0f / (1.0f + exp_neg_x);
        } else {
            float exp_x = expf(x2);
            sigmoid2 = exp_x / (1.0f + exp_x);
        }
        
        if (x3 >= 0.0f) {
            float exp_neg_x = expf(-x3);
            sigmoid3 = 1.0f / (1.0f + exp_neg_x);
        } else {
            float exp_x = expf(x3);
            sigmoid3 = exp_x / (1.0f + exp_x);
        }
        
        // 存储结果
        y[base_idx + 0] = sigmoid0 * (1.0f - sigmoid0);
        y[base_idx + 1] = sigmoid1 * (1.0f - sigmoid1);
        y[base_idx + 2] = sigmoid2 * (1.0f - sigmoid2);
        y[base_idx + 3] = sigmoid3 * (1.0f - sigmoid3);
    }
    
    // 处理剩余元素
    int remainder = size % 4;
    if (tid == 0 && remainder > 0) {
        int start_idx = vec_size * 4;
        for (int i = start_idx; i < size; i++) {
            float xi = x[i];
            float sigmoid_x;
            if (xi >= 0.0f) {
                float exp_neg_x = expf(-xi);
                sigmoid_x = 1.0f / (1.0f + exp_neg_x);
            } else {
                float exp_x = expf(xi);
                sigmoid_x = exp_x / (1.0f + exp_x);
            }
            y[i] = sigmoid_x * (1.0f - sigmoid_x);
        }
    }
}

torch::Tensor sigmoid_derivative_cuda(
    torch::Tensor x,
    std::string mode = "vectorized"
) {
    auto x_contig = x.contiguous();
    auto y = torch::empty_like(x_contig);
    int size = x_contig.numel();
    
    if (mode == "vectorized") {
        // 向量化版本
        const int block_size = 256;
        int num_blocks = (size / 4 + block_size - 1) / block_size;
        int grid_size = std::min(num_blocks, 65535);
        
        sigmoid_derivative_vectorized4_kernel<<<grid_size, block_size>>>(
            x_contig.data_ptr<float>(),
            y.data_ptr<float>(),
            size
        );
    }
    
    return y;
}
"""

sigmoid_derivative_cpp_source = """
torch::Tensor sigmoid_derivative_cuda(torch::Tensor x, std::string mode);
"""

# 编译CUDA代码
sigmoid_derivative = load_inline(
    name="sigmoid_derivative_opt",
    cpp_sources=sigmoid_derivative_cpp_source,
    cuda_sources=sigmoid_derivative_source,
    functions=["sigmoid_derivative_cuda"],
    extra_cuda_cflags=[
        "-O3", 
        "--use_fast_math",
        "-std=c++17",
        "-maxrregcount=64"  # 增加寄存器使用
    ],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, mode="vectorized"):
        super(ModelNew, self).__init__()
        self.mode = mode
        self.sigmoid_derivative = sigmoid_derivative

    def forward(self, x):
        return self.sigmoid_derivative.sigmoid_derivative_cuda(x, self.mode)

# 测试数据配置
batch_size = 16
channels = 64
height = 32
width = 32

def get_inputs():
    x = torch.randn(batch_size, channels, height, width).cuda() * 2.0
    return [x]

def get_init_inputs():
    return []
