forked from ccf-ai-infra/GPUCodeForces
197 lines
6.4 KiB
Python
197 lines
6.4 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
from torch.utils.cpp_extension import load_inline
|
||
|
||
N, D = 32, 128
|
||
|
||
assert D % 4 == 0, "Embedding dimension D must be a multiple of 4 for vectorization"
|
||
|
||
|
||
class ModelNew(nn.Module):
|
||
|
||
def __init__(self, margin=1.0, swap=False):
|
||
super().__init__()
|
||
self.margin = float(margin)
|
||
self.swap = swap
|
||
self.block_size = 256
|
||
self._compile_cuda_kernel()
|
||
|
||
def _compile_cuda_kernel(self):
|
||
cpp_source = """
|
||
#include <torch/extension.h>
|
||
|
||
torch::Tensor triplet_forward_cuda(
|
||
torch::Tensor anchor,
|
||
torch::Tensor positive,
|
||
torch::Tensor negative,
|
||
float margin,
|
||
bool swap,
|
||
int N,
|
||
int D);
|
||
"""
|
||
|
||
cuda_source = f"""
|
||
#include <cuda_runtime.h>
|
||
#include <cmath>
|
||
|
||
#define BLOCK_SIZE {self.block_size}
|
||
#define WARP_SIZE 32
|
||
|
||
// Warp 归约工具
|
||
__inline__ __device__ float warp_reduce_sum(float val) {{
|
||
#pragma unroll
|
||
for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {{
|
||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||
}}
|
||
return val;
|
||
}}
|
||
|
||
|
||
__global__ void triplet_l2_kernel(
|
||
const float* __restrict__ anchor,
|
||
const float* __restrict__ positive,
|
||
const float* __restrict__ negative,
|
||
float* __restrict__ output,
|
||
float margin,
|
||
bool swap,
|
||
int D_vec // D / 4
|
||
) {{
|
||
const int n_idx = blockIdx.x;
|
||
const int tid = threadIdx.x;
|
||
|
||
|
||
const int offset = n_idx * D_vec * 4;
|
||
const float4* a_ptr = reinterpret_cast<const float4*>(anchor + offset);
|
||
const float4* p_ptr = reinterpret_cast<const float4*>(positive + offset);
|
||
const float4* n_ptr = reinterpret_cast<const float4*>(negative + offset);
|
||
|
||
|
||
float sum_sq_ap = 0.0f;
|
||
float sum_sq_an = 0.0f;
|
||
float sum_sq_pn = 0.0f;
|
||
|
||
|
||
for (int i = tid; i < D_vec; i += BLOCK_SIZE) {{
|
||
float4 a = __ldg(&a_ptr[i]);
|
||
float4 p = __ldg(&p_ptr[i]);
|
||
float4 n = __ldg(&n_ptr[i]);
|
||
|
||
|
||
float4 diff_ap, diff_an, diff_pn;
|
||
|
||
diff_ap.x = a.x - p.x; diff_ap.y = a.y - p.y; diff_ap.z = a.z - p.z; diff_ap.w = a.w - p.w;
|
||
diff_an.x = a.x - n.x; diff_an.y = a.y - n.y; diff_an.z = a.z - n.z; diff_an.w = a.w - n.w;
|
||
|
||
|
||
sum_sq_ap += diff_ap.x*diff_ap.x + diff_ap.y*diff_ap.y + diff_ap.z*diff_ap.z + diff_ap.w*diff_ap.w;
|
||
sum_sq_an += diff_an.x*diff_an.x + diff_an.y*diff_an.y + diff_an.z*diff_an.z + diff_an.w*diff_an.w;
|
||
|
||
if (swap) {{
|
||
diff_pn.x = p.x - n.x; diff_pn.y = p.y - n.y; diff_pn.z = p.z - n.z; diff_pn.w = p.w - n.w;
|
||
sum_sq_pn += diff_pn.x*diff_pn.x + diff_pn.y*diff_pn.y + diff_pn.z*diff_pn.z + diff_pn.w*diff_pn.w;
|
||
}}
|
||
}}
|
||
|
||
|
||
__shared__ float shared_data[32][3];
|
||
|
||
int lane = tid % WARP_SIZE;
|
||
int wid = tid / WARP_SIZE;
|
||
|
||
|
||
sum_sq_ap = warp_reduce_sum(sum_sq_ap);
|
||
sum_sq_an = warp_reduce_sum(sum_sq_an);
|
||
if (swap) sum_sq_pn = warp_reduce_sum(sum_sq_pn);
|
||
|
||
|
||
if (lane == 0) {{
|
||
shared_data[wid][0] = sum_sq_ap;
|
||
shared_data[wid][1] = sum_sq_an;
|
||
if (swap) shared_data[wid][2] = sum_sq_pn;
|
||
}}
|
||
__syncthreads();
|
||
|
||
|
||
if (wid == 0) {{
|
||
sum_sq_ap = (tid < blockDim.x / WARP_SIZE) ? shared_data[lane][0] : 0.0f;
|
||
sum_sq_an = (tid < blockDim.x / WARP_SIZE) ? shared_data[lane][1] : 0.0f;
|
||
sum_sq_pn = (tid < blockDim.x / WARP_SIZE && swap) ? shared_data[lane][2] : 0.0f;
|
||
|
||
sum_sq_ap = warp_reduce_sum(sum_sq_ap);
|
||
sum_sq_an = warp_reduce_sum(sum_sq_an);
|
||
if (swap) sum_sq_pn = warp_reduce_sum(sum_sq_pn);
|
||
|
||
|
||
if (tid == 0) {{
|
||
// 开根号得到 L2 距离 (加上 epsilon 防止梯度爆炸通常在backward处理，前向计算通常加个极小值)
|
||
float dist_ap = sqrtf(sum_sq_ap + 1e-8f);
|
||
float dist_an = sqrtf(sum_sq_an + 1e-8f);
|
||
|
||
if (swap) {{
|
||
float dist_pn = sqrtf(sum_sq_pn + 1e-8f);
|
||
if (dist_pn < dist_an) {{
|
||
dist_an = dist_pn;
|
||
}}
|
||
}}
|
||
|
||
// loss = max(d_ap - d_an + margin, 0)
|
||
float loss = fmaxf(dist_ap - dist_an + margin, 0.0f);
|
||
output[n_idx] = loss;
|
||
}}
|
||
}}
|
||
}}
|
||
|
||
torch::Tensor triplet_forward_cuda(
|
||
torch::Tensor anchor,
|
||
torch::Tensor positive,
|
||
torch::Tensor negative,
|
||
float margin,
|
||
bool swap,
|
||
int N,
|
||
int D)
|
||
{{
|
||
anchor = anchor.contiguous();
|
||
positive = positive.contiguous();
|
||
negative = negative.contiguous();
|
||
|
||
auto output = torch::empty({{N}}, anchor.options());
|
||
|
||
int D_vec = D / 4;
|
||
|
||
dim3 blocks(N);
|
||
dim3 threads(BLOCK_SIZE);
|
||
|
||
triplet_l2_kernel<<<blocks, threads>>>(
|
||
anchor.data_ptr<float>(),
|
||
positive.data_ptr<float>(),
|
||
negative.data_ptr<float>(),
|
||
output.data_ptr<float>(),
|
||
margin,
|
||
swap,
|
||
D_vec
|
||
);
|
||
|
||
return output;
|
||
}}
|
||
"""
|
||
|
||
self.op = load_inline(
|
||
name='triplet_loss_cuda_v1',
|
||
cpp_sources=cpp_source,
|
||
cuda_sources=cuda_source,
|
||
functions=['triplet_forward_cuda'],
|
||
extra_cuda_cflags=['-O3', '--use_fast_math'],
|
||
verbose=False
|
||
)
|
||
|
||
def forward(self, a: torch.Tensor, p: torch.Tensor, n: torch.Tensor) -> torch.Tensor:
|
||
|
||
if not a.is_cuda: a = a.cuda()
|
||
if not p.is_cuda: p = p.cuda()
|
||
if not n.is_cuda: n = n.cuda()
|
||
|
||
N, D = a.shape
|
||
|
||
losses = self.op.triplet_forward_cuda(a, p, n, self.margin, self.swap, N, D)
|
||
|
||
return losses.mean() |