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

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

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 cossim_forward_cuda(torch::Tensor x1, torch::Tensor x2);
        """

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

        #define BLOCK_SIZE 256
        #define WARP_SIZE 32
        #define VEC_SIZE 4 

        __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 cossim_fused_vectorized_kernel(
            const float* __restrict__ x1,
            const float* __restrict__ x2,
            float* __restrict__ y_out, 
            int num_batches,
            int feature_dim
        ) {{
            __shared__ float s_data[BLOCK_SIZE * 3]; 

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

            const float4* a4_ptr = reinterpret_cast<const float4*>(x1 + batch_idx * feature_dim);
            const float4* b4_ptr = reinterpret_cast<const float4*>(x2 + batch_idx * feature_dim);

            int n_work_items = feature_dim / VEC_SIZE;

            float thread_dot = 0.0f;
            float thread_norm_sq_a = 0.0f;
            float thread_norm_sq_b = 0.0f;
            
            for (int i = threadIdx.x; i < n_work_items; i += blockDim.x) {{

                float4 a_vec = a4_ptr[i];
                float4 b_vec = b4_ptr[i];

                thread_dot += a_vec.x * b_vec.x + a_vec.y * b_vec.y + a_vec.z * b_vec.z + a_vec.w * b_vec.w;
                thread_norm_sq_a += a_vec.x * a_vec.x + a_vec.y * a_vec.y + a_vec.z * a_vec.z + a_vec.w * a_vec.w;
                thread_norm_sq_b += b_vec.x * b_vec.x + b_vec.y * b_vec.y + b_vec.z * b_vec.z + b_vec.w * b_vec.w;
            }}

            thread_dot = warp_reduce_sum(thread_dot);
            thread_norm_sq_a = warp_reduce_sum(thread_norm_sq_a);
            thread_norm_sq_b = warp_reduce_sum(thread_norm_sq_b);

            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_dot;
                s_data[warp_id + num_warps] = thread_norm_sq_a;
                s_data[warp_id + 2*num_warps] = thread_norm_sq_b;
            }}
            __syncthreads();

            if (threadIdx.x < num_warps) {
                float total_dot = warp_reduce_sum(s_data[threadIdx.x]);
                float total_norm_sq_a = warp_reduce_sum(s_data[threadIdx.x + num_warps]);
                float total_norm_sq_b = warp_reduce_sum(s_data[threadIdx.x + 2*num_warps]);

                if (threadIdx.x == 0) {
                    float norm_a = sqrtf(total_norm_sq_a);
                    float norm_b = sqrtf(total_norm_sq_b);
                    
                    float result = total_dot / (norm_a * norm_b + 1e-8f); // 加 eps 防止除以零
                    
                    y_out[batch_idx] = result; 
                }
            }
        }}

        torch::Tensor cossim_forward_cuda(torch::Tensor x1, torch::Tensor x2) {
            TORCH_CHECK(x1.is_cuda(), "Input must be a CUDA tensor");

            TORCH_CHECK(x1.size(1) % VEC_SIZE == 0, "Feature dimension must be divisible by 4 for float4 vectorization");
            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) * 3 * sizeof(float);

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

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        return self.cossim_op.cossim_forward_cuda(x1, x2)