# example_cudacode.py

import torch
from torch.utils.cpp_extension import load_inline

# Triplet Margin Loss的自定义CUDA实现
triplet_loss_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>

// 定义共享内存和块的大小
#define BLOCK_SIZE 256

// 计算两个向量之间的L2距离的平方
__device__ float squared_l2_distance(const float* v1, const float* v2, int dim, float* sdata) {
    float my_sum = 0.0f;
    // 每个线程计算一部分元素的平方差之和
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        float diff = v1[i] - v2[i];
        my_sum += diff * diff;
    }
    sdata[threadIdx.x] = my_sum;
    __syncthreads(); // 确保所有线程都完成了它们的初始求和

    // 使用共享内存执行并行规约
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            sdata[threadIdx.x] += sdata[threadIdx.x + s];
        }
        __syncthreads();
    }

    // 最终的平方和在sdata[0]中
    return sdata[0];
}

__global__ void triplet_loss_kernel(
    const float* anchor, const float* positive, const float* negative,
    float* loss, int batch_size, int dim, float margin) {

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

    __shared__ float sdata[BLOCK_SIZE];

    // 计算当前三元组的向量指针
    const float* anchor_ptr = anchor + batch_idx * dim;
    const float* positive_ptr = positive + batch_idx * dim;
    const float* negative_ptr = negative + batch_idx * dim;

    // 计算正对和负对的距离
    float dist_pos_sq = squared_l2_distance(anchor_ptr, positive_ptr, dim, sdata);
    float dist_neg_sq = squared_l2_distance(anchor_ptr, negative_ptr, dim, sdata);

    // 只有块中的第一个线程执行最终的计算和写入操作
    if (threadIdx.x == 0) {
        float dist_pos = sqrtf(dist_pos_sq);
        float dist_neg = sqrtf(dist_neg_sq);
        float loss_val = dist_pos - dist_neg + margin;
 
        loss[batch_idx] = fmaxf(0.0f, loss_val);
    }
}

torch::Tensor triplet_loss_cuda(
    torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, float margin) {
    
    int batch_size = anchor.size(0);
    int dim = anchor.size(1);
    
    auto options = torch::TensorOptions().device(anchor.device()).dtype(anchor.dtype());
    auto loss = torch::empty({batch_size}, options);

    dim3 grid(batch_size);

    dim3 block(BLOCK_SIZE);

    triplet_loss_kernel<<<grid, block>>>(
        anchor.data_ptr<float>(), positive.data_ptr<float>(), negative.data_ptr<float>(),
        loss.data_ptr<float>(), batch_size, dim, margin);
    
    return loss.sum()/batch_size;
}
"""

triplet_loss_cpp_source = """
torch::Tensor triplet_loss_cuda(
    torch::Tensor anchor, torch::Tensor positive, torch::Tensor negative, float margin);
"""

# 编译内联CUDA代码
triplet_loss_module = load_inline(
    name="triplet_loss_module",
    cpp_sources=triplet_loss_cpp_source,
    cuda_sources=triplet_loss_source,
    functions=["triplet_loss_cuda"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, margin: float = 1.0):
        super(ModelNew, self).__init__()
        self.margin = margin
        self.triplet_loss = triplet_loss_module

    def forward(self, anchor, positive, negative):
        return self.triplet_loss.triplet_loss_cuda(anchor, positive, negative, self.margin)