From d6c100a2f55626691c607e4288f03d5e8ec71c11 Mon Sep 17 00:00:00 2001 From: hli28146 Date: Thu, 13 Nov 2025 10:58:25 +0800 Subject: [PATCH] finish multilabelmarginloss #44 --- S1/44/multilabelmarginloss_cuda.py | 156 ++++++++++++++++++++++++++++ S1/44/multilabelmarginloss_torch.py | 34 ++++++ S1/44/prompt.txt | 66 ++++++++++++ S1/44/run_code.py | 88 ++++++++++++++++ 4 files changed, 344 insertions(+) create mode 100644 S1/44/multilabelmarginloss_cuda.py create mode 100644 S1/44/multilabelmarginloss_torch.py create mode 100644 S1/44/prompt.txt create mode 100644 S1/44/run_code.py diff --git a/S1/44/multilabelmarginloss_cuda.py b/S1/44/multilabelmarginloss_cuda.py new file mode 100644 index 0000000..ebccb5d --- /dev/null +++ b/S1/44/multilabelmarginloss_cuda.py @@ -0,0 +1,156 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +cpp_source = """ +#include +#include + +// 函数前向声明 +torch::Tensor multi_label_margin_loss_cuda_forward( + const torch::Tensor& input, + const torch::Tensor& target, + const std::string& reduction +); +""" + +cuda_source = """ +#include +#include + +#define BLOCK_SIZE 256 +// 设置一个合理的单个样本最大正类标签数，用于共享内存数组 +#define MAX_POSITIVE_LABELS 32 + +template +__global__ void multi_label_margin_loss_kernel( + T* output, // (N) + const T* input, // (N, C) + const long* target, // (N, C) + const int N, + const int C) +{ + int sample_idx = blockIdx.x; + if (sample_idx >= N) return; + + __shared__ long positive_indices[MAX_POSITIVE_LABELS]; + __shared__ int num_positives; + + // 线程 0 初始化共享内存计数器 + if (threadIdx.x == 0) { + num_positives = 0; + } + __syncthreads(); + + for (int i = threadIdx.x; i < C; i += blockDim.x) { + long label = target[sample_idx * C + i]; + if (label != -1) { + int index = atomicAdd(&num_positives, 1); + if (index < MAX_POSITIVE_LABELS) { + positive_indices[index] = label; + } + } + } + __syncthreads(); + + __shared__ T sdata[BLOCK_SIZE]; + int tid = threadIdx.x; + T my_sum = 0.0f; + const T* input_row = input + sample_idx * C; + + for (int neg_class_idx = tid; neg_class_idx < C; neg_class_idx += blockDim.x) { + // 检查当前类别是否为正类 + bool is_positive = false; + for (int j = 0; j < num_positives; ++j) { + if (positive_indices[j] == neg_class_idx) { + is_positive = true; + break; + } + } + + // 如果是负类，则计算与所有正类的损失 + if (!is_positive) { + T x_neg = input_row[neg_class_idx]; + for (int j = 0; j < num_positives; ++j) { + long pos_class_idx = positive_indices[j]; + T x_pos = input_row[pos_class_idx]; + T loss_term = 1.0f - (x_pos - x_neg); + if (loss_term > 0) { + my_sum += loss_term; + } + } + } + } + sdata[tid] = my_sum; + __syncthreads(); + + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + sdata[tid] += sdata[tid + s]; + } + __syncthreads(); + } + + if (tid == 0) { + output[sample_idx] = sdata[0] / C; + } +} + +torch::Tensor multi_label_margin_loss_cuda_forward( + const torch::Tensor& input, + const torch::Tensor& target, + const std::string& reduction) +{ + TORCH_CHECK(input.is_cuda() && target.is_cuda(), "Tensors must be on CUDA"); + TORCH_CHECK(input.dim() == 2, "Input must be 2D"); + TORCH_CHECK(target.dim() == 2, "Target must be 2D"); + TORCH_CHECK(input.size(0) == target.size(0), "Batch sizes must match"); + TORCH_CHECK(input.is_contiguous() && target.is_contiguous(), "Tensors must be contiguous"); + + const int N = input.size(0); + const int C = input.size(1); + + auto options = torch::TensorOptions().device(input.device()).dtype(input.dtype()); + auto sample_losses = torch::empty({N}, options); + + dim3 grid(N); + dim3 block(BLOCK_SIZE); + + AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "multi_label_margin_loss_kernel", ([&] {{ + multi_label_margin_loss_kernel<<>>( + sample_losses.data_ptr(), + input.data_ptr(), + target.data_ptr(), + N, C + ); + }})); + + if (reduction == "none") { + return sample_losses; + } else if (reduction == "sum") { + return sample_losses.sum(); + } else {{ // "mean" + return sample_losses.mean(); + }} +} +""" + +class ModelNew(nn.Module): + def __init__(self, reduction='mean'): + super(ModelNew, self).__init__() + self.reduction = reduction + + self.op = load_inline( + name='multi_label_margin_loss_op', + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=['multi_label_margin_loss_cuda_forward'], + verbose=False + ) + + def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor: + return self.op.multi_label_margin_loss_cuda_forward( + input_tensor, + target_tensor, + self.reduction + ) \ No newline at end of file diff --git a/S1/44/multilabelmarginloss_torch.py b/S1/44/multilabelmarginloss_torch.py new file mode 100644 index 0000000..cc42cfa --- /dev/null +++ b/S1/44/multilabelmarginloss_torch.py @@ -0,0 +1,34 @@ +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] \ No newline at end of file diff --git a/S1/44/prompt.txt b/S1/44/prompt.txt new file mode 100644 index 0000000..2382961 --- /dev/null +++ b/S1/44/prompt.txt @@ -0,0 +1,66 @@ +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] \ No newline at end of file diff --git a/S1/44/run_code.py b/S1/44/run_code.py new file mode 100644 index 0000000..1dd4831 --- /dev/null +++ b/S1/44/run_code.py @@ -0,0 +1,88 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from multilabelmarginloss_torch import Model, get_inputs, get_init_inputs +from multilabelmarginloss_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) + + # 更严格的精度检查 + abs_diff = (output_torch - output_cuda).abs() + max_diff = abs_diff.max().item() + mean_diff = abs_diff.mean().item() + + print(f"最大差异: {max_diff:.6f}") + print(f"平均差异: {mean_diff:.6f}") + + precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-05, atol=1e-05) + if precision_flag: + print("✅ 精度对齐：两个模型的输出结果非常接近。") + else: + print("❌ 精度不一致！") + + print("\n-------------------- 性能加速比测试 --------------------") + num_iterations = 1000 # 增加迭代次数以获得更准确的时间测量 + + # Warm up + for _ in range(100): + _ = torch_model(*inputs) + _ = cuda_model(*inputs) + + # 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 MultiLabelMarginLoss 平均执行时间: {torch_time:.6f} 秒") + print(f"自定义 CUDA MultiLabelMarginLoss 平均执行时间: {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