forked from ccf-ai-infra/GPUCodeForces
66 lines
4.4 KiB
Plaintext
66 lines
4.4 KiB
Plaintext
Write a custom CUDA kernel to optimize `torch.nn.MultiLabelMarginLoss`.
|
||
|
||
The original operation is defined by the formula:
|
||
`loss(x, y) = sum_{j,i} max(0, 1 - (x[y[j]] - x[i])) / (x.size(0) * x.size(1))`
|
||
where `y[j]` are the positive class indices for a sample and `i` are the negative class indices. This is computed per sample and then reduced.
|
||
|
||
**Problem Analysis:**
|
||
`MultiLabelMarginLoss` is notoriously difficult to vectorize efficiently in PyTorch. The performance bottlenecks are severe:
|
||
1. **Irregular Data Access**: Each sample has a variable number of positive labels defined in `y`, which are padded with -1. Identifying the set of positive and negative classes for each sample requires complex, non-vectorized logic (e.g., masks, loops, or boolean indexing), which is slow.
|
||
2. **Massive Intermediate Tensors**: A naive vectorized approach would require gathering scores for positive classes and broadcasting them for subtraction against scores of negative classes. This would create huge intermediate tensors and is highly memory-inefficient.
|
||
3. **Complex Nested Loop Logic**: The core formula is a nested loop (`for each positive class`, `for each negative class`) for every sample, which is antithetical to efficient GPU execution without a custom kernel.
|
||
|
||
**Optimization Strategy: Fused Block-Level Parallelism with Shared Memory Caching**
|
||
|
||
The strategy is to implement the entire complex logic within a single CUDA kernel, using a block-per-sample parallelization model.
|
||
|
||
1. **Parallelization Model**: A grid of `N` blocks is launched, where `N` is the batch size. Each thread block is assigned to compute the total loss for one sample.
|
||
|
||
2. **Shared Memory Caching**: For each sample (block), the kernel first collaboratively reads the list of positive class indices from the `target` tensor. These indices (and their count) are cached in **shared memory**. This makes the critical metadata for the sample instantly accessible to all threads in the block.
|
||
|
||
3. **Fused Computation Loop**: The threads within a block then work together to iterate through all `C` possible classes. For each class `i`, a thread checks if it's a positive or negative class using the cached shared memory data.
|
||
* If `i` is a negative class, the thread then iterates through the *positive class indices cached in shared memory*.
|
||
* For each positive-negative pair, it calculates the hinge loss term `max(0, 1 - (x_pos - x_neg))` and accumulates it into a thread-local register. This fuses the nested loops, indexing, subtraction, and `max` operations.
|
||
|
||
4. **Efficient Intra-Block Reduction**: Once all classes are processed, a fast parallel reduction is performed using shared memory to sum the partial results from all threads within the block into a single total loss for that sample.
|
||
|
||
5. **Finalization**: The first thread of each block performs the final division and writes the result to the output tensor. The kernel directly produces the per-sample losses (`reduction='none'`). The final batch reduction (`'mean'` or `'sum'`) is efficiently handled by a single PyTorch call on the small 1D output tensor.
|
||
|
||
This approach transforms the complex, memory-bound, and hard-to-vectorize PyTorch operation into a single, efficient, compute-focused CUDA kernel.
|
||
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
|
||
import numpy as np
|
||
|
||
BATCH_SIZE = 512
|
||
NUM_CLASSES = 1024
|
||
REDUCTION = 'mean'
|
||
# 每个样本的正类标签数量范围
|
||
MIN_LABELS = 1
|
||
MAX_LABELS = 10
|
||
|
||
class Model(nn.Module):
|
||
def __init__(self, reduction='mean'):
|
||
super(Model, self).__init__()
|
||
self.loss_fn = nn.MultiLabelMarginLoss(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每行包含正类索引，并用 -1 填充
|
||
target_np = np.full((BATCH_SIZE, NUM_CLASSES), -1, dtype=np.int64)
|
||
for i in range(BATCH_SIZE):
|
||
num_labels = np.random.randint(MIN_LABELS, MAX_LABELS + 1)
|
||
labels = np.random.choice(NUM_CLASSES, num_labels, replace=False)
|
||
target_np[i, :num_labels] = labels
|
||
target_tensor = torch.from_numpy(target_np)
|
||
|
||
return [input_tensor.contiguous(), target_tensor.contiguous()]
|
||
|
||
def get_init_inputs():
|
||
return [REDUCTION] |