forked from ccf-ai-infra/GPUCodeForces
122 lines
3.9 KiB
Python
122 lines
3.9 KiB
Python
# marginrankingloss_cuda.py
|
|
import torch
|
|
from torch.utils.cpp_extension import load_inline
|
|
from marginrankingloss_torch import BATCH_SIZE, FEATURE_DIM, MARGIN
|
|
|
|
|
|
|
|
N_ELEMENTS = 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 mrl_forward_cuda(torch::Tensor x1, torch::Tensor x2, torch::Tensor target, float margin_val);
|
|
"""
|
|
|
|
cuda_source = """
|
|
#include <cuda_runtime.h>
|
|
#include <cmath>
|
|
#include <float.h>
|
|
|
|
#define BLOCK_SIZE 256
|
|
#define MARGIN_VAL {margin_val}f
|
|
|
|
__device__ __forceinline__ float warp_reduce_sum(float val) {{
|
|
for (int offset = BLOCK_SIZE / 2; offset > 0; offset /= 2) {{
|
|
val += __shfl_down_sync(0xffffffff, val, offset);
|
|
}}
|
|
return val;
|
|
}}
|
|
|
|
|
|
__global__ void mrl_fused_kernel(
|
|
const float* __restrict__ x1,
|
|
const float* __restrict__ x2,
|
|
const float* __restrict__ target,
|
|
float* __restrict__ loss_sum_out,
|
|
int n_elements,
|
|
float margin_val
|
|
) {{
|
|
__shared__ float s_data[BLOCK_SIZE];
|
|
|
|
// Grid-Stride Loop
|
|
float thread_loss = 0.0f;
|
|
int grid_stride = gridDim.x * blockDim.x;
|
|
|
|
for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
idx < n_elements;
|
|
idx += grid_stride)
|
|
{{
|
|
float val1 = x1[idx];
|
|
float val2 = x2[idx];
|
|
float y = target[idx];
|
|
|
|
// Loss = max(0, -y * (x1 - x2) + margin)
|
|
|
|
float diff = val1 - val2;
|
|
float term = -y * diff + MARGIN_VAL;
|
|
|
|
float loss_val = fmaxf(0.0f, term);
|
|
|
|
thread_loss += loss_val;
|
|
}}
|
|
|
|
s_data[threadIdx.x] = thread_loss;
|
|
__syncthreads();
|
|
|
|
for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {{
|
|
if (threadIdx.x < offset) {{
|
|
s_data[threadIdx.x] += s_data[threadIdx.x + offset];
|
|
}}
|
|
__syncthreads();
|
|
}}
|
|
|
|
if (threadIdx.x == 0) {{
|
|
loss_sum_out[blockIdx.x] = s_data[0];
|
|
}}
|
|
}}
|
|
|
|
torch::Tensor mrl_forward_cuda(torch::Tensor x1, torch::Tensor x2, torch::Tensor target, float margin_val) {{
|
|
TORCH_CHECK(x1.is_cuda(), "Input must be a CUDA tensor");
|
|
x1 = x1.contiguous();
|
|
x2 = x2.contiguous();
|
|
target = target.contiguous();
|
|
|
|
int n_elements = x1.numel();
|
|
|
|
const int block_size = BLOCK_SIZE;
|
|
const int grid_size = std::max(1, (n_elements + block_size - 1) / block_size);
|
|
|
|
auto block_loss_sums = torch::empty({{grid_size}}, x1.options());
|
|
|
|
mrl_fused_kernel<<<grid_size, block_size>>>(
|
|
x1.data_ptr<float>(),
|
|
x2.data_ptr<float>(),
|
|
target.data_ptr<float>(),
|
|
block_loss_sums.data_ptr<float>(),
|
|
n_elements,
|
|
margin_val
|
|
);
|
|
|
|
|
|
return block_loss_sums.sum() / n_elements;
|
|
}}
|
|
""".format(margin_val=MARGIN)
|
|
|
|
self.mrl_op = load_inline(
|
|
name="mrl_fused_op_safest",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["mrl_forward_cuda"],
|
|
extra_cuda_cflags=["-O3", "--use_fast_math"],
|
|
verbose=True
|
|
)
|
|
|
|
def forward(self, x1: torch.Tensor, x2: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
return self.mrl_op.mrl_forward_cuda(x1, x2, target, MARGIN) |