import torch
from torch.utils.cpp_extension import load_inline

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

// 修复后的融合sigmoid + BCE kernel
__global__ void bce_sigmoid_fused_kernel(
    const float* __restrict__ logits,
    const float* __restrict__ targets,
    float* __restrict__ loss,
    int size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 局部累加器
    float local_sum = 0.0f;
    
    // 每个线程处理多个元素
    int stride = blockDim.x * gridDim.x;
    for (int i = idx; i < size; i += stride) {
        float logit = logits[i];
        float target = targets[i];
        
        // 修复：正确的BCE with logits公式
        // BCE = log(1 + exp(-logit)) + (1-target)*logit 当 logit >= 0
        // BCE = -logit + log(1 + exp(logit)) + target*logit 当 logit < 0
        float bce_loss;
        if (logit >= 0.0f) {
            // 对于正logit：使用数值稳定的计算
            float exp_neg_logit = expf(-logit);
            bce_loss = log1pf(exp_neg_logit) + (1.0f - target) * logit;
        } else {
            // 对于负logit：使用数值稳定的计算
            float exp_logit = expf(logit);
            bce_loss = -logit + log1pf(exp_logit) + target * logit;
        }
        
        local_sum += bce_loss;
    }
    
    // 使用共享内存进行块内归约
    extern __shared__ float shared_mem[];
    shared_mem[threadIdx.x] = local_sum;
    
    __syncthreads();
    
    // 块内归约
    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
        if (threadIdx.x < stride) {
            shared_mem[threadIdx.x] += shared_mem[threadIdx.x + stride];
        }
        __syncthreads();
    }
    
    // 第一个线程将结果写入全局内存
    if (threadIdx.x == 0) {
        atomicAdd(loss, shared_mem[0]);
    }
}

torch::Tensor bce_sigmoid_cuda(torch::Tensor logits, torch::Tensor targets) {
    TORCH_CHECK(logits.scalar_type() == torch::kFloat32, "Input must be float32");
    TORCH_CHECK(targets.scalar_type() == torch::kFloat32, "Target must be float32");
    TORCH_CHECK(logits.sizes() == targets.sizes(), "Input and target must have same shape");
    
    auto logits_contig = logits.contiguous();
    auto targets_contig = targets.contiguous();
    int size = logits_contig.numel();
    
    // 创建输出tensor，初始化为0
    auto loss = torch::zeros({1}, logits.options());
    
    // 优化的kernel配置
    const int block_size = 256;
    int num_blocks = min(65535, (size + block_size - 1) / block_size);
    
    size_t shared_mem = block_size * sizeof(float);
    
    // 启动修复后的kernel
    bce_sigmoid_fused_kernel<<<num_blocks, block_size, shared_mem>>>(
        logits_contig.data_ptr<float>(),
        targets_contig.data_ptr<float>(),
        loss.data_ptr<float>(),
        size
    );
    
    // 检查CUDA错误
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        AT_ERROR("CUDA kernel failed: ", cudaGetErrorString(err));
    }
    
    return loss;
}
"""

bce_sigmoid_cpp_source = """
torch::Tensor bce_sigmoid_cuda(torch::Tensor logits, torch::Tensor targets);
"""

# 编译修复后的CUDA代码
bce_sigmoid = load_inline(
    name="bce_sigmoid_fixed",
    cpp_sources=bce_sigmoid_cpp_source,
    cuda_sources=bce_sigmoid_source,
    functions=["bce_sigmoid_cuda"],
    extra_cuda_cflags=[
        "-O3", 
        "--use_fast_math",
        "-gencode=arch=compute_80,code=sm_80"
    ],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.bce_sigmoid = bce_sigmoid

    def forward(self, logits, targets):
        return self.bce_sigmoid.bce_sigmoid_cuda(logits, targets)
