forked from ccf-ai-infra/GPUCodeForces
finish ListMLELoss #48
This commit is contained in:
parent
f876a28ada
commit
6446db5e43
|
|
@ -0,0 +1,241 @@
|
|||
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 listmle_loss_cuda_forward(
|
||||
const torch::Tensor& scores,
|
||||
const torch::Tensor& labels,
|
||||
std::string reduction);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
#define BLOCK_SIZE 512
|
||||
#define LIST_SIZE 1024 // Fixed for simplicity
|
||||
|
||||
struct Item {
|
||||
float score;
|
||||
float label;
|
||||
};
|
||||
|
||||
template<typename T>
|
||||
__device__ __forceinline__ T warp_reduce_sum(T val) {
|
||||
#pragma unroll
|
||||
for (int offset = 16; offset > 0; offset /= 2) {
|
||||
val += __shfl_down_sync(0xffffffff, val, offset);
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float block_reduce_sum(float val) {
|
||||
static __shared__ float shared[16]; // 512/32 = 16 warps
|
||||
int lane = threadIdx.x % 32;
|
||||
int wid = threadIdx.x / 32;
|
||||
val = warp_reduce_sum(val);
|
||||
if (lane == 0) shared[wid] = val;
|
||||
__syncthreads();
|
||||
val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;
|
||||
if (wid == 0) val = warp_reduce_sum(val);
|
||||
return val;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float block_reduce_max(float val) {
|
||||
// Simplification: assumes val initialized to -inf
|
||||
// Not implemented fully for brevity, standard pattern
|
||||
// We can use a loop for Max easily since data is in Shared Mem
|
||||
return 0.0f; // Placeholder, logic implemented inline
|
||||
}
|
||||
|
||||
__global__ void listmle_kernel(
|
||||
float* __restrict__ output,
|
||||
const float* __restrict__ scores,
|
||||
const float* __restrict__ labels,
|
||||
int num_lists)
|
||||
{
|
||||
// Shared Memory for Sorting
|
||||
__shared__ Item s_data[LIST_SIZE];
|
||||
|
||||
int row_idx = blockIdx.x;
|
||||
if (row_idx >= num_lists) return;
|
||||
|
||||
int tid = threadIdx.x;
|
||||
|
||||
// 1. Load Data (2 items per thread)
|
||||
const float* row_scores = scores + row_idx * LIST_SIZE;
|
||||
const float* row_labels = labels + row_idx * LIST_SIZE;
|
||||
|
||||
int idx1 = tid;
|
||||
int idx2 = tid + BLOCK_SIZE;
|
||||
|
||||
s_data[idx1] = {row_scores[idx1], row_labels[idx1]};
|
||||
s_data[idx2] = {row_scores[idx2], row_labels[idx2]};
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// 2. Bitonic Sort (Descending by Label)
|
||||
for (int size = 2; size <= LIST_SIZE; size <<= 1) {
|
||||
for (int stride = size / 2; stride > 0; stride >>= 1) {
|
||||
__syncthreads();
|
||||
|
||||
// Thread handles comparators
|
||||
int pos = 2 * tid - (tid & (stride - 1));
|
||||
int partner = pos + stride;
|
||||
|
||||
// Standard Bitonic direction:
|
||||
// desc if ((pos & size) == 0)
|
||||
// We want full descending -> invert logic or final stage
|
||||
// Let's just use monotonic logic: sort blocks Descending/Ascending
|
||||
|
||||
bool sort_desc = ((pos & size) == 0);
|
||||
|
||||
Item a = s_data[pos];
|
||||
Item b = s_data[partner];
|
||||
|
||||
bool swap = false;
|
||||
if (sort_desc) {
|
||||
if (a.label < b.label) swap = true;
|
||||
} else {
|
||||
if (a.label > b.label) swap = true;
|
||||
}
|
||||
|
||||
if (swap) {
|
||||
s_data[pos] = b;
|
||||
s_data[partner] = a;
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// 3. Find Max Score
|
||||
float my_max = fmaxf(s_data[idx1].score, s_data[idx2].score);
|
||||
|
||||
// Simple Block Reduce Max
|
||||
static __shared__ float s_max_buffer[16];
|
||||
float warp_max = my_max;
|
||||
for (int off=16; off>0; off/=2) warp_max = fmaxf(warp_max, __shfl_down_sync(0xffffffff, warp_max, off));
|
||||
if ((tid % 32) == 0) s_max_buffer[tid/32] = warp_max;
|
||||
__syncthreads();
|
||||
if (tid < 16) {
|
||||
warp_max = s_max_buffer[tid];
|
||||
for (int off=8; off>0; off/=2) warp_max = fmaxf(warp_max, __shfl_down_sync(0xffffffff, warp_max, off));
|
||||
if (tid == 0) s_max_buffer[0] = warp_max;
|
||||
}
|
||||
__syncthreads();
|
||||
float global_max = s_max_buffer[0];
|
||||
|
||||
// 4. Compute Exp (in place or separate buffer? Reuse s_data label field for exp value to save memory?)
|
||||
// s_data[i].label is not needed anymore.
|
||||
double exp1 = exp((double)(s_data[idx1].score - global_max));
|
||||
double exp2 = exp((double)(s_data[idx2].score - global_max));
|
||||
|
||||
// Reuse shared memory for Suffix Scan
|
||||
// But s_data is struct. Let's just cast pointer or use label field.
|
||||
// Using label field (float) for exp value.
|
||||
s_data[idx1].label = (float)exp1;
|
||||
s_data[idx2].label = (float)exp2;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// 5. Parallel Suffix Scan (Reverse Inclusive Scan)
|
||||
__shared__ double scan_data[2][LIST_SIZE]; // Double precision for sum
|
||||
|
||||
scan_data[0][idx1] = exp1;
|
||||
scan_data[0][idx2] = exp2;
|
||||
|
||||
int in = 0;
|
||||
int out = 1;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Suffix Scan: out[i] = in[i] + in[i + stride]
|
||||
for (int stride = 1; stride < LIST_SIZE; stride *= 2) {
|
||||
__syncthreads();
|
||||
|
||||
// idx1
|
||||
if (idx1 + stride < LIST_SIZE)
|
||||
scan_data[out][idx1] = scan_data[in][idx1] + scan_data[in][idx1 + stride];
|
||||
else
|
||||
scan_data[out][idx1] = scan_data[in][idx1];
|
||||
|
||||
// idx2
|
||||
if (idx2 + stride < LIST_SIZE)
|
||||
scan_data[out][idx2] = scan_data[in][idx2] + scan_data[in][idx2 + stride];
|
||||
else
|
||||
scan_data[out][idx2] = scan_data[in][idx2];
|
||||
|
||||
int temp = in; in = out; out = temp;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Now scan_data[in] holds suffix sums
|
||||
|
||||
// 6. Compute Loss
|
||||
// term = log(suffix) + M - z_i
|
||||
double l1 = log(scan_data[in][idx1] + 1e-10) + global_max - s_data[idx1].score;
|
||||
double l2 = log(scan_data[in][idx2] + 1e-10) + global_max - s_data[idx2].score;
|
||||
|
||||
float local_loss = (float)(l1 + l2);
|
||||
|
||||
// 7. Reduce Loss
|
||||
float total_loss = block_reduce_sum(local_loss);
|
||||
|
||||
if (tid == 0) {
|
||||
output[row_idx] = total_loss;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor listmle_loss_cuda_forward(
|
||||
const torch::Tensor& scores,
|
||||
const torch::Tensor& labels,
|
||||
std::string reduction)
|
||||
{
|
||||
TORCH_CHECK(scores.is_cuda() && labels.is_cuda(), "Inputs must be CUDA");
|
||||
TORCH_CHECK(scores.is_contiguous() && labels.is_contiguous(), "Inputs must be contiguous");
|
||||
TORCH_CHECK(scores.size(1) == LIST_SIZE, "Fixed list size 1024");
|
||||
|
||||
int batch_size = scores.size(0);
|
||||
auto output = torch::empty({batch_size}, scores.options());
|
||||
|
||||
listmle_kernel<<<batch_size, BLOCK_SIZE>>>(
|
||||
output.data_ptr<float>(),
|
||||
scores.data_ptr<float>(),
|
||||
labels.data_ptr<float>(),
|
||||
batch_size
|
||||
);
|
||||
|
||||
if (reduction == "mean") {
|
||||
return output.mean();
|
||||
} else if (reduction == "sum") {
|
||||
return output.sum();
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
"""
|
||||
|
||||
class ModelNew(nn.Module):
|
||||
def __init__(self, reduction='none'):
|
||||
super(ModelNew, self).__init__()
|
||||
self.reduction = reduction
|
||||
self.op = load_inline(
|
||||
name='listmle_op',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['listmle_loss_cuda_forward'],
|
||||
verbose=False,
|
||||
extra_cuda_cflags=['-O3']
|
||||
)
|
||||
|
||||
def forward(self, scores: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
||||
return self.op.listmle_loss_cuda_forward(
|
||||
scores.contiguous(),
|
||||
labels.contiguous(),
|
||||
self.reduction
|
||||
)
|
||||
|
|
@ -0,0 +1,59 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
BATCH_SIZE = 4096
|
||||
LIST_SIZE = 1024
|
||||
SHAPE = (BATCH_SIZE, LIST_SIZE)
|
||||
|
||||
class ListMLELoss(nn.Module):
|
||||
def __init__(self, reduction='mean'):
|
||||
super(ListMLELoss, self).__init__()
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, scores: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
||||
# scores: (B, N)
|
||||
# labels: (B, N)
|
||||
|
||||
sorted_indices = torch.argsort(labels, dim=1, descending=True)
|
||||
|
||||
sorted_scores = torch.gather(scores, 1, sorted_indices)
|
||||
|
||||
max_val, _ = sorted_scores.max(dim=1, keepdim=True)
|
||||
sorted_scores_stable = sorted_scores - max_val
|
||||
|
||||
exp_scores = torch.exp(sorted_scores_stable)
|
||||
|
||||
exp_sum_reverse = torch.cumsum(torch.flip(exp_scores, [1]), dim=1)
|
||||
exp_sum_reverse = torch.flip(exp_sum_reverse, [1])
|
||||
|
||||
log_cumsum = torch.log(exp_sum_reverse + 1e-10) # eps
|
||||
|
||||
# Restore scale: log(sum(e^(x-m))) = log(sum) + m
|
||||
# Loss term_i = (log_cumsum_i + M) - sorted_score_i
|
||||
loss_per_item = (log_cumsum + max_val) - sorted_scores
|
||||
|
||||
loss = loss_per_item.sum(dim=1)
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
return loss
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, reduction='none'):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = ListMLELoss(reduction=reduction)
|
||||
|
||||
def forward(self, scores, labels):
|
||||
return self.loss_fn(scores, labels)
|
||||
|
||||
def get_inputs():
|
||||
scores = torch.randn(SHAPE, dtype=torch.float32)
|
||||
labels = torch.randint(0, 5, SHAPE, dtype=torch.float32)
|
||||
labels += torch.rand(SHAPE, dtype=torch.float32) * 1e-3
|
||||
|
||||
return [scores.contiguous(), labels.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return ['none']
|
||||
|
|
@ -0,0 +1,94 @@
|
|||
Write a custom CUDA kernel to optimize `ListMLE Loss`.
|
||||
|
||||
Formula: L = Sum_{i=0}^{N-1} [ log( Sum_{k=i}^{N-1} exp(z_pi(k)) ) - z_pi(i) ]
|
||||
Where `z` are predicted scores, and `pi` is the permutation that sorts the ground truth labels in descending order.
|
||||
Essentially: Sort scores based on labels -> Compute LogSumExp of the suffix -> Subtract score -> Sum.
|
||||
|
||||
Problem Analysis:
|
||||
1. Sorting Bottleneck: The standard implementation requires `torch.argsort` on labels for every sample in the batch, which is computationally expensive and memory-intensive.
|
||||
2. Memory Traffic: Calculating the suffix sum (denominator) involves multiple passes: exp, flip, cumsum, flip, log.
|
||||
|
||||
Optimization Strategy: Fused Sort-Scan Kernel in Shared Memory
|
||||
|
||||
Constraint: List size `N` is fixed to a power of 2 (e.g., 1024) for efficient Bitonic Sort.
|
||||
|
||||
1. Block-per-Query: Assign one CUDA block to process one query (list of items).
|
||||
|
||||
2. Shared Memory Staging: Load scores and labels into Shared Memory structures.
|
||||
|
||||
3. Parallel Bitonic Sort:
|
||||
- Sort the data in Shared Memory based on **labels** in descending order.
|
||||
- Use `score` as the payload moved along with labels.
|
||||
- This replaces the global `argsort` + `gather` pattern.
|
||||
|
||||
4. Numerical Stability & Suffix Scan:
|
||||
- Find the max score `M` in the sorted list for stability.
|
||||
- Compute `exp(score - M)`.
|
||||
- Perform a **Reverse Parallel Scan** (Suffix Sum) in Shared Memory to calculate `CumSumExp_i = sum_{k=i}^{N-1} exp(...)`.
|
||||
|
||||
5. Fused Loss:
|
||||
- `Loss_i = (log(CumSumExp_i) + M) - score_i`.
|
||||
- Sum `Loss_i` across the block to get total loss for the query.
|
||||
|
||||
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 = 4096
|
||||
LIST_SIZE = 1024
|
||||
SHAPE = (BATCH_SIZE, LIST_SIZE)
|
||||
|
||||
class ListMLELoss(nn.Module):
|
||||
def __init__(self, reduction='mean'):
|
||||
super(ListMLELoss, self).__init__()
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, scores: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
||||
# scores: (B, N)
|
||||
# labels: (B, N)
|
||||
|
||||
sorted_indices = torch.argsort(labels, dim=1, descending=True)
|
||||
|
||||
sorted_scores = torch.gather(scores, 1, sorted_indices)
|
||||
|
||||
max_val, _ = sorted_scores.max(dim=1, keepdim=True)
|
||||
sorted_scores_stable = sorted_scores - max_val
|
||||
|
||||
exp_scores = torch.exp(sorted_scores_stable)
|
||||
|
||||
exp_sum_reverse = torch.cumsum(torch.flip(exp_scores, [1]), dim=1)
|
||||
exp_sum_reverse = torch.flip(exp_sum_reverse, [1])
|
||||
|
||||
log_cumsum = torch.log(exp_sum_reverse + 1e-10) # eps
|
||||
|
||||
# Restore scale: log(sum(e^(x-m))) = log(sum) + m
|
||||
# Loss term_i = (log_cumsum_i + M) - sorted_score_i
|
||||
loss_per_item = (log_cumsum + max_val) - sorted_scores
|
||||
|
||||
loss = loss_per_item.sum(dim=1)
|
||||
|
||||
if self.reduction == 'mean':
|
||||
return loss.mean()
|
||||
elif self.reduction == 'sum':
|
||||
return loss.sum()
|
||||
return loss
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, reduction='none'):
|
||||
super(Model, self).__init__()
|
||||
self.loss_fn = ListMLELoss(reduction=reduction)
|
||||
|
||||
def forward(self, scores, labels):
|
||||
return self.loss_fn(scores, labels)
|
||||
|
||||
def get_inputs():
|
||||
scores = torch.randn(SHAPE, dtype=torch.float32)
|
||||
labels = torch.randint(0, 5, SHAPE, dtype=torch.float32)
|
||||
labels += torch.rand(SHAPE, dtype=torch.float32) * 1e-3
|
||||
|
||||
return [scores.contiguous(), labels.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return ['none']
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from mleloss_torch import Model,get_inputs,get_init_inputs
|
||||
from mleloss_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)
|
||||
|
||||
precision_flag = torch.allclose(output_torch, output_cuda,rtol=1e-03)
|
||||
if precision_flag:
|
||||
print("✅ 精度对齐：两个模型的输出结果非常接近。")
|
||||
else:
|
||||
print("❌ 精度不一致！")
|
||||
|
||||
print("\n-------------------- 性能加速比测试 --------------------")
|
||||
num_iterations = 100
|
||||
|
||||
# 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 torch.relu 平均执行时间: {torch_time:.6f} 秒")
|
||||
print(f"自定义 CUDA 内核 平均执行时间: {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