forked from ccf-ai-infra/GPUCodeForces
134 lines
3.9 KiB
Python
134 lines
3.9 KiB
Python
import torch
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
hingeloss_source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
// 高效归约版本 - 最优性能和精度
|
|
__global__ void hingeloss_kernel_optimal(
|
|
const float* __restrict__ input,
|
|
const float* __restrict__ target,
|
|
float* __restrict__ partial_sums,
|
|
int total_elements,
|
|
float margin
|
|
) {
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
// 局部累加器
|
|
float local_loss = 0.0f;
|
|
|
|
// 每个线程处理多个元素
|
|
int stride = blockDim.x * gridDim.x;
|
|
for (int element_idx = idx; element_idx < total_elements; element_idx += stride) {
|
|
// Hinge Loss: max(0, margin - y_true * y_pred)
|
|
float prediction = input[element_idx];
|
|
float label = target[element_idx];
|
|
|
|
// 计算 margin - y_true * y_pred
|
|
float margin_diff = margin - label * prediction;
|
|
|
|
// Hinge Loss: max(0, margin_diff)
|
|
float loss = (margin_diff > 0.0f) ? margin_diff : 0.0f;
|
|
|
|
local_loss += loss;
|
|
}
|
|
|
|
// 使用共享内存进行块内归约
|
|
extern __shared__ float shared_mem[];
|
|
shared_mem[threadIdx.x] = local_loss;
|
|
|
|
__syncthreads();
|
|
|
|
// 块内归约
|
|
for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
|
|
if (threadIdx.x < stride) {
|
|
shared_mem[threadIdx.x] += shared_mem[threadIdx.x + stride];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
// 每个block写入部分和到全局内存
|
|
if (threadIdx.x == 0) {
|
|
partial_sums[blockIdx.x] = shared_mem[0];
|
|
}
|
|
}
|
|
|
|
torch::Tensor hingeloss_cuda(
|
|
torch::Tensor input,
|
|
torch::Tensor target,
|
|
float margin,
|
|
std::string mode = "optimal"
|
|
) {
|
|
TORCH_CHECK(input.scalar_type() == torch::kFloat32, "Input must be float32");
|
|
TORCH_CHECK(target.scalar_type() == torch::kFloat32, "Target must be float32");
|
|
TORCH_CHECK(input.sizes() == target.sizes(), "Input and target must have same shape");
|
|
|
|
auto input_contig = input.contiguous();
|
|
auto target_contig = target.contiguous();
|
|
|
|
int total_elements = input_contig.numel();
|
|
|
|
// 只使用最优的efficient模式
|
|
const int block_size = 256;
|
|
|
|
// 针对不同数据量的优化block分配
|
|
int num_blocks;
|
|
if (total_elements <= 2048) {
|
|
num_blocks = 4;
|
|
} else if (total_elements <= 8192) {
|
|
num_blocks = 8;
|
|
} else if (total_elements <= 32768) {
|
|
num_blocks = 16;
|
|
} else if (total_elements <= 131072) {
|
|
num_blocks = 32;
|
|
} else {
|
|
num_blocks = 64;
|
|
}
|
|
|
|
// 创建部分和数组
|
|
auto partial_sums = torch::zeros({num_blocks}, input.options());
|
|
|
|
size_t shared_mem = block_size * sizeof(float);
|
|
hingeloss_kernel_optimal<<<num_blocks, block_size, shared_mem>>>(
|
|
input_contig.data_ptr<float>(),
|
|
target_contig.data_ptr<float>(),
|
|
partial_sums.data_ptr<float>(),
|
|
total_elements,
|
|
margin
|
|
);
|
|
|
|
// 在GPU上完成最终归约
|
|
auto total_loss = torch::sum(partial_sums);
|
|
return total_loss;
|
|
}
|
|
"""
|
|
|
|
hingeloss_cpp_source = """
|
|
torch::Tensor hingeloss_cuda(torch::Tensor input, torch::Tensor target, float margin, std::string mode);
|
|
"""
|
|
|
|
# 编译CUDA代码 - 平衡精度和性能
|
|
hingeloss = load_inline(
|
|
name="hingeloss",
|
|
cpp_sources=hingeloss_cpp_source,
|
|
cuda_sources=hingeloss_source,
|
|
functions=["hingeloss_cuda"],
|
|
extra_cuda_cflags=[
|
|
"-O3",
|
|
"--use_fast_math",
|
|
"-gencode=arch=compute_80,code=sm_80"
|
|
],
|
|
verbose=True
|
|
)
|
|
|
|
class ModelNew(torch.nn.Module):
|
|
def __init__(self, margin=1.0, mode="optimal"):
|
|
super(ModelNew, self).__init__()
|
|
self.margin = margin
|
|
self.mode = mode
|
|
self.hingeloss = hingeloss
|
|
|
|
def forward(self, input, target):
|
|
return self.hingeloss.hingeloss_cuda(input, target, self.margin, self.mode)
|