forked from ccf-ai-infra/GPUCodeForces
65 lines
3.5 KiB
Plaintext
65 lines
3.5 KiB
Plaintext
Write a custom CUDA kernel to optimize `torch.nn.SoftMarginLoss`.
|
||
|
||
The original operation is defined by the formula (for 'mean' reduction):
|
||
`Loss = sum_i(log(1 + exp(-y_i * x_i))) / N`
|
||
where `y_i` is either 1 or -1.
|
||
|
||
**Problem Analysis:**
|
||
1. **Memory-Bound**: The PyTorch implementation is a chain of element-wise operations (`y * x`, `neg`, `exp`, `add`, `log`) followed by a reduction. Each step materializes a full-sized intermediate tensor in global memory, leading to excessive, wasteful memory bandwidth consumption.
|
||
2. **Kernel Launch Overhead**: Each operation in the chain launches a separate CUDA kernel, accumulating latency.
|
||
3. **Numerical Instability**: The `exp(val)` term in the formula can easily overflow for large positive `val` (where `val = -y*x`), resulting in `inf`. A robust implementation must handle this.
|
||
|
||
**Optimization Strategy: Fused Computation with Numerical Stability and Parallel Reduction**
|
||
|
||
The goal is to create a single, highly-efficient CUDA kernel pass that addresses all the above issues.
|
||
|
||
1. **Full Kernel Fusion**: All element-wise operations are fused into a single calculation within the kernel. Each thread computes the final loss value for an element `i` directly from `x_i` and `y_i`, eliminating all intermediate tensors.
|
||
|
||
2. **Numerical Stability**: The numerically unstable `log(1 + exp(x))` is replaced with a stable equivalent form inside the kernel:
|
||
- If `x > 0`, compute `x + log(1 + exp(-x))`.
|
||
- Otherwise, compute `log(1 + exp(x))`.
|
||
This ensures correct results even for large inputs and makes the custom implementation more robust than a naive one.
|
||
|
||
3. **Fused Parallel Reduction (for 'mean'/'sum')**: For reduction modes, the element-wise calculation is fused with a classic parallel reduction algorithm:
|
||
* **Stage 1**: A kernel is launched where each thread block calculates the sum of losses for a subset of the data. This intra-block reduction is performed efficiently using **shared memory**. Each block then writes its partial sum to a small temporary buffer.
|
||
* **Stage 2**: The small buffer of partial sums is summed to get the final total loss. This can be done with a second, tiny kernel or a simple `sum()` call on the GPU, as the buffer is small. For 'mean' reduction, a final division by the element count is performed.
|
||
|
||
This strategy transforms a multi-stage, memory-bound operation into a single-pass, compute-efficient, and numerically stable kernel, yielding substantial performance improvements.
|
||
|
||
You are given the following architecture:
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
BATCH_SIZE = 512
|
||
DIM = 4096
|
||
SHAPE = (BATCH_SIZE, DIM)
|
||
REDUCTION = 'mean'
|
||
|
||
class Model(nn.Module):
|
||
"""
|
||
使用 PyTorch 内置的 torch.nn.SoftMarginLoss 作为基准模型。
|
||
"""
|
||
def __init__(self, reduction='mean'):
|
||
super(Model, self).__init__()
|
||
self.loss_fn = nn.SoftMarginLoss(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(SHAPE, dtype=torch.float32)
|
||
# target 张量必须只包含 1 和 -1
|
||
# 使用 randint 生成 0 或 1，然后映射到 -1 或 1
|
||
target_tensor = torch.randint(0, 2, SHAPE, dtype=torch.float32) * 2 - 1
|
||
|
||
return [input_tensor.contiguous(), target_tensor.contiguous()]
|
||
|
||
def get_init_inputs():
|
||
"""
|
||
提供模型初始化所需的参数。
|
||
"""
|
||
return [REDUCTION] |