forked from ccf-ai-infra/GPUCodeForces
119 lines
3.9 KiB
Python
119 lines
3.9 KiB
Python
import torch
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
from CrossEntropyLoss_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 cel_forward_cuda(torch::Tensor input, torch::Tensor target);
|
|
"""
|
|
|
|
cuda_source = """
|
|
#include <cuda_runtime.h>
|
|
#include <float.h>
|
|
#include <cmath>
|
|
|
|
#define BLOCK_SIZE 256
|
|
|
|
__global__ void cel_fused_kernel(
|
|
const float* __restrict__ input, // [B, C]
|
|
const int64_t* __restrict__ target,// [B]
|
|
float* __restrict__ loss_partial, // [B]
|
|
int num_classes)
|
|
{
|
|
extern __shared__ float smem[]; // 动态 shared mem: 存 logits & exp
|
|
float* logits = smem;
|
|
float* exp_logits = smem + num_classes;
|
|
|
|
int sample_idx = blockIdx.x;
|
|
const float* row = input + sample_idx * num_classes;
|
|
int label = target[sample_idx];
|
|
|
|
// 1. load logits to shared memory
|
|
for (int j = threadIdx.x; j < num_classes; j += blockDim.x) {
|
|
logits[j] = row[j];
|
|
}
|
|
__syncthreads();
|
|
|
|
// 2. compute row max for numerical stability
|
|
float local_max = -FLT_MAX;
|
|
for (int j = threadIdx.x; j < num_classes; j += blockDim.x)
|
|
local_max = fmaxf(local_max, logits[j]);
|
|
|
|
// block reduce max
|
|
__shared__ float row_max;
|
|
if (threadIdx.x == 0) row_max = -FLT_MAX;
|
|
__syncthreads();
|
|
|
|
atomicMax((int*)&row_max, __float_as_int(local_max));
|
|
__syncthreads();
|
|
|
|
// 3. compute exp(x - max) and sum
|
|
float local_sum = 0.0f;
|
|
for (int j = threadIdx.x; j < num_classes; j += blockDim.x) {
|
|
float e = expf(logits[j] - row_max);
|
|
exp_logits[j] = e;
|
|
local_sum += e;
|
|
}
|
|
|
|
// block reduce sum
|
|
__shared__ float row_sum;
|
|
if (threadIdx.x == 0) row_sum = 0.0f;
|
|
__syncthreads();
|
|
|
|
atomicAdd(&row_sum, local_sum);
|
|
__syncthreads();
|
|
|
|
// 4. compute -log(p_correct)
|
|
float loss_val = 0.0f;
|
|
if (threadIdx.x == 0) {
|
|
float p_correct = exp_logits[label] / row_sum;
|
|
loss_val = -logf(p_correct);
|
|
loss_partial[sample_idx] = loss_val;
|
|
}
|
|
}
|
|
|
|
torch::Tensor cel_forward_cuda(torch::Tensor input, torch::Tensor target) {
|
|
TORCH_CHECK(input.is_cuda(), "input must be CUDA tensor");
|
|
TORCH_CHECK(target.is_cuda(), "target must be CUDA tensor");
|
|
|
|
input = input.contiguous();
|
|
target = target.contiguous();
|
|
|
|
const int batch_size = input.size(0);
|
|
const int num_classes = input.size(1);
|
|
|
|
auto loss_buf = torch::empty({batch_size}, input.options());
|
|
|
|
const dim3 grid(batch_size);
|
|
const dim3 block(BLOCK_SIZE);
|
|
const size_t shmem_bytes = 2 * num_classes * sizeof(float);
|
|
|
|
cel_fused_kernel<<<grid, block, shmem_bytes>>>(
|
|
input.data_ptr<float>(),
|
|
target.data_ptr<int64_t>(),
|
|
loss_buf.data_ptr<float>(),
|
|
num_classes
|
|
);
|
|
|
|
return loss_buf.mean();
|
|
}
|
|
"""
|
|
|
|
self.cel_op = load_inline(
|
|
name="cel_fused_op_v2",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["cel_forward_cuda"],
|
|
extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"],
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
return self.cel_op.cel_forward_cuda(input, target) |