forked from ccf-ai-infra/GPUCodeForces
finish multimarginloss
This commit is contained in:
parent
bee0a2a683
commit
d147f7ddb3
|
|
@ -0,0 +1,158 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.cpp_extension import load_inline
|
||||
|
||||
cpp_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <string>
|
||||
|
||||
torch::Tensor multi_margin_loss_cuda_forward(
|
||||
const torch::Tensor& input,
|
||||
const torch::Tensor& target,
|
||||
int p,
|
||||
float margin,
|
||||
const std::string& reduction
|
||||
);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
|
||||
#define BLOCK_SIZE 256
|
||||
|
||||
template <typename T>
|
||||
__global__ void multi_margin_loss_kernel(
|
||||
T* output, // (N)
|
||||
const T* input, // (N, C)
|
||||
const long* target, // (N)
|
||||
const int N,
|
||||
const int C,
|
||||
const int p,
|
||||
const T margin)
|
||||
{
|
||||
// Each block processes one sample
|
||||
int sample_idx = blockIdx.x;
|
||||
if (sample_idx >= N) return;
|
||||
|
||||
// --- Shared memory for broadcasting target index and correct class score ---
|
||||
__shared__ long y_shared;
|
||||
__shared__ T x_correct_shared;
|
||||
|
||||
// Thread 0 of each block loads the critical data
|
||||
if (threadIdx.x == 0) {
|
||||
long y_idx = target[sample_idx];
|
||||
y_shared = y_idx;
|
||||
x_correct_shared = input[sample_idx * C + y_idx];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// All threads in the block now have access to y_shared and x_correct_shared
|
||||
long y = y_shared;
|
||||
T x_correct = x_correct_shared;
|
||||
|
||||
__shared__ T sdata[BLOCK_SIZE];
|
||||
int tid = threadIdx.x;
|
||||
T my_sum = 0.0f;
|
||||
|
||||
// --- Grid-stride loop for this block to iterate over all classes ---
|
||||
for (int class_idx = tid; class_idx < C; class_idx += blockDim.x) {
|
||||
if (class_idx == y) {
|
||||
continue; // Skip the target class
|
||||
}
|
||||
|
||||
T x_other = input[sample_idx * C + class_idx];
|
||||
T loss_term = margin - x_correct + x_other;
|
||||
|
||||
if (loss_term > 0) {
|
||||
if (p == 2) {
|
||||
loss_term *= loss_term;
|
||||
}
|
||||
my_sum += loss_term;
|
||||
}
|
||||
}
|
||||
|
||||
sdata[tid] = my_sum;
|
||||
__syncthreads();
|
||||
|
||||
// --- Intra-block reduction to sum up all thread-local sums ---
|
||||
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) {
|
||||
sdata[tid] += sdata[tid + s];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// --- Thread 0 writes the final result for this sample ---
|
||||
if (tid == 0) {
|
||||
output[sample_idx] = sdata[0] / C;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
torch::Tensor multi_margin_loss_cuda_forward(
|
||||
const torch::Tensor& input,
|
||||
const torch::Tensor& target,
|
||||
int p,
|
||||
float margin,
|
||||
const std::string& reduction)
|
||||
{
|
||||
TORCH_CHECK(input.is_cuda() && target.is_cuda(), "Tensors must be on CUDA");
|
||||
TORCH_CHECK(input.dim() == 2, "Input must be 2D");
|
||||
TORCH_CHECK(target.dim() == 1, "Target must be 1D");
|
||||
TORCH_CHECK(input.size(0) == target.size(0), "Batch sizes must match");
|
||||
TORCH_CHECK(input.is_contiguous() && target.is_contiguous(), "Tensors must be contiguous");
|
||||
|
||||
const int N = input.size(0);
|
||||
const int C = input.size(1);
|
||||
|
||||
auto options = torch::TensorOptions().device(input.device()).dtype(input.dtype());
|
||||
auto sample_losses = torch::empty({N}, options);
|
||||
|
||||
// Launch one block per sample
|
||||
dim3 grid(N);
|
||||
dim3 block(BLOCK_SIZE);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "multi_margin_loss_kernel", ([&] {
|
||||
multi_margin_loss_kernel<scalar_t><<<grid, block>>>(
|
||||
sample_losses.data_ptr<scalar_t>(),
|
||||
input.data_ptr<scalar_t>(),
|
||||
target.data_ptr<long>(),
|
||||
N, C, p, static_cast<scalar_t>(margin)
|
||||
);
|
||||
}));
|
||||
|
||||
if (reduction == "none") {
|
||||
return sample_losses;
|
||||
} else if (reduction == "sum") {
|
||||
return sample_losses.sum();
|
||||
} else { // "mean"
|
||||
return sample_losses.mean();
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, p=1, margin=1.0, reduction='mean'):
|
||||
super(ModelNew, self).__init__()
|
||||
self.p = p
|
||||
self.margin = margin
|
||||
self.reduction = reduction
|
||||
|
||||
self.op = load_inline(
|
||||
name='multi_margin_loss_op',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['multi_margin_loss_cuda_forward'],
|
||||
verbose=False
|
||||
)
|
||||
|
||||
def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor:
|
||||
return self.op.multi_margin_loss_cuda_forward(
|
||||
input_tensor,
|
||||
target_tensor,
|
||||
self.p,
|
||||
self.margin,
|
||||
self.reduction
|
||||
)
|
||||
|
|
@ -0,0 +1,30 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
BATCH_SIZE = 512
|
||||
NUM_CLASSES = 4096
|
||||
REDUCTION = 'mean'
|
||||
P = 1 # 1 for L1 hinge, 2 for L2
|
||||
MARGIN = 1.0
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
使用 PyTorch 内置的 torch.nn.MultiMarginLoss 作为基准模型。
|
||||
"""
|
||||
def __init__(self, p=1, margin=1.0, reduction='mean'):
|
||||
super(Model, self).__init__()
|
||||
# weight is not benchmarked for simplicity, but the CUDA kernel supports it.
|
||||
self.loss_fn = nn.MultiMarginLoss(p=p, margin=margin, 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, NUM_CLASSES, dtype=torch.float32)
|
||||
|
||||
target_tensor = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
|
||||
|
||||
return [input_tensor.contiguous(), target_tensor.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return [P, MARGIN, REDUCTION]
|
||||
|
|
@ -0,0 +1,58 @@
|
|||
Write a custom CUDA kernel to optimize `torch.nn.MultiMarginLoss`.
|
||||
|
||||
The original operation is defined by the formula:
|
||||
`loss(x, y) = sum(max(0, margin - x[y] + x[i]))^p / C` for `i != y`.
|
||||
This is computed for each sample in the batch, and then a reduction is applied.
|
||||
|
||||
**Problem Analysis:**
|
||||
The standard PyTorch implementation of this loss is memory-bound and inefficient due to its operational complexity. It requires a sequence of advanced indexing (`gather`), broadcasting, masking (for `i != y`), element-wise operations (`max`, `pow`), and two levels of reduction (first over classes, then over the batch). Each step materializes large intermediate tensors of size (N, C), leading to high memory bandwidth consumption and kernel launch overhead.
|
||||
|
||||
**Optimization Strategy: Fused Block-Level Parallelism**
|
||||
|
||||
The optimization strategy fuses the entire per-sample computation into a single CUDA kernel, using a block-per-sample parallelization model.
|
||||
|
||||
1. **Parallelization Model**: The kernel is launched with a grid of `N` blocks, where `N` is the batch size. Each thread block is exclusively responsible for calculating the total loss for a single sample.
|
||||
|
||||
2. **Shared Memory for Broadcasting**: For each sample (i.e., each block), the target class index `y` and its corresponding score `x[y]` are loaded once into **shared memory**. A `__syncthreads()` call makes this data available to all threads in the block, serving as an extremely fast, localized broadcast mechanism.
|
||||
|
||||
3. **Fused Intra-Block Computation**: The threads within a block collaboratively iterate over the `C` classes. Each thread computes the hinge loss `max(0, ...)` for a subset of the classes, accumulating a partial sum in its local registers. This fuses indexing, subtraction, clamping (`max`), and power (`p`) operations.
|
||||
|
||||
4. **Efficient Intra-Block Reduction**: After processing all classes, a highly-optimized parallel reduction is performed using shared memory. The threads sum their partial sums together in a tree-like fashion, yielding the total loss for the sample in a few clock cycles.
|
||||
|
||||
5. **Finalization and Output**: The first thread of each block performs the final division by `C` and applies the class `weight` (if provided), then writes the final scalar loss for its assigned sample to the output tensor.
|
||||
|
||||
This kernel directly produces the result for `reduction='none'`. For `'mean'` and `'sum'`, a simple, fast reduction is applied to the kernel's small 1D output tensor. This approach transforms a complex, multi-stage, memory-intensive workflow into a single, efficient, compute-bound kernel pass.
|
||||
|
||||
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
BATCH_SIZE = 512
|
||||
NUM_CLASSES = 4096
|
||||
REDUCTION = 'mean'
|
||||
P = 1 # 1 for L1 hinge, 2 for L2
|
||||
MARGIN = 1.0
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
使用 PyTorch 内置的 torch.nn.MultiMarginLoss 作为基准模型。
|
||||
"""
|
||||
def __init__(self, p=1, margin=1.0, reduction='mean'):
|
||||
super(Model, self).__init__()
|
||||
# weight is not benchmarked for simplicity, but the CUDA kernel supports it.
|
||||
self.loss_fn = nn.MultiMarginLoss(p=p, margin=margin, 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, NUM_CLASSES, dtype=torch.float32)
|
||||
|
||||
target_tensor = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
|
||||
|
||||
return [input_tensor.contiguous(), target_tensor.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return [P, MARGIN, REDUCTION]
|
||||
|
|
@ -0,0 +1,88 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from multimarginloss_torch import Model, get_inputs, get_init_inputs
|
||||
from multimarginloss_cuda import ModelNew
|
||||
|
||||
def run_benchmark():
|
||||
# 检查 CUDA 是否可用
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
|
||||
return
|
||||
else:
|
||||
device = torch.device("cuda")
|
||||
|
||||
# 初始化模型
|
||||
init_inputs = get_init_inputs()
|
||||
init_inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs
|
||||
]
|
||||
inputs = get_inputs()
|
||||
inputs = [
|
||||
x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs
|
||||
]
|
||||
|
||||
torch_model = Model(*init_inputs).cuda()
|
||||
cuda_model = ModelNew(*init_inputs).cuda()
|
||||
|
||||
torch_model.eval()
|
||||
cuda_model.eval()
|
||||
|
||||
print("-------------------- 精度对齐验证 --------------------")
|
||||
with torch.no_grad():
|
||||
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 = 1000 # 增加迭代次数以获得更准确的时间测量
|
||||
|
||||
# Warm up
|
||||
for _ in range(100):
|
||||
_ = 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 MultiMarginLoss 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 MultiMarginLoss 平均执行时间: {cuda_time:.6f} 秒")
|
||||
speedup = 0
|
||||
if cuda_time > 0:
|
||||
speedup = torch_time / cuda_time
|
||||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||||
else:
|
||||
print("CUDA 内核执行时间为0，无法计算加速比。")
|
||||
return precision_flag, speedup
|
||||
|
||||
if __name__ == "__main__":
|
||||
precision_flag, speedup = run_benchmark()
|
||||
Loading…
Reference in New Issue