forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish huberloss #34' (#61) from hli28146/GPUCodeForces:huberloss into main
This commit is contained in:
commit
3c9d445a58
|
|
@ -0,0 +1,183 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
import math
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <string>
|
||||
|
||||
torch::Tensor huber_loss_cuda_forward(
|
||||
const torch::Tensor& input,
|
||||
const torch::Tensor& target,
|
||||
double delta,
|
||||
const std::string& reduction
|
||||
);
|
||||
"""
|
||||
|
||||
# CUDA 源代码，包含核函数和其调用封装
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <limits>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// Element-wise Kernel for reduction='none'
|
||||
// ----------------------------------------------------------------------------
|
||||
template <typename T>
|
||||
__global__ void huber_loss_elementwise_kernel(
|
||||
T* output,
|
||||
const T* input,
|
||||
const T* target,
|
||||
T delta,
|
||||
int64_t n_elements)
|
||||
{
|
||||
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= n_elements) return;
|
||||
|
||||
T diff = input[idx] - target[idx];
|
||||
T abs_diff = fabsf(diff);
|
||||
if (abs_diff < delta) {
|
||||
output[idx] = 0.5f * diff * diff;
|
||||
} else {
|
||||
output[idx] = delta * (abs_diff - 0.5f * delta);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// ----------------------------------------------------------------------------
|
||||
// Reduction Kernel: Stage 1 (calculate loss and reduce within blocks)
|
||||
// ----------------------------------------------------------------------------
|
||||
template <typename T>
|
||||
__global__ void huber_loss_reduce_kernel_stage1(
|
||||
T* block_results,
|
||||
const T* input,
|
||||
const T* target,
|
||||
T delta,
|
||||
int64_t n_elements)
|
||||
{
|
||||
__shared__ T sdata[BLOCK_SIZE];
|
||||
|
||||
int64_t tid = threadIdx.x;
|
||||
int64_t i = blockIdx.x * (blockDim.x * 2) + tid;
|
||||
int64_t gridSize = blockDim.x * 2 * gridDim.x;
|
||||
|
||||
T my_sum = 0.0f;
|
||||
|
||||
while (i < n_elements) {
|
||||
T diff1 = input[i] - target[i];
|
||||
T abs_diff1 = fabsf(diff1);
|
||||
my_sum += (abs_diff1 < delta) ? (0.5f * diff1 * diff1) : (delta * (abs_diff1 - 0.5f * delta));
|
||||
|
||||
if (i + blockDim.x < n_elements) {
|
||||
T diff2 = input[i + blockDim.x] - target[i + blockDim.x];
|
||||
T abs_diff2 = fabsf(diff2);
|
||||
my_sum += (abs_diff2 < delta) ? (0.5f * diff2 * diff2) : (delta * (abs_diff2 - 0.5f * delta));
|
||||
}
|
||||
i += gridSize;
|
||||
}
|
||||
sdata[tid] = my_sum;
|
||||
__syncthreads();
|
||||
|
||||
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) {
|
||||
sdata[tid] += sdata[tid + s];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// 第一个线程将该块的结果写入全局内存
|
||||
if (tid == 0) {
|
||||
block_results[blockIdx.x] = sdata[0];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
torch::Tensor huber_loss_cuda_forward(
|
||||
const torch::Tensor& input,
|
||||
const torch::Tensor& target,
|
||||
double delta_d,
|
||||
const std::string& reduction)
|
||||
{
|
||||
TORCH_CHECK(input.is_cuda(), "Input tensor must be a CUDA tensor");
|
||||
TORCH_CHECK(target.is_cuda(), "Target tensor must be a CUDA tensor");
|
||||
TORCH_CHECK(input.sizes() == target.sizes(), "Input and target shapes must match");
|
||||
TORCH_CHECK(input.is_contiguous(), "Input tensor must be contiguous");
|
||||
TORCH_CHECK(target.is_contiguous(), "Target tensor must be contiguous");
|
||||
|
||||
const int64_t n_elements = input.numel();
|
||||
auto scalar_type = input.scalar_type();
|
||||
|
||||
if (reduction == "none") {
|
||||
auto output = torch::empty_like(input);
|
||||
const int num_blocks = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
||||
AT_DISPATCH_FLOATING_TYPES(scalar_type, "huber_loss_elementwise", ([&] {
|
||||
huber_loss_elementwise_kernel<scalar_t><<<num_blocks, BLOCK_SIZE>>>(
|
||||
output.data_ptr<scalar_t>(),
|
||||
input.data_ptr<scalar_t>(),
|
||||
target.data_ptr<scalar_t>(),
|
||||
static_cast<scalar_t>(delta_d),
|
||||
n_elements);
|
||||
}));
|
||||
return output;
|
||||
}
|
||||
else // 'sum' or 'mean'
|
||||
{
|
||||
int max_grid_size = 1024; // A reasonable limit for partial results
|
||||
int num_blocks = std::min(max_grid_size, (int)((n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE));
|
||||
|
||||
auto options = torch::TensorOptions().device(input.device()).dtype(input.dtype());
|
||||
auto block_results = torch::empty({num_blocks}, options);
|
||||
auto output = torch::empty({}, options);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(scalar_type, "huber_loss_reduce_stage1", ([&] {
|
||||
huber_loss_reduce_kernel_stage1<scalar_t><<<num_blocks, BLOCK_SIZE>>>(
|
||||
block_results.data_ptr<scalar_t>(),
|
||||
input.data_ptr<scalar_t>(),
|
||||
target.data_ptr<scalar_t>(),
|
||||
static_cast<scalar_t>(delta_d),
|
||||
n_elements);
|
||||
}));
|
||||
|
||||
// Stage 2: final reduction of block_results.
|
||||
// For simplicity and robustness, we can sum the small block_results tensor on the CPU
|
||||
// or launch another kernel. A simple sum() is often fast enough here.
|
||||
// A full GPU stage2 implementation would be similar to stage1 but on `block_results`.
|
||||
torch::Tensor total_sum_tensor = block_results.sum();
|
||||
|
||||
if (reduction == "mean") {
|
||||
return total_sum_tensor / n_elements;
|
||||
}
|
||||
return total_sum_tensor;
|
||||
}
|
||||
}
|
||||
|
||||
"""
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
"""
|
||||
使用自定义 CUDA 内核进行优化的 HuberLoss 模型。
|
||||
"""
|
||||
def __init__(self, delta=1.0, reduction='mean'):
|
||||
super(ModelNew, self).__init__()
|
||||
self.delta = delta
|
||||
self.reduction = reduction
|
||||
|
||||
self.huber_loss_op = load_inline(
|
||||
name='huber_loss_op',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['huber_loss_cuda_forward'],
|
||||
verbose=False,
|
||||
extra_cuda_cflags=["-O3"]
|
||||
)
|
||||
|
||||
def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
|
||||
return self.huber_loss_op.huber_loss_cuda_forward(
|
||||
input_tensor,
|
||||
target_tensor,
|
||||
self.delta,
|
||||
self.reduction
|
||||
)
|
||||
|
|
@ -0,0 +1,33 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
# --- 用于基准测试的配置 ---
|
||||
batch_size = 512
|
||||
dim = 4096
|
||||
DELTA = 1.0
|
||||
REDUCTION = 'mean'
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
使用 PyTorch 内置的 torch.nn.HuberLoss 作为基准模型。
|
||||
"""
|
||||
def __init__(self, delta=1.0, reduction='mean'):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = nn.HuberLoss(delta=delta, reduction=reduction)
|
||||
|
||||
def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
|
||||
return self.loss_fn(input_tensor, target_tensor)
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成用于测试的输入张量。
|
||||
"""
|
||||
input_tensor = torch.randn(batch_size, dim, dtype=torch.float32)
|
||||
target_tensor = input_tensor + torch.randn(batch_size, dim, dtype=torch.float32) * 0.5
|
||||
return [input_tensor.contiguous(), target_tensor.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
"""
|
||||
提供模型初始化所需的参数。
|
||||
"""
|
||||
return [DELTA, REDUCTION]
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
Write a custom CUDA kernel to optimize `torch.nn.HuberLoss`.
|
||||
|
||||
The original operation is defined by the formula:
|
||||
loss(x, y)_i =
|
||||
- 0.5 * (x_i - y_i)^2, if |x_i - y_i| < delta
|
||||
- delta * (|x_i - y_i| - 0.5 * delta), otherwise
|
||||
|
||||
This is followed by a reduction operation over all elements ('none', 'mean', or 'sum').
|
||||
|
||||
**Problem Analysis:**
|
||||
The standard PyTorch implementation of HuberLoss is memory-bound. It executes a chain of element-wise operations (subtraction, absolute value, comparison, multiplication, etc.) and a final reduction. Each step materializes a full-sized intermediate tensor in global memory, which is immediately read back by the next operation. This results in excessive memory traffic and multiple kernel launch overheads, which are the primary performance bottlenecks.
|
||||
|
||||
**Optimization Strategy: Fused Computation and Parallel Reduction**
|
||||
|
||||
The goal is to create a CUDA implementation that fuses all stages into one or two kernel launches.
|
||||
|
||||
1. **Fusion for `reduction='none'`**:
|
||||
A single element-wise kernel is implemented. Each thread is assigned to one element of the input tensors. It performs the entire Huber Loss calculation (diff, abs, condition, formula) in registers and writes the final result directly to the output tensor. This completely eliminates intermediate memory traffic.
|
||||
|
||||
2. **Fusion for `reduction='mean'` or `'sum'`**:
|
||||
A highly-optimized, two-stage parallel reduction strategy is employed:
|
||||
* **Kernel 1 (Calculation & Block-Level Reduction)**: This kernel is launched with a grid size large enough to cover all elements.
|
||||
- Each thread computes the Huber loss for one or more elements.
|
||||
- The results within a thread block are then efficiently summed up using **shared memory** in a tree-like reduction pattern. This avoids slow global memory atomics.
|
||||
- The first thread of each block writes its block's partial sum to a temporary intermediate buffer in global memory.
|
||||
* **Kernel 2 (Final Reduction)**: A second, much smaller kernel (often a single block) is launched. It reads the partial sums from the intermediate buffer and performs the final reduction, again using shared memory, to produce a single scalar result.
|
||||
* For `reduction='mean'`, the final sum is divided by the total number of elements.
|
||||
|
||||
This comprehensive fusion strategy minimizes global memory access to a single pass over the input data, drastically reducing bandwidth usage and kernel launch overhead, leading to significant performance gains.
|
||||
|
||||
You are given the following architecture:
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
# --- 用于基准测试的配置 ---
|
||||
batch_size = 512
|
||||
dim = 4096
|
||||
DELTA = 1.0
|
||||
REDUCTION = 'mean'
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
使用 PyTorch 内置的 torch.nn.HuberLoss 作为基准模型。
|
||||
"""
|
||||
def __init__(self, delta=1.0, reduction='mean'):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = nn.HuberLoss(delta=delta, reduction=reduction)
|
||||
|
||||
def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
|
||||
return self.loss_fn(input_tensor, target_tensor)
|
||||
|
||||
def get_inputs():
|
||||
"""
|
||||
生成用于测试的输入张量。
|
||||
"""
|
||||
input_tensor = torch.randn(batch_size, dim, dtype=torch.float32)
|
||||
target_tensor = input_tensor + torch.randn(batch_size, dim, dtype=torch.float32) * 0.5
|
||||
return [input_tensor.contiguous(), target_tensor.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
"""
|
||||
提供模型初始化所需的参数。
|
||||
"""
|
||||
return [DELTA, REDUCTION]
|
||||
|
|
@ -0,0 +1,80 @@
|
|||
import torch
|
||||
import time
|
||||
from huberloss_torch import Model, get_inputs, get_init_inputs
|
||||
from huberloss_cuda import ModelNew
|
||||
|
||||
def run_benchmark():
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用")
|
||||
return
|
||||
|
||||
device = torch.device("cuda")
|
||||
|
||||
# 准备输入数据
|
||||
inputs = [x.cuda(device=device) for x in get_inputs()]
|
||||
init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
|
||||
|
||||
# 初始化模型
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
# 预热GPU
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# 正式测试
|
||||
output_torch = torch_model(*inputs)
|
||||
output_cuda = cuda_model(*inputs)
|
||||
|
||||
# 更严格的精度检查
|
||||
abs_diff = (output_torch - output_cuda).abs()
|
||||
max_diff = abs_diff.max().item()
|
||||
mean_diff = abs_diff.mean().item()
|
||||
|
||||
print(f"最大差异: {max_diff:.6f}")
|
||||
print(f"平均差异: {mean_diff:.6f}")
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-05, atol=1e-05)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 100
|
||||
|
||||
# 预热GPU
|
||||
for _ in range(10):
|
||||
_ = torch_model(*inputs)
|
||||
_ = cuda_model(*inputs)
|
||||
|
||||
# PyTorch模型计时
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
_ = torch_model(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
torch_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
# 自定义CUDA内核计时
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
_ = cuda_model(*inputs)
|
||||
torch.cuda.synchronize()
|
||||
cuda_time = (time.time() - start_time) / num_iterations
|
||||
|
||||
print(f"PyTorch内置HuberLoss平均执行时间: {torch_time:.6f}秒")
|
||||
print(f"自定义CUDA HuberLoss平均执行时间: {cuda_time:.6f}秒")
|
||||
speedup = torch_time / cuda_time if cuda_time > 0 else 0
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
|
||||
return precision_flag, speedup
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag, speedup = run_benchmark()
|
||||
Loading…
Reference in New Issue