diff --git a/S1/hli28146_#55/GDL_cuda.py b/S1/hli28146_#55/GDL_cuda.py new file mode 100644 index 00000000..27683021 --- /dev/null +++ b/S1/hli28146_#55/GDL_cuda.py @@ -0,0 +1,183 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +cpp_source = """ +#include +#include + +torch::Tensor gdl_loss_cuda_forward( + const torch::Tensor& logits, + const torch::Tensor& targets, + float eps, + std::string reduction); +""" + +cuda_source = """ +#include +#include +#include +#include + +#define BLOCK_SIZE 256 +// For GDL, C is usually small (Brain: 4-5 classes) +#define MAX_CLASSES 32 + +struct __align__(16) Float4 { + float x, y, z, w; +}; + +// --- Fused GDL Kernel --- +// Grid: (B, 1, 1), Block: (256, 1, 1) +// Since C is small, we can keep C accumulators in registers/shared mem per thread. +__global__ void gdl_loss_kernel( + float* __restrict__ output, + const float* __restrict__ logits, + const int64_t* __restrict__ targets, + int C, + int spatial_dim, + float eps) +{ + int batch_idx = blockIdx.x; + int tid = threadIdx.x; + + float acc_r[MAX_CLASSES]; + float acc_inter[MAX_CLASSES]; + float acc_union[MAX_CLASSES]; + + #pragma unroll + for (int c = 0; c < MAX_CLASSES; ++c) { + acc_r[c] = 0.0f; + acc_inter[c] = 0.0f; + acc_union[c] = 0.0f; + } + + // Offset for this batch + const float* b_logits = logits + batch_idx * C * spatial_dim; + const int64_t* b_targets = targets + batch_idx * spatial_dim; + + // Grid-Stride Loop over spatial pixels + for (int i = tid; i < spatial_dim; i += blockDim.x) { + // Read Target + int64_t target_label = b_targets[i]; + + // Read Logits & Compute Softmax for this pixel + float local_logits[MAX_CLASSES]; + float max_l = -FLT_MAX; + + for (int c = 0; c < C; ++c) { + float val = b_logits[c * spatial_dim + i]; + local_logits[c] = val; + if (val > max_l) max_l = val; + } + + float sum_exp = 0.0f; + for (int c = 0; c < C; ++c) { + float e = __expf(local_logits[c] - max_l); + local_logits[c] = e; // Store exp temporarily + sum_exp += e; + } + + // 3. Accumulate stats + float inv_sum = 1.0f / sum_exp; + + for (int c = 0; c < C; ++c) { + float p = local_logits[c] * inv_sum; // prob + float r = (c == target_label) ? 1.0f : 0.0f; // one-hot + + acc_r[c] += r; + acc_inter[c] += r * p; + acc_union[c] += r + p; + } + } + + // Block Reduction + __shared__ float s_stats[3 * MAX_CLASSES]; + if (tid < 3 * C) { + s_stats[tid] = 0.0f; + } + __syncthreads(); + + for (int c = 0; c < C; ++c) { + atomicAdd(&s_stats[0 * C + c], acc_r[c]); + atomicAdd(&s_stats[1 * C + c], acc_inter[c]); + atomicAdd(&s_stats[2 * C + c], acc_union[c]); + } + __syncthreads(); + + if (tid == 0) { + float num = 0.0f; + float den = 0.0f; + + for (int c = 0; c < C; ++c) { + float sum_r = s_stats[0 * C + c]; + float sum_inter = s_stats[1 * C + c]; + float sum_union = s_stats[2 * C + c]; + + float w = 1.0f / (sum_r * sum_r + eps); + + num += w * sum_inter; + den += w * sum_union; + } + + float dice = (2.0f * num) / (den + eps); + output[batch_idx] = 1.0f - dice; + } +} + +torch::Tensor gdl_loss_cuda_forward( + const torch::Tensor& logits, + const torch::Tensor& targets, + float eps, + std::string reduction) +{ + TORCH_CHECK(logits.is_cuda() && targets.is_cuda(), "Inputs must be CUDA"); + TORCH_CHECK(logits.is_contiguous() && targets.is_contiguous(), "Inputs must be contiguous"); + + // (B, C, Spatial) + int batch_size = logits.size(0); + int num_classes = logits.size(1); + + // Total spatial pixels + int spatial_dim = 1; + for (int i=2; i>>( + output.data_ptr(), + logits.data_ptr(), + targets.data_ptr(), + num_classes, + spatial_dim, + eps + ); + + if (reduction == "mean") { + return output.mean(); + } else if (reduction == "sum") { + return output.sum(); + } + + return output; +} +""" + +class ModelNew(nn.Module): + def __init__(self, eps=1e-6, reduction='none'): + super(ModelNew, self).__init__() + self.eps = eps + self.reduction = reduction + self.op = load_inline( + name='gdl_loss_op', + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=['gdl_loss_cuda_forward'], + verbose=False, + extra_cuda_cflags=['-O3'] + ) + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + return self.op.gdl_loss_cuda_forward(logits.contiguous(), targets.contiguous(), self.eps, self.reduction) \ No newline at end of file diff --git a/S1/hli28146_#55/GDL_torch.py b/S1/hli28146_#55/GDL_torch.py new file mode 100644 index 00000000..dea99f06 --- /dev/null +++ b/S1/hli28146_#55/GDL_torch.py @@ -0,0 +1,76 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +BATCH_SIZE = 16 +NUM_CLASSES = 4 +DEPTH = 32 +HEIGHT = 128 +WIDTH = 128 +SPATIAL_DIM = DEPTH * HEIGHT * WIDTH +SHAPE_LOGITS = (BATCH_SIZE, NUM_CLASSES, DEPTH, HEIGHT, WIDTH) +SHAPE_TARGET = (BATCH_SIZE, DEPTH, HEIGHT, WIDTH) + +EPS = 1e-6 +REDUCTION = 'none' + +class GeneralizedDiceLoss(nn.Module): + """ + Generalized Dice Loss (Sudre et al. MICCAI 2017) + https://arxiv.org/pdf/1707.03237 + """ + def __init__(self, eps=1e-6, reduction='mean'): + super(GeneralizedDiceLoss, self).__init__() + self.eps = eps + self.reduction = reduction + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + # logits: (B, C, D, H, W) + # targets: (B, D, H, W) -> indices + probs = F.softmax(logits, dim=1) + + # One-hot encoding + targets_onehot = F.one_hot(targets.long(), num_classes=logits.shape[1]) + targets_onehot = targets_onehot.permute(0, 4, 1, 2, 3).float() + + # Compute weights w_l = 1 / (sum(r_ln)^2) + sum_r = targets_onehot.sum(dim=(2, 3, 4)) # (B, C) + weights = 1.0 / (sum_r * sum_r + self.eps) + + # Compute Intersection & Union + # Intersection: r * p + intersection = (targets_onehot * probs).sum(dim=(2, 3, 4)) # (B, C) + # Union: r + p + union = (targets_onehot + probs).sum(dim=(2, 3, 4)) # (B, C) + + # Weighted Sum + # Numerator: 2 * sum_l (w_l * inter_l) + # Denominator: sum_l (w_l * union_l) + + numerator = 2.0 * (weights * intersection).sum(dim=1) + denominator = (weights * union).sum(dim=1) + + dice_score = numerator / (denominator + self.eps) + loss = 1.0 - dice_score + + if self.reduction == 'mean': + return loss.mean() + elif self.reduction == 'sum': + return loss.sum() + return loss + +class Model(nn.Module): + def __init__(self, eps=1e-6, reduction='none'): + super(Model, self).__init__() + self.loss_fn = GeneralizedDiceLoss(eps=eps, reduction=reduction) + + def forward(self, logits, targets): + return self.loss_fn(logits, targets) + +def get_inputs(): + logits = torch.randn(SHAPE_LOGITS, dtype=torch.float32) + targets = torch.randint(0, NUM_CLASSES, SHAPE_TARGET, dtype=torch.long) + return [logits.contiguous(), targets.contiguous()] + +def get_init_inputs(): + return [EPS, REDUCTION] \ No newline at end of file diff --git a/S1/hli28146_#55/prompt.txt b/S1/hli28146_#55/prompt.txt new file mode 100644 index 00000000..63e38cf3 --- /dev/null +++ b/S1/hli28146_#55/prompt.txt @@ -0,0 +1,122 @@ +Write a custom CUDA kernel to optimize `Generalized Dice Loss` (GDL). + +Formula: +GDL = 1 - 2 * (Sum_l w_l * Sum_n (r_ln * p_ln)) / (Sum_l w_l * Sum_n (r_ln + p_ln)) +Where: +- l is class index, n is spatial index (pixels/voxels). +- p_ln is softmax probability. +- r_ln is one-hot ground truth. +- w_l = 1 / (Sum_n r_ln)^2. + +Problem Analysis: +1. Memory Bandwidth: Standard implementation involves Softmax, One-hot generation, Summation per class for weights, and then weighted intersection/union sums. This requires multiple passes over the large (N, C, Spatial) tensors. +2. Intermediate Storage: Storing probability maps and one-hot targets consumes significant memory. + +Optimization Strategy: Fused Softmax-Reduction Kernel + +1. One-Block-per-Sample: Each block processes one image/volume in the batch. + +2. On-the-fly Calculation: + - Compute Softmax probabilities `p` from logits on-the-fly using cached Max/SumExp. + - Generate `r` (one-hot) from target indices on-the-fly. + +3. Fused Accumulation: + - Iterate over all spatial pixels `n` and classes `l`. + - Maintain 3 accumulators per class in Shared Memory/Registers: + - `sum_r`: sum(r_ln) [for weights] + - `sum_inter`: sum(r_ln * p_ln) + - `sum_union`: sum(r_ln + p_ln) + + Wait, `w_l` depends on `sum_r` across the WHOLE spatial dimension. So we must accumulate `sum_r` for all pixels first? + Yes, standard GDL weights depend on the GT volume. + However, `sum_inter` and `sum_union` also require summation over `n`. + We can accumulate `sum_r`, `sum_inter`, `sum_union` simultaneously in one pass over spatial dimensions. + +4. Shared Memory Reduction: + - Use atomicAdd or tree reduction in Shared Memory to aggregate these sums for each class across threads. + +5. Final Composition: + - Thread 0 reads the aggregated per-class sums. + - Computes `w_l = 1 / (sum_r^2 + eps)`. + - Computes numerator `2 * sum(w_l * sum_inter)` and denominator `sum(w_l * sum_union)`. + - Writes the final loss per sample. + +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 torch.nn.functional as F + +BATCH_SIZE = 16 +NUM_CLASSES = 4 +DEPTH = 32 +HEIGHT = 128 +WIDTH = 128 +SPATIAL_DIM = DEPTH * HEIGHT * WIDTH +SHAPE_LOGITS = (BATCH_SIZE, NUM_CLASSES, DEPTH, HEIGHT, WIDTH) +SHAPE_TARGET = (BATCH_SIZE, DEPTH, HEIGHT, WIDTH) + +EPS = 1e-6 +REDUCTION = 'none' + +class GeneralizedDiceLoss(nn.Module): + """ + Generalized Dice Loss (Sudre et al. MICCAI 2017) + https://arxiv.org/pdf/1707.03237 + """ + def __init__(self, eps=1e-6, reduction='mean'): + super(GeneralizedDiceLoss, self).__init__() + self.eps = eps + self.reduction = reduction + + def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: + # logits: (B, C, D, H, W) + # targets: (B, D, H, W) -> indices + probs = F.softmax(logits, dim=1) + + # One-hot encoding + targets_onehot = F.one_hot(targets.long(), num_classes=logits.shape[1]) + targets_onehot = targets_onehot.permute(0, 4, 1, 2, 3).float() + + # Compute weights w_l = 1 / (sum(r_ln)^2) + sum_r = targets_onehot.sum(dim=(2, 3, 4)) # (B, C) + weights = 1.0 / (sum_r * sum_r + self.eps) + + # Compute Intersection & Union + # Intersection: r * p + intersection = (targets_onehot * probs).sum(dim=(2, 3, 4)) # (B, C) + # Union: r + p + union = (targets_onehot + probs).sum(dim=(2, 3, 4)) # (B, C) + + # Weighted Sum + # Numerator: 2 * sum_l (w_l * inter_l) + # Denominator: sum_l (w_l * union_l) + + numerator = 2.0 * (weights * intersection).sum(dim=1) + denominator = (weights * union).sum(dim=1) + + dice_score = numerator / (denominator + self.eps) + loss = 1.0 - dice_score + + if self.reduction == 'mean': + return loss.mean() + elif self.reduction == 'sum': + return loss.sum() + return loss + +class Model(nn.Module): + def __init__(self, eps=1e-6, reduction='none'): + super(Model, self).__init__() + self.loss_fn = GeneralizedDiceLoss(eps=eps, reduction=reduction) + + def forward(self, logits, targets): + return self.loss_fn(logits, targets) + +def get_inputs(): + logits = torch.randn(SHAPE_LOGITS, dtype=torch.float32) + targets = torch.randint(0, NUM_CLASSES, SHAPE_TARGET, dtype=torch.long) + return [logits.contiguous(), targets.contiguous()] + +def get_init_inputs(): + return [EPS, REDUCTION] \ No newline at end of file diff --git a/S1/hli28146_#55/run_code.py b/S1/hli28146_#55/run_code.py new file mode 100644 index 00000000..29550a33 --- /dev/null +++ b/S1/hli28146_#55/run_code.py @@ -0,0 +1,74 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from GDL_torch import Model,get_inputs,get_init_inputs +from GDL_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