From 699d514aea5f74ea1a7c858d62c047415abab581 Mon Sep 17 00:00:00 2001 From: hli28146 Date: Fri, 5 Dec 2025 00:15:11 +0800 Subject: [PATCH] finish sparsemaxloss #47 --- S1/hli28146_#47/prompt.txt | 93 +++++++++ S1/hli28146_#47/run_code.py | 74 +++++++ S1/hli28146_#47/sparsemaxloss_cuda.py | 270 +++++++++++++++++++++++++ S1/hli28146_#47/sparsemaxloss_torch.py | 64 ++++++ 4 files changed, 501 insertions(+) create mode 100644 S1/hli28146_#47/prompt.txt create mode 100644 S1/hli28146_#47/run_code.py create mode 100644 S1/hli28146_#47/sparsemaxloss_cuda.py create mode 100644 S1/hli28146_#47/sparsemaxloss_torch.py diff --git a/S1/hli28146_#47/prompt.txt b/S1/hli28146_#47/prompt.txt new file mode 100644 index 00000000..5c533104 --- /dev/null +++ b/S1/hli28146_#47/prompt.txt @@ -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'] \ No newline at end of file diff --git a/S1/hli28146_#47/run_code.py b/S1/hli28146_#47/run_code.py new file mode 100644 index 00000000..0a10a0f6 --- /dev/null +++ b/S1/hli28146_#47/run_code.py @@ -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() \ No newline at end of file diff --git a/S1/hli28146_#47/sparsemaxloss_cuda.py b/S1/hli28146_#47/sparsemaxloss_cuda.py new file mode 100644 index 00000000..42972356 --- /dev/null +++ b/S1/hli28146_#47/sparsemaxloss_cuda.py @@ -0,0 +1,270 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +cpp_source = """ +#include +#include + +torch::Tensor sparsemax_loss_cuda_forward( + const torch::Tensor& logits, + const torch::Tensor& targets, + std::string reduction); +""" + +cuda_source = """ +#include +#include +#include + +#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<<>>( + output.data_ptr(), + logits.data_ptr(), + targets.data_ptr(), + 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 + ) \ No newline at end of file diff --git a/S1/hli28146_#47/sparsemaxloss_torch.py b/S1/hli28146_#47/sparsemaxloss_torch.py new file mode 100644 index 00000000..e349e563 --- /dev/null +++ b/S1/hli28146_#47/sparsemaxloss_torch.py @@ -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'] \ No newline at end of file