forked from ccf-ai-infra/GPUCodeForces
156 lines
4.9 KiB
Python
156 lines
4.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
BATCH_SIZE = 4096
|
|
N_CLASSES = 1024
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
cpp_source = """
|
|
#include <torch/extension.h>
|
|
|
|
// C++ 接口
|
|
torch::Tensor cross_entropy_forward_cuda(
|
|
torch::Tensor logits,
|
|
torch::Tensor target
|
|
);
|
|
"""
|
|
|
|
cuda_source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
#include <cmath>
|
|
#include <float.h>
|
|
|
|
#define BLOCK_SIZE 256
|
|
|
|
__global__ void cross_entropy_fused_kernel(
|
|
const float* __restrict__ logits_data, // (N, C)
|
|
const int64_t* __restrict__ target_data, // (N,)
|
|
float* __restrict__ loss_per_row_out, // (N,)
|
|
int N,
|
|
int C
|
|
) {
|
|
// 当前处理的 Batch 索引
|
|
int row_idx = blockIdx.x;
|
|
if (row_idx >= N) return;
|
|
|
|
// 当前行的指针
|
|
const float* row_logits = logits_data + row_idx * C;
|
|
int tid = threadIdx.x;
|
|
|
|
// 共享内存：用于 Max 和 Sum 的归约
|
|
__shared__ float s_data[BLOCK_SIZE];
|
|
|
|
float thread_max = -FLT_MAX;
|
|
|
|
// Grid-Stride Loop 遍历类别 C
|
|
for (int c = tid; c < C; c += BLOCK_SIZE) {
|
|
float val = row_logits[c];
|
|
if (val > thread_max) {
|
|
thread_max = val;
|
|
}
|
|
}
|
|
s_data[tid] = thread_max;
|
|
__syncthreads();
|
|
|
|
// 块内归约 (Max)
|
|
for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
|
|
if (tid < offset) {
|
|
if (s_data[tid + offset] > s_data[tid]) {
|
|
s_data[tid] = s_data[tid + offset];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
|
|
float row_max_val = s_data[0];
|
|
__syncthreads();
|
|
|
|
float thread_sum_exp = 0.0f;
|
|
|
|
for (int c = tid; c < C; c += BLOCK_SIZE) {
|
|
float val = row_logits[c];
|
|
thread_sum_exp += expf(val - row_max_val);
|
|
}
|
|
s_data[tid] = thread_sum_exp;
|
|
__syncthreads();
|
|
|
|
// 块内归约 (Sum)
|
|
for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
|
|
if (tid < offset) {
|
|
s_data[tid] += s_data[tid + offset];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
if (tid == 0) {
|
|
float row_sum_exp = s_data[0];
|
|
float log_sum_exp = logf(row_sum_exp) + row_max_val;
|
|
|
|
int64_t target_class = target_data[row_idx];
|
|
float target_logit = row_logits[target_class];
|
|
|
|
// Cross Entropy Formula
|
|
loss_per_row_out[row_idx] = -target_logit + log_sum_exp;
|
|
}
|
|
}
|
|
|
|
// C++ 封装函数
|
|
torch::Tensor cross_entropy_forward_cuda(
|
|
torch::Tensor logits,
|
|
torch::Tensor target
|
|
) {
|
|
TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor");
|
|
TORCH_CHECK(target.is_cuda(), "target must be a CUDA tensor");
|
|
TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
|
|
TORCH_CHECK(target.dim() == 1, "target must be 1D");
|
|
|
|
// 确保连续
|
|
logits = logits.contiguous();
|
|
target = target.contiguous();
|
|
|
|
int N = logits.size(0); // Batch Size
|
|
int C = logits.size(1); // Num Classes
|
|
|
|
TORCH_CHECK(target.size(0) == N, "Target size mismatch");
|
|
|
|
auto losses = torch::empty({N}, logits.options());
|
|
|
|
// 启动配置:
|
|
// Grid: N (每个 Batch 一个 Block)
|
|
// Block: 256
|
|
dim3 grid_dim(N);
|
|
dim3 block_dim(BLOCK_SIZE);
|
|
|
|
cross_entropy_fused_kernel<<<grid_dim, block_dim>>>(
|
|
logits.data_ptr<float>(),
|
|
target.data_ptr<int64_t>(),
|
|
losses.data_ptr<float>(),
|
|
N, C
|
|
);
|
|
|
|
// 返回 Mean Reduction
|
|
return losses.mean();
|
|
}
|
|
"""
|
|
|
|
self.ce_op = load_inline(
|
|
name="cross_entropy_op_v1",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=cuda_source,
|
|
functions=["cross_entropy_forward_cuda"],
|
|
extra_cuda_cflags=["-O3"],
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
return self.ce_op.cross_entropy_forward_cuda(logits, target) |