# pairwisedistance_cuda.py
import torch
from torch.utils.cpp_extension import load_inline

from pairwisedistance_torch import BATCH_SIZE, FEATURE_DIM, P_NORM # 导入维度常量

EPS = 1e-6 

class ModelNew(torch.nn.Module):
 
    def __init__(self):
        super().__init__()
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
  
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor pdist_forward_cuda(torch::Tensor x1, torch::Tensor x2, float eps_val);
        """

        cuda_source = """
        #include <cuda_runtime.h>
        #include <cmath>
        #include <float.h>

        #define BLOCK_SIZE 256
        #define WARP_SIZE 32
        #define FEATURE_DIM_VAL {FEATURE_DIM} 

        __device__ __forceinline__ float warp_reduce_sum(float val) {{
            for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {{
                val += __shfl_down_sync(0xffffffff, val, offset);
            }}
            return val;
        }}

        __global__ void pdist_fused_kernel(
            const float* __restrict__ x1,
            const float* __restrict__ x2,
            float* __restrict__ y_out, 
            int num_batches,
            int feature_dim,
            float eps_val 
        ) {{
     
            __shared__ float s_data[BLOCK_SIZE]; 

            int batch_idx = blockIdx.x;
            if (batch_idx >= num_batches) return;

            const float* a_ptr = x1 + batch_idx * feature_dim;
            const float* b_ptr = x2 + batch_idx * feature_dim;

            float thread_diff_sq_sum = 0.0f;
            
            for (int i = threadIdx.x; i < feature_dim; i += blockDim.x) {{
                float a_val = a_ptr[i];
                float b_val = b_ptr[i];
                
                float diff = a_val - b_val;
                // L2 Norm: |diff|^2
                thread_diff_sq_sum += diff * diff;
            }}
  
            thread_diff_sq_sum = warp_reduce_sum(thread_diff_sq_sum);

            int warp_id = threadIdx.x / WARP_SIZE;
            int lane_id = threadIdx.x % WARP_SIZE;
            int num_warps = blockDim.x / WARP_SIZE;

            if (lane_id == 0) {{
                s_data[warp_id] = thread_diff_sq_sum;
            }}
            __syncthreads();

            if (threadIdx.x < num_warps) {
                float total_diff_sq_sum = warp_reduce_sum(s_data[threadIdx.x]);

                if (threadIdx.x == 0) {
                    float result = sqrtf(total_diff_sq_sum + eps_val); 
                    y_out[batch_idx] = result; 
                }
            }
        }}

        torch::Tensor pdist_forward_cuda(torch::Tensor x1, torch::Tensor x2, float eps_val) {
            TORCH_CHECK(x1.is_cuda(), "Input must be a CUDA tensor");
            x1 = x1.contiguous();
            x2 = x2.contiguous();

            int num_batches = x1.size(0);
            int feature_dim = x1.size(1);
      
            auto output = torch::empty({num_batches}, x1.options()); 

            const int block_size = BLOCK_SIZE;
            const int grid_size = num_batches; 

            size_t shared_mem_size = (block_size / WARP_SIZE) * sizeof(float);

            pdist_fused_kernel<<<grid_size, block_size, shared_mem_size>>>(
                x1.data_ptr<float>(),
                x2.data_ptr<float>(),
                output.data_ptr<float>(),
                num_batches,
                feature_dim,
                eps_val 
            );
            
            return output;
        }
        """
        
        self.pdist_op = load_inline(
            name="pdist_fused_op_safest",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["pdist_forward_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math"],
            verbose=True
        )

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        return self.pdist_op.pdist_forward_cuda(x1, x2, EPS)