import torch
from torch.utils.cpp_extension import load_inline

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

// 向量化DiceLoss内核 - 8元素并行处理
__global__ void dice_kernel_vectorized(
    const float* __restrict__ pred,
    const float* __restrict__ target,
    float* __restrict__ intersection,
    float* __restrict__ pred_sum,
    float* __restrict__ target_sum,
    int size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 8元素向量化处理
    int vec_idx = idx * 8;
    
    // 局部累加器 - 避免原子操作
    float local_intersection = 0.0f;
    float local_pred_sum = 0.0f;
    float local_target_sum = 0.0f;
    
    if (vec_idx + 7 < size) {
        // 加载8个元素（2个float4）
        float4 pred_vec1 = *reinterpret_cast<const float4*>(&pred[vec_idx]);
        float4 pred_vec2 = *reinterpret_cast<const float4*>(&pred[vec_idx + 4]);
        float4 target_vec1 = *reinterpret_cast<const float4*>(&target[vec_idx]);
        float4 target_vec2 = *reinterpret_cast<const float4*>(&target[vec_idx + 4]);
        
        // 计算前4个元素
        local_intersection += pred_vec1.x * target_vec1.x;
        local_intersection += pred_vec1.y * target_vec1.y;
        local_intersection += pred_vec1.z * target_vec1.z;
        local_intersection += pred_vec1.w * target_vec1.w;
        
        local_pred_sum += pred_vec1.x + pred_vec1.y + pred_vec1.z + pred_vec1.w;
        local_target_sum += target_vec1.x + target_vec1.y + target_vec1.z + target_vec1.w;
        
        // 计算后4个元素
        local_intersection += pred_vec2.x * target_vec2.x;
        local_intersection += pred_vec2.y * target_vec2.y;
        local_intersection += pred_vec2.z * target_vec2.z;
        local_intersection += pred_vec2.w * target_vec2.w;
        
        local_pred_sum += pred_vec2.x + pred_vec2.y + pred_vec2.z + pred_vec2.w;
        local_target_sum += target_vec2.x + target_vec2.y + target_vec2.z + target_vec2.w;
    } else {
        // 边界处理
        for (int i = vec_idx; i < size && i < vec_idx + 8; i++) {
            float p = pred[i];
            float t = target[i];
            local_intersection += p * t;
            local_pred_sum += p;
            local_target_sum += t;
        }
    }
    
    // 使用共享内存进行块内归约，减少原子操作
    extern __shared__ float shared_mem[];
    float* shared_intersection = shared_mem;
    float* shared_pred_sum = shared_mem + blockDim.x;
    float* shared_target_sum = shared_mem + blockDim.x * 2;
    
    // 写入共享内存
    shared_intersection[threadIdx.x] = local_intersection;
    shared_pred_sum[threadIdx.x] = local_pred_sum;
    shared_target_sum[threadIdx.x] = local_target_sum;
    
    __syncthreads();
    
    // 块内归约
    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
        if (threadIdx.x < stride) {
            shared_intersection[threadIdx.x] += shared_intersection[threadIdx.x + stride];
            shared_pred_sum[threadIdx.x] += shared_pred_sum[threadIdx.x + stride];
            shared_target_sum[threadIdx.x] += shared_target_sum[threadIdx.x + stride];
        }
        __syncthreads();
    }
    
    // 只有第一个线程做原子操作
    if (threadIdx.x == 0) {
        atomicAdd(intersection, shared_intersection[0]);
        atomicAdd(pred_sum, shared_pred_sum[0]);
        atomicAdd(target_sum, shared_target_sum[0]);
    }
}

// Warp级归约版本 - 更高效的并行归约
__global__ void dice_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;
    
    // 处理数据
    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);
    }
}

// 大块处理版本 - 减少kernel启动开销
__global__ void dice_kernel_large_blocks(
    const float* __restrict__ pred,
    const float* __restrict__ target,
    float* __restrict__ intersection,
    float* __restrict__ pred_sum,
    float* __restrict__ target_sum,
    int size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 每个线程处理更多元素
    int elements_per_thread = 16;
    int start_idx = idx * elements_per_thread;
    
    // 局部累加
    float local_intersection = 0.0f;
    float local_pred_sum = 0.0f;
    float local_target_sum = 0.0f;
    
    for (int i = 0; i < elements_per_thread; i++) {
        int cur_idx = start_idx + i;
        if (cur_idx < size) {
            float p = pred[cur_idx];
            float t = target[cur_idx];
            local_intersection += p * t;
            local_pred_sum += p;
            local_target_sum += t;
        }
    }
    
    // 只有部分线程做原子操作
    if (threadIdx.x % 32 == 0) {  // 每32个线程只有一个做原子操作
        atomicAdd(intersection, local_intersection);
        atomicAdd(pred_sum, local_pred_sum);
        atomicAdd(target_sum, local_target_sum);
    }
}

torch::Tensor dice_cuda(torch::Tensor pred, torch::Tensor target) {
    auto size = pred.numel();
    
    // 创建中间结果张量
    auto intersection = torch::zeros(1, pred.options());
    auto pred_sum = torch::zeros(1, pred.options());
    auto target_sum = torch::zeros(1, pred.options());
    
    // 使用向量化版本
    const int block_size = 256;
    int num_blocks = (size + block_size * 8 - 1) / (block_size * 8);
    size_t shared_mem = block_size * 3 * sizeof(float);
    
    dice_kernel_vectorized<<<num_blocks, block_size, shared_mem>>>(
        pred.data_ptr<float>(), 
        target.data_ptr<float>(),
        intersection.data_ptr<float>(),
        pred_sum.data_ptr<float>(),
        target_sum.data_ptr<float>(),
        size
    );
    
    // 计算最终DiceLoss
    float epsilon = 1e-6f;
    auto dice_score = (2.0f * intersection) / (pred_sum + target_sum + epsilon);
    auto dice_loss = 1.0f - dice_score;
    
    return dice_loss;
}
"""

dice_cpp_source = """
torch::Tensor dice_cuda(torch::Tensor pred, torch::Tensor target);
"""

# Compile the inline CUDA code
dice = load_inline(
    name="dice",
    cpp_sources=dice_cpp_source,
    cuda_sources=dice_source,
    functions=["dice_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 = dice  # The module containing the kernel

    def forward(self, pred, target):
        return self.dice.dice_cuda(pred, target)
