import torch
from torch.utils.cpp_extension import load_inline

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

// Warp级归约版本 - 更高效的并行归约，并融合了Flatten操作
__global__ void dice_from_2d_kernel_warp_reduction(
    const float* __restrict__ pred,
    const float* __restrict__ target,
    float* __restrict__ intersection,
    float* __restrict__ pred_sum,
    float* __restrict__ target_sum,
    int size
) {
    int tid = threadIdx.x;
    int warp_id = (blockIdx.x * blockDim.x + tid) / 32;
    int lane_id = tid % 32;
    
    // 每个warp处理一块数据
    int elements_per_warp = 1024;  // 每个warp处理1K元素
    int warp_start = warp_id * elements_per_warp;
    int warp_end = min(warp_start + elements_per_warp, size);
    
    // 寄存器累加
    float reg_intersection = 0.0f;
    float reg_pred_sum = 0.0f;
    float reg_target_sum = 0.0f;
    
    // 处理数据 - 这里的size是总元素数，实现了隐式flatten
    for (int i = warp_start + lane_id; i < warp_end; i += 32) {
        float p = pred[i];
        float t = target[i];
        reg_intersection += p * t;
        reg_pred_sum += p;
        reg_target_sum += t;
    }
    
    // Warp内归约 - 使用warp shuffle指令
    #pragma unroll
    for (int offset = 16; offset > 0; offset /= 2) {
        reg_intersection += __shfl_down_sync(0xffffffff, reg_intersection, offset);
        reg_pred_sum += __shfl_down_sync(0xffffffff, reg_pred_sum, offset);
        reg_target_sum += __shfl_down_sync(0xffffffff, reg_target_sum, offset);
    }
    
    // warp内第一个线程写入结果
    if (lane_id == 0) {
        atomicAdd(intersection, reg_intersection);
        atomicAdd(pred_sum, reg_pred_sum);
        atomicAdd(target_sum, reg_target_sum);
    }
}

torch::Tensor dice_from_2d_cuda(torch::Tensor pred, torch::Tensor target) {
    // --- 关键修改：直接计算总元素数，实现隐式flatten ---
    auto size = pred.numel();
    
    // 创建中间结果张量
    auto intersection = torch::zeros(1, pred.options());
    auto pred_sum_tensor = torch::zeros(1, pred.options());
    auto target_sum_tensor = torch::zeros(1, pred.options());
    
    // 使用Warp级归约版本
    const int block_size = 256;
    int num_warps_needed = (size + 1024 - 1) / 1024;
    int num_blocks = (num_warps_needed * 32 + block_size - 1) / block_size;
    
    dice_from_2d_kernel_warp_reduction<<<num_blocks, block_size>>>(
        pred.data_ptr<float>(), 
        target.data_ptr<float>(),
        intersection.data_ptr<float>(),
        pred_sum_tensor.data_ptr<float>(),
        target_sum_tensor.data_ptr<float>(),
        size
    );
    
    // 计算最终DiceLoss
    float epsilon = 1e-6f;
    auto dice_score = (2.0f * intersection) / (pred_sum_tensor + target_sum_tensor + epsilon);
    auto dice_loss = 1.0f - dice_score;
    
    return dice_loss;
}
"""

dice_from_2d_cpp_source = """
torch::Tensor dice_from_2d_cuda(torch::Tensor pred, torch::Tensor target);
"""

# 编译CUDA代码
dice_from_2d = load_inline(
    name="dice_from_2d",
    cpp_sources=dice_from_2d_cpp_source,
    cuda_sources=dice_from_2d_source,
    functions=["dice_from_2d_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.dice_from_2d = dice_from_2d

    def forward(self, pred, target):
        return self.dice_from_2d.dice_from_2d_cuda(pred, target)
