import torch
from torch.utils.cpp_extension import load_inline

# 高效的成对距离 CUDA 源码
pairwise_distance_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

// Kernel to compute pairwise Euclidean distance
// Each thread computes one element of the output distance matrix.
__global__ void pairwise_distance_kernel(
    const float* __restrict__ x, 
    const float* __restrict__ y, 
    float* __restrict__ dist,
    int N, int M, int D
) {
    // 2D thread mapping
    int row = blockIdx.y * blockDim.y + threadIdx.y; // Maps to N
    int col = blockIdx.x * blockDim.x + threadIdx.x; // Maps to M

    if (row < N && col < M) {
        // Pointer to the start of the vectors for this thread's computation
        const float* x_vec = x + row * D;
        const float* y_vec = y + col * D;
        
        float sum_sq = 0.0f;
        // Compute the sum of squared differences
        for (int k = 0; k < D; ++k) {
            float diff = x_vec[k] - y_vec[k];
            sum_sq += diff * diff;
        }
        
        // Store the final distance
        dist[row * M + col] = sqrtf(sum_sq);
    }
}

torch::Tensor pairwise_distance_cuda(torch::Tensor x, torch::Tensor y) {
    TORCH_CHECK(x.is_cuda() && x.is_contiguous());
    TORCH_CHECK(y.is_cuda() && y.is_contiguous());
    TORCH_CHECK(x.dim() == 2 && y.dim() == 2);
    TORCH_CHECK(x.size(1) == y.size(1), "Vector dimensions must match");

    int N = x.size(0);
    int M = y.size(0);
    int D = x.size(1);
    
    auto dist = torch::empty({N, M}, x.options());

    // Use a 2D grid and 2D blocks
    dim3 block_size(16, 16); // A 16x16 thread block is a good starting point
    dim3 grid_size((M + block_size.x - 1) / block_size.x, 
                   (N + block_size.y - 1) / block_size.y);

    pairwise_distance_kernel<<<grid_size, block_size>>>(
        x.data_ptr<float>(),
        y.data_ptr<float>(),
        dist.data_ptr<float>(),
        N, M, D
    );
    
    return dist;
}
"""

pairwise_distance_cpp_source = """
torch::Tensor pairwise_distance_cuda(torch::Tensor x, torch::Tensor y);
"""

# 编译模块
pairwise_distance_module = load_inline(
    name="pairwise_distance",
    cpp_sources=pairwise_distance_cpp_source,
    cuda_sources=pairwise_distance_source,
    functions=["pairwise_distance_cuda"],
    verbose=True,
    extra_cuda_cflags=["-O3", "--use_fast_math"]
)

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

    def forward(self, x, y):
        if not x.is_contiguous(): x = x.contiguous()
        if not y.is_contiguous(): y = y.contiguous()
        return self.pairwise_distance.pairwise_distance_cuda(x, y)
