forked from ccf-ai-infra/GPUCodeForces
Merge pull request 'finish sparsemaxloss #47' (#375) from hli28146/GPUCodeForces:h47 into main
This commit is contained in:
commit
020794b324
|
|
@ -0,0 +1,93 @@
|
|||
Write a custom CUDA kernel to optimize `Sparsemax Loss` (ICML 2016).
|
||||
|
||||
Formula: L = 0.5 * sum_{j in Support} (z_j^2 - tau^2) + 0.5 - z_target
|
||||
Algorithm to find Support and tau:
|
||||
1. Sort logits z in descending order.
|
||||
2. Find largest k such that 1 + k * z_k > sum(z_1...z_k).
|
||||
3. tau = (sum(z_1...z_k) - 1) / k.
|
||||
4. Support set is indices where z_j > tau.
|
||||
|
||||
Problem Analysis:
|
||||
1. Sorting Overhead: The standard implementation uses `torch.sort`, which operates in global memory and is expensive for the subsequent logic flow.
|
||||
2. Memory Traffic: Calculating cumsum and masks after sorting requires multiple passes over global memory tensors.
|
||||
|
||||
Optimization Strategy: Fused Shared-Memory Sort & Reduction
|
||||
|
||||
Constraint: Assume `num_classes` is a power of 2 (e.g., 2048) to facilitate efficient Bitonic Sort.
|
||||
|
||||
1. Block-per-Row: Launch one block per sample.
|
||||
2. Shared Memory Loading: Load the entire row of logits into Shared Memory.
|
||||
3. Bitonic Sort (Descending): Implement parallel Bitonic Sort in Shared Memory to order the logits. This avoids global memory sorting.
|
||||
4. Parallel Scan (Cumsum): Compute the prefix sum of the sorted logits in Shared Memory to evaluate the condition `1 + k * z_k > cumsum_k`.
|
||||
5. Threshold Detection: Identify the threshold index `k` and compute `tau`.
|
||||
6. Fused Loss Calculation:
|
||||
- Calculate sum of squares for the top-k elements (using reduction).
|
||||
- Calculate final loss using the pre-loaded target logit (read from global memory initially).
|
||||
|
||||
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 = 2048
|
||||
NUM_CLASSES = 2048
|
||||
SHAPE = (BATCH_SIZE, NUM_CLASSES)
|
||||
|
||||
class SparsemaxLoss(nn.Module):
|
||||
"""
|
||||
Sparsemax Loss (Martins & Astudillo, 2016)
|
||||
L = 0.5 * sum(z_j^2 - tau^2) + 0.5 - z_y
|
||||
"""
|
||||
def __init__(self, reduction='mean'):
|
||||
super(SparsemaxLoss, self).__init__()
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
# logits: (N, C)
|
||||
# targets: (N)
|
||||
|
||||
# Sort (Descending)
|
||||
z_sorted, _ = torch.sort(logits, dim=1, descending=True)
|
||||
|
||||
z_cumsum = torch.cumsum(z_sorted, dim=1)
|
||||
|
||||
k = torch.arange(1, logits.size(1) + 1, device=logits.device)
|
||||
|
||||
support = (1 + k * z_sorted) > z_cumsum
|
||||
k_z = torch.sum(support, dim=1, keepdim=True) # (N, 1)
|
||||
|
||||
zs_sum = torch.gather(z_cumsum, 1, k_z - 1)
|
||||
tau = (zs_sum - 1) / k_z
|
||||
|
||||
mask = torch.arange(NUM_CLASSES, device=logits.device).unsqueeze(0) < k_z
|
||||
z_support = z_sorted * mask
|
||||
|
||||
sum_sq_z = (z_support ** 2).sum(dim=1)
|
||||
sum_sq_tau = (tau.squeeze(1) ** 2) * k_z.squeeze(1).float()
|
||||
|
||||
z_y = logits.gather(1, targets.unsqueeze(1)).squeeze(1)
|
||||
|
||||
loss = 0.5 * (sum_sq_z - sum_sq_tau) + 0.5 - z_y
|
||||
|
||||
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 = SparsemaxLoss(reduction=reduction)
|
||||
|
||||
def forward(self, logits, targets):
|
||||
return self.loss_fn(logits, targets)
|
||||
|
||||
def get_inputs():
|
||||
logits = torch.randn(SHAPE, dtype=torch.float32)
|
||||
targets = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
|
||||
return [logits.contiguous(), targets.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return ['none']
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
###########################################################
|
||||
# 性能和精度验证程序
|
||||
###########################################################
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
from sparsemaxloss_torch import Model,get_inputs,get_init_inputs
|
||||
from sparsemaxloss_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()
|
||||
|
|
@ -0,0 +1,270 @@
|
|||
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 sparsemax_loss_cuda_forward(
|
||||
const torch::Tensor& logits,
|
||||
const torch::Tensor& targets,
|
||||
std::string reduction);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
#define BLOCK_SIZE 1024
|
||||
#define NUM_ELEM 2048
|
||||
|
||||
__device__ inline void swap(float& a, float& b) {
|
||||
float tmp = a; a = b; b = tmp;
|
||||
}
|
||||
|
||||
__global__ void sparsemax_loss_kernel(
|
||||
float* __restrict__ output,
|
||||
const float* __restrict__ logits,
|
||||
const int64_t* __restrict__ targets,
|
||||
int cols)
|
||||
{
|
||||
// 1. Load Logits into Shared Memory
|
||||
__shared__ float s_val[NUM_ELEM];
|
||||
|
||||
int row_idx = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
const float* row_logits = logits + row_idx * cols;
|
||||
|
||||
// Retrieve Target Logit early (before we mess up indices or sort)
|
||||
int64_t target_idx = targets[row_idx];
|
||||
float z_y = row_logits[target_idx]; // Global read
|
||||
|
||||
// Load 2 elements per thread
|
||||
int idx1 = tid;
|
||||
int idx2 = tid + BLOCK_SIZE;
|
||||
|
||||
s_val[idx1] = row_logits[idx1];
|
||||
s_val[idx2] = row_logits[idx2];
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// 2. Bitonic Sort (Descending)
|
||||
for (int size = 2; size <= NUM_ELEM; size <<= 1) {
|
||||
// Bitonic Merge
|
||||
// Descending order: means we want largest first.
|
||||
// dir = ( (tid & (size / 2)) == 0 ) check is for alternating up/down
|
||||
// But for the final full merge, we want one direction.
|
||||
// Bitonic sort produces a monotonic sequence only at the very end.
|
||||
|
||||
for (int stride = size / 2; stride > 0; stride >>= 1) {
|
||||
__syncthreads();
|
||||
// Emulate 2048 threads with 1024 threads loop
|
||||
// Effective thread ID mapping for bitonic network
|
||||
|
||||
// Algorithm:
|
||||
// For each pair (pos, pos+stride)
|
||||
// Logic for "Standard" Bitonic Sort implementation (iterative):
|
||||
// There are NUM_ELEM / 2 comparators. We have 1024 threads. Perfect match.
|
||||
|
||||
int pos = 2 * tid - (tid & (stride - 1));
|
||||
int partner = pos + stride;
|
||||
|
||||
if (partner < NUM_ELEM) {
|
||||
float a = s_val[pos];
|
||||
float b = s_val[partner];
|
||||
|
||||
// Direction logic:
|
||||
// Full sort descending: goal is for final stage to be descending
|
||||
// The XOR trick determines direction for sub-blocks
|
||||
bool sort_descending = ((pos & size) == 0);
|
||||
|
||||
// If size == NUM_ELEM, we force direction to be Descending (or Ascending depending on what we want)
|
||||
// Actually for standard bitonic sort, the direction flag flips.
|
||||
// To get a fully Descending array:
|
||||
// We essentially run standard sort but invert compare.
|
||||
|
||||
// Let's keep it simple: Standard Bitonic creates Ascending.
|
||||
// To get Descending, we swap logic.
|
||||
if (sort_descending) {
|
||||
if (a < b) { s_val[pos] = b; s_val[partner] = a; }
|
||||
} else {
|
||||
if (a > b) { s_val[pos] = b; s_val[partner] = a; }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Now s_val is Sorted Descending: z_(1) >= z_(2) ... >= z_(K)
|
||||
|
||||
// 3. Parallel Prefix Sum (Scan) - Inclusive
|
||||
// We need cumsum to check condition: 1 + k * z_k > cumsum_k
|
||||
// Use Hillis-Steele double buffering
|
||||
__shared__ float s_sum[2][NUM_ELEM];
|
||||
|
||||
// Init scan buffer
|
||||
s_sum[0][idx1] = s_val[idx1];
|
||||
s_sum[0][idx2] = s_val[idx2];
|
||||
__syncthreads();
|
||||
|
||||
int in_buf = 0;
|
||||
int out_buf = 1;
|
||||
|
||||
for (int stride = 1; stride < NUM_ELEM; stride <<= 1) {
|
||||
__syncthreads(); // barrier between steps
|
||||
|
||||
// Process idx1
|
||||
if (idx1 >= stride)
|
||||
s_sum[out_buf][idx1] = s_sum[in_buf][idx1] + s_sum[in_buf][idx1 - stride];
|
||||
else
|
||||
s_sum[out_buf][idx1] = s_sum[in_buf][idx1];
|
||||
|
||||
// Process idx2
|
||||
if (idx2 >= stride)
|
||||
s_sum[out_buf][idx2] = s_sum[in_buf][idx2] + s_sum[in_buf][idx2 - stride];
|
||||
else
|
||||
s_sum[out_buf][idx2] = s_sum[in_buf][idx2];
|
||||
|
||||
// Swap
|
||||
in_buf = 1 - in_buf;
|
||||
out_buf = 1 - out_buf;
|
||||
}
|
||||
__syncthreads();
|
||||
// Result is in in_buf
|
||||
|
||||
// 4. Find Threshold k
|
||||
// Condition: 1 + k * z_k > cumsum_k
|
||||
// k is 1-based index (1..C). Array is 0-based (0..C-1).
|
||||
// So for index i: k = i + 1.
|
||||
// Cond: 1 + (i+1) * s_val[i] > s_sum[in_buf][i]
|
||||
|
||||
// We need to find the LARGEST i satisfying this.
|
||||
// Since z is sorted, this property is monotonic.
|
||||
// We can simply count how many elements satisfy this.
|
||||
|
||||
int satisfy1 = (1.0f + (float)(idx1 + 1) * s_val[idx1] > s_sum[in_buf][idx1]) ? 1 : 0;
|
||||
int satisfy2 = (1.0f + (float)(idx2 + 1) * s_val[idx2] > s_sum[in_buf][idx2]) ? 1 : 0;
|
||||
|
||||
// Let's put satisfy counts into s_sum[0] and reduce
|
||||
s_sum[0][idx1] = (float)satisfy1;
|
||||
s_sum[0][idx2] = (float)satisfy2;
|
||||
__syncthreads();
|
||||
|
||||
// Tree reduction for K
|
||||
for (int s = NUM_ELEM / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) {
|
||||
// Each thread sums 2 nodes, but stride handling needs care for > BLOCK_SIZE
|
||||
// Our threads cover 0..1023. Total 2048.
|
||||
// Standard reduction:
|
||||
// Iter 1: s=1024. tid 0..1023. Add [tid] and [tid+1024].
|
||||
// Iter 2: s=512. tid 0..511. Add [tid] and [tid+512].
|
||||
|
||||
// Note: Initial mapping was s_sum[idx1] and s_sum[idx2] where idx2 = idx1 + 1024.
|
||||
// So step 1 is just:
|
||||
s_sum[0][tid] += s_sum[0][tid + s];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Now s_sum[0][0] holds k(z)
|
||||
__shared__ float k_z_val;
|
||||
__shared__ float tau;
|
||||
|
||||
if (tid == 0) {
|
||||
k_z_val = s_sum[0][0];
|
||||
// tau = (cumsum[k-1] - 1) / k
|
||||
int k_idx = (int)k_z_val - 1;
|
||||
// Retrieve cumsum from buffer. buffer index is in_buf
|
||||
float cumsum_val = s_sum[in_buf][k_idx];
|
||||
tau = (cumsum_val - 1.0f) / k_z_val;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// 5. Calculate Loss
|
||||
// L = 0.5 * sum_{j in S} (z_j^2 - tau^2) + 0.5 - z_y
|
||||
// S is indices 0 to k-1
|
||||
|
||||
// Each thread calculates z^2 - tau^2 for its elements IF they are in support
|
||||
float local_loss_part = 0.0f;
|
||||
float t = tau;
|
||||
int k_limit = (int)k_z_val;
|
||||
|
||||
if (idx1 < k_limit) {
|
||||
float z = s_val[idx1];
|
||||
local_loss_part += (z * z - t * t);
|
||||
}
|
||||
if (idx2 < k_limit) {
|
||||
float z = s_val[idx2];
|
||||
local_loss_part += (z * z - t * t);
|
||||
}
|
||||
|
||||
// Reduce loss parts
|
||||
s_sum[0][idx1] = local_loss_part; // Reusing buffer 0
|
||||
s_sum[0][idx2] = 0.0f; // Clear second slot (since reduction below assumes sum of tid and tid+s)
|
||||
// Actually, better: store local sum in s_sum[0][tid] = local_loss_part (which includes idx1 and idx2)
|
||||
__syncthreads();
|
||||
|
||||
s_sum[0][tid] = local_loss_part;
|
||||
__syncthreads();
|
||||
|
||||
// Reduction
|
||||
for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) {
|
||||
s_sum[0][tid] += s_sum[0][tid + s];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
float support_term = s_sum[0][0];
|
||||
output[row_idx] = 0.5f * support_term + 0.5f - z_y;
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor sparsemax_loss_cuda_forward(
|
||||
const torch::Tensor& logits,
|
||||
const torch::Tensor& targets,
|
||||
std::string reduction)
|
||||
{
|
||||
TORCH_CHECK(logits.is_cuda() && targets.is_cuda(), "Inputs must be CUDA");
|
||||
TORCH_CHECK(logits.is_contiguous(), "Logits must be contiguous");
|
||||
TORCH_CHECK(logits.size(1) == NUM_ELEM, "Kernel optimized for 2048 classes");
|
||||
|
||||
int batch_size = logits.size(0);
|
||||
auto output = torch::empty({batch_size}, logits.options());
|
||||
|
||||
sparsemax_loss_kernel<<<batch_size, BLOCK_SIZE>>>(
|
||||
output.data_ptr<float>(),
|
||||
logits.data_ptr<float>(),
|
||||
targets.data_ptr<int64_t>(),
|
||||
NUM_ELEM
|
||||
);
|
||||
|
||||
if (reduction == "mean") return output.mean();
|
||||
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='sparsemax_loss_op',
|
||||
cpp_sources=cpp_source,
|
||||
cuda_sources=cuda_source,
|
||||
functions=['sparsemax_loss_cuda_forward'],
|
||||
verbose=False,
|
||||
extra_cuda_cflags=['-O3']
|
||||
)
|
||||
|
||||
def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
return self.op.sparsemax_loss_cuda_forward(
|
||||
logits.contiguous(),
|
||||
targets.contiguous(),
|
||||
self.reduction
|
||||
)
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
BATCH_SIZE = 2048
|
||||
NUM_CLASSES = 2048
|
||||
SHAPE = (BATCH_SIZE, NUM_CLASSES)
|
||||
|
||||
class SparsemaxLoss(nn.Module):
|
||||
"""
|
||||
Sparsemax Loss (Martins & Astudillo, 2016)
|
||||
L = 0.5 * sum(z_j^2 - tau^2) + 0.5 - z_y
|
||||
"""
|
||||
def __init__(self, reduction='mean'):
|
||||
super(SparsemaxLoss, self).__init__()
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
||||
# logits: (N, C)
|
||||
# targets: (N)
|
||||
|
||||
# Sort (Descending)
|
||||
z_sorted, _ = torch.sort(logits, dim=1, descending=True)
|
||||
|
||||
z_cumsum = torch.cumsum(z_sorted, dim=1)
|
||||
|
||||
k = torch.arange(1, logits.size(1) + 1, device=logits.device)
|
||||
|
||||
support = (1 + k * z_sorted) > z_cumsum
|
||||
k_z = torch.sum(support, dim=1, keepdim=True) # (N, 1)
|
||||
|
||||
zs_sum = torch.gather(z_cumsum, 1, k_z - 1)
|
||||
tau = (zs_sum - 1) / k_z
|
||||
|
||||
mask = torch.arange(NUM_CLASSES, device=logits.device).unsqueeze(0) < k_z
|
||||
z_support = z_sorted * mask
|
||||
|
||||
sum_sq_z = (z_support ** 2).sum(dim=1)
|
||||
sum_sq_tau = (tau.squeeze(1) ** 2) * k_z.squeeze(1).float()
|
||||
|
||||
z_y = logits.gather(1, targets.unsqueeze(1)).squeeze(1)
|
||||
|
||||
loss = 0.5 * (sum_sq_z - sum_sq_tau) + 0.5 - z_y
|
||||
|
||||
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 = SparsemaxLoss(reduction=reduction)
|
||||
|
||||
def forward(self, logits, targets):
|
||||
return self.loss_fn(logits, targets)
|
||||
|
||||
def get_inputs():
|
||||
logits = torch.randn(SHAPE, dtype=torch.float32)
|
||||
targets = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
|
||||
return [logits.contiguous(), targets.contiguous()]
|
||||
|
||||
def get_init_inputs():
|
||||
return ['none']
|
||||
Loading…
Reference in New Issue