import torch
from torch.utils.cpp_extension import load_inline

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

// Warp级优化版本 - 最佳性能方案
__global__ void euclidean_kernel_warp(
    const float* __restrict__ x,
    const float* __restrict__ y,
    float* __restrict__ distances,
    int batch_size,
    int feature_dim
) {
    int sample_idx = blockIdx.x;
    if (sample_idx >= batch_size) return;
    
    int tid = threadIdx.x;
    int warp_id = tid / 32;
    int lane_id = tid % 32;
    
    // 预计算指针
    const float* x_ptr = x + sample_idx * feature_dim;
    const float* y_ptr = y + sample_idx * feature_dim;
    
    // 每个warp处理一部分数据
    int elements_per_warp = (feature_dim + 8 - 1) / 8; // 最多8个warps
    int start_dim = warp_id * elements_per_warp;
    int end_dim = min(start_dim + elements_per_warp, feature_dim);
    
    float warp_sum = 0.0f;
    
    // 每个线程处理4个元素
    for (int dim = start_dim + lane_id * 4; dim < end_dim; dim += 32 * 4) {
        if (dim + 3 < end_dim) {
            float4 x_val = *reinterpret_cast<const float4*>(&x_ptr[dim]);
            float4 y_val = *reinterpret_cast<const float4*>(&y_ptr[dim]);
            
            float diff0 = x_val.x - y_val.x;
            float diff1 = x_val.y - y_val.y;
            float diff2 = x_val.z - y_val.z;
            float diff3 = x_val.w - y_val.w;
            
            warp_sum += diff0*diff0 + diff1*diff1 + diff2*diff2 + diff3*diff3;
        } else {
            for (int j = dim; j < end_dim; j++) {
                float diff = x_ptr[j] - y_ptr[j];
                warp_sum += diff * diff;
            }
        }
    }
    
    // Warp级归约
    for (int offset = 16; offset > 0; offset /= 2) {
        warp_sum += __shfl_down_sync(0xffffffff, warp_sum, offset);
    }
    
    // 使用共享内存进行跨warp归约
    extern __shared__ float shared_warp_sums[];
    if (lane_id == 0) {
        shared_warp_sums[warp_id] = warp_sum;
    }
    __syncthreads();
    
    // 第一个线程进行最终归约
    if (tid == 0) {
        float total_sum = 0.0f;
        int num_warps = blockDim.x / 32;
        for (int i = 0; i < num_warps; i++) {
            total_sum += shared_warp_sums[i];
        }
        distances[sample_idx] = sqrtf(total_sum);
    }
}

torch::Tensor euclidean_cuda(
    torch::Tensor x,
    torch::Tensor y
) {
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "X must be float32");
    TORCH_CHECK(y.scalar_type() == torch::kFloat32, "Y must be float32");
    
    auto x_contig = x.contiguous();
    auto y_contig = y.contiguous();
    
    int batch_size = x_contig.size(0);
    int feature_dim = x_contig.size(1);
    
    auto distances = torch::zeros({batch_size}, x.options());
    
    // Warp级优化版本 - 最佳性能
    const int block_size = 256; // 8个warps
    size_t shared_mem = 8 * sizeof(float); // 8个warp的结果
    
    euclidean_kernel_warp<<<batch_size, block_size, shared_mem>>>(
        x_contig.data_ptr<float>(),
        y_contig.data_ptr<float>(),
        distances.data_ptr<float>(),
        batch_size,
        feature_dim
    );
    
    return distances;
}
"""

euclidean_cpp_source = """
torch::Tensor euclidean_cuda(torch::Tensor x, torch::Tensor y);
"""

# 编译CUDA代码
euclidean = load_inline(
    name="euclidean",
    cpp_sources=euclidean_cpp_source,
    cuda_sources=euclidean_source,
    functions=["euclidean_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.euclidean = euclidean

    def forward(self, x, y):
        return self.euclidean.euclidean_cuda(x, y)
