forked from ccf-ai-infra/GPUCodeForces
280 lines
9.8 KiB
Python
280 lines
9.8 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.cpp_extension import load_inline
|
|
import math
|
|
|
|
# -------------------------------------------------------------
|
|
# 常量定义 (同上)
|
|
# -------------------------------------------------------------
|
|
N, C, H, W = 8, 10, 16, 16
|
|
# 确保 weight 在 CUDA 上
|
|
WEIGHT = torch.rand(C, dtype=torch.float32).cuda()
|
|
IGNORE_INDEX = -100
|
|
REDUCTION = 'mean'
|
|
BLOCK_SIZE = 256
|
|
|
|
|
|
# -------------------------------------------------------------
|
|
|
|
class ModelNew(nn.Module):
|
|
|
|
def __init__(self, weight=None, size_average=None, ignore_index=-100,
|
|
reduce=None, reduction='mean'):
|
|
super().__init__()
|
|
self.reduction_str = reduction
|
|
self.ignore_index = ignore_index
|
|
|
|
if weight is not None:
|
|
self.register_buffer('weight', weight.contiguous())
|
|
else:
|
|
self.weight = None
|
|
|
|
self.block_size = BLOCK_SIZE
|
|
self._compile_cuda_kernel()
|
|
|
|
def _compile_cuda_kernel(self):
|
|
|
|
cpp_header = f"""
|
|
#include <torch/extension.h>
|
|
|
|
// C++ 接口
|
|
torch::Tensor nll_loss_forward_cuda(
|
|
torch::Tensor input,
|
|
torch::Tensor target,
|
|
c10::optional<torch::Tensor> weight,
|
|
int64_t ignore_index,
|
|
std::string reduction
|
|
);
|
|
"""
|
|
|
|
cuda_source = f"""
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
#define BLOCK_SIZE {self.block_size}
|
|
|
|
/* * 辅助函数：计算单个元素的 loss 和 weight
|
|
* (在两个核函数之间共享)
|
|
*/
|
|
__device__ inline void compute_nll_loss_item(
|
|
const float* __restrict__ input_data,
|
|
const int64_t* __restrict__ target_data,
|
|
const float* __restrict__ weight_data,
|
|
int i, int N_spatial, int C, int64_t ignore_index,
|
|
float& loss_val, // 输出
|
|
float& weight_val // 输出
|
|
) {{
|
|
const int64_t target_idx = target_data[i];
|
|
|
|
weight_val = 1.0f;
|
|
loss_val = 0.0f;
|
|
|
|
if (target_idx == ignore_index) {{
|
|
weight_val = 0.0f;
|
|
}} else if (target_idx < 0 || target_idx >= C) {{
|
|
weight_val = 0.0f;
|
|
}} else {{
|
|
if (weight_data != nullptr) {{
|
|
weight_val = weight_data[target_idx];
|
|
}}
|
|
|
|
const int n = i / N_spatial;
|
|
const int s = i % N_spatial;
|
|
// const int s = i - n * N_spatial; // 优化的 modulo
|
|
|
|
const int64_t input_idx =
|
|
(int64_t)n * C * N_spatial +
|
|
(int64_t)target_idx * N_spatial +
|
|
(int64_t)s;
|
|
|
|
loss_val = -weight_val * input_data[input_idx];
|
|
}}
|
|
}}
|
|
|
|
|
|
__global__ void nll_loss_kernel_no_reduce(
|
|
const float* __restrict__ input_data,
|
|
const int64_t* __restrict__ target_data,
|
|
const float* __restrict__ weight_data,
|
|
float* __restrict__ loss_out_data,
|
|
int N_spatial, int C, int N_total, int64_t ignore_index
|
|
) {{
|
|
for (int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
i < N_total;
|
|
i += gridDim.x * blockDim.x)
|
|
{{
|
|
float loss_val, weight_val;
|
|
compute_nll_loss_item(
|
|
input_data, target_data, weight_data,
|
|
i, N_spatial, C, ignore_index,
|
|
loss_val, weight_val
|
|
);
|
|
// 写入未归约的损失
|
|
loss_out_data[i] = loss_val;
|
|
}}
|
|
}}
|
|
|
|
|
|
__global__ void nll_loss_kernel_reduce(
|
|
const float* __restrict__ input_data,
|
|
const int64_t* __restrict__ target_data,
|
|
const float* __restrict__ weight_data,
|
|
float* __restrict__ partial_loss_out, // Block-level
|
|
float* __restrict__ partial_weight_out, // Block-level
|
|
int N_spatial, int C, int N_total, int64_t ignore_index
|
|
) {{
|
|
// 共享内存用于 Block 内部归约
|
|
__shared__ float s_loss[BLOCK_SIZE];
|
|
__shared__ float s_weight[BLOCK_SIZE];
|
|
|
|
const int tid = threadIdx.x;
|
|
float thread_loss_sum = 0.0f;
|
|
float thread_weight_sum = 0.0f;
|
|
|
|
// 1. Grid-Stride Loop: 计算和累加
|
|
for (int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
i < N_total;
|
|
i += gridDim.x * blockDim.x)
|
|
{{
|
|
float loss_val, weight_val;
|
|
compute_nll_loss_item(
|
|
input_data, target_data, weight_data,
|
|
i, N_spatial, C, ignore_index,
|
|
loss_val, weight_val
|
|
);
|
|
thread_loss_sum += loss_val;
|
|
thread_weight_sum += weight_val;
|
|
}}
|
|
|
|
s_loss[tid] = thread_loss_sum;
|
|
s_weight[tid] = thread_weight_sum;
|
|
__syncthreads();
|
|
|
|
// 2. 共享内存归约 (Sum)
|
|
for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {{
|
|
if (tid < offset) {{
|
|
s_loss[tid] += s_loss[tid + offset];
|
|
s_weight[tid] += s_weight[tid + offset];
|
|
}}
|
|
__syncthreads();
|
|
}}
|
|
|
|
// 3. 线程 0 写入 Block 的总和
|
|
if (tid == 0) {{
|
|
partial_loss_out[blockIdx.x] = s_loss[0];
|
|
partial_weight_out[blockIdx.x] = s_weight[0];
|
|
}}
|
|
}}
|
|
|
|
// C++ 封装函数
|
|
torch::Tensor nll_loss_forward_cuda(
|
|
torch::Tensor input,
|
|
torch::Tensor target,
|
|
c10::optional<torch::Tensor> weight_opt,
|
|
int64_t ignore_index,
|
|
std::string reduction
|
|
) {{
|
|
// 检查
|
|
TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
|
|
TORCH_CHECK(target.is_cuda(), "target must be a CUDA tensor");
|
|
TORCH_CHECK(input.dim() >= 2, "input must be >= 2D");
|
|
TORCH_CHECK(input.dim() == target.dim() + 1, "input/target dim mismatch");
|
|
TORCH_CHECK(target.scalar_type() == torch::kLong, "target must be torch.long");
|
|
|
|
const int64_t N = input.size(0);
|
|
const int64_t C = input.size(1);
|
|
|
|
int64_t N_spatial = 1;
|
|
if (input.dim() > 2) {{
|
|
for (int d = 2; d < input.dim(); ++d) {{
|
|
TORCH_CHECK(input.size(d) == target.size(d-1), "spatial dim mismatch");
|
|
N_spatial *= input.size(d);
|
|
}}
|
|
}}
|
|
const int64_t N_total = N * N_spatial;
|
|
|
|
if (N_total == 0) {{
|
|
return torch::tensor(0.0, input.options());
|
|
}}
|
|
|
|
auto input_flat = input.contiguous().view({{N, C, N_spatial}});
|
|
auto target_flat = target.contiguous().view({{N_total}});
|
|
|
|
const float* weight_ptr = nullptr;
|
|
if (weight_opt.has_value()) {{
|
|
auto& weight = weight_opt.value();
|
|
TORCH_CHECK(weight.is_cuda() && weight.is_contiguous() &&
|
|
weight.dim() == 1 && weight.size(0) == C, "weight size mismatch");
|
|
weight_ptr = weight.data_ptr<float>();
|
|
}}
|
|
|
|
// 4.核函数路由
|
|
dim3 block_dim(BLOCK_SIZE);
|
|
dim3 grid_dim((N_total + BLOCK_SIZE - 1) / BLOCK_SIZE);
|
|
|
|
if (reduction == "none") {{
|
|
auto output_loss = torch::empty_like(target_flat, input.options());
|
|
|
|
nll_loss_kernel_no_reduce<<<grid_dim, block_dim>>>(
|
|
input_flat.data_ptr<float>(),
|
|
target_flat.data_ptr<long>(),
|
|
weight_ptr,
|
|
output_loss.data_ptr<float>(),
|
|
N_spatial, C, N_total, ignore_index
|
|
);
|
|
return output_loss.view(target.sizes());
|
|
|
|
}} else {{ // "mean" or "sum"
|
|
|
|
// 创建小的部分和张量
|
|
auto partial_loss_out = torch::empty({{grid_dim.x}}, input.options());
|
|
auto partial_weight_out = torch::empty({{grid_dim.x}}, input.options());
|
|
|
|
nll_loss_kernel_reduce<<<grid_dim, block_dim>>>(
|
|
input_flat.data_ptr<float>(),
|
|
target_flat.data_ptr<long>(),
|
|
weight_ptr,
|
|
partial_loss_out.data_ptr<float>(),
|
|
partial_weight_out.data_ptr<float>(),
|
|
N_spatial, C, N_total, ignore_index
|
|
);
|
|
|
|
// 5. 在 C++ 中对*小的*部分和张量进行归约
|
|
torch::Tensor total_loss = partial_loss_out.sum();
|
|
|
|
if (reduction == "sum") {{
|
|
return total_loss;
|
|
}}
|
|
|
|
// reduction == "mean"
|
|
double total_weight = partial_weight_out.sum().item<double>();
|
|
if (total_weight == 0.0) {{
|
|
return torch::tensor(0.0, input.options());
|
|
}}
|
|
return total_loss / total_weight;
|
|
}}
|
|
}}
|
|
"""
|
|
|
|
nvcc_flags = ['-O3', '--use_fast_math']
|
|
|
|
# JIT (Just-In-Time) 编译
|
|
self.loss_op = load_inline(
|
|
name="nll_loss_op_v2_reduce", # 更改名称
|
|
cpp_sources=cpp_header,
|
|
cuda_sources=cuda_source,
|
|
functions=["nll_loss_forward_cuda"],
|
|
extra_cuda_cflags=nvcc_flags,
|
|
verbose=False
|
|
)
|
|
|
|
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
|
|
# C++ wrapper 现在处理 .contiguous()
|
|
return self.loss_op.nll_loss_forward_cuda(
|
|
input,
|
|
target,
|
|
self.weight,
|
|
self.ignore_index,
|
|
self.reduction_str
|
|
) |