diff --git a/S1/35/prompt.txt b/S1/35/prompt.txt new file mode 100644 index 0000000..d0bcc70 --- /dev/null +++ b/S1/35/prompt.txt @@ -0,0 +1,64 @@ +Write a custom CUDA kernel to optimize `torch.nn.SoftMarginLoss`. + +The original operation is defined by the formula (for 'mean' reduction): +`Loss = sum_i(log(1 + exp(-y_i * x_i))) / N` +where `y_i` is either 1 or -1. + +**Problem Analysis:** +1. **Memory-Bound**: The PyTorch implementation is a chain of element-wise operations (`y * x`, `neg`, `exp`, `add`, `log`) followed by a reduction. Each step materializes a full-sized intermediate tensor in global memory, leading to excessive, wasteful memory bandwidth consumption. +2. **Kernel Launch Overhead**: Each operation in the chain launches a separate CUDA kernel, accumulating latency. +3. **Numerical Instability**: The `exp(val)` term in the formula can easily overflow for large positive `val` (where `val = -y*x`), resulting in `inf`. A robust implementation must handle this. + +**Optimization Strategy: Fused Computation with Numerical Stability and Parallel Reduction** + +The goal is to create a single, highly-efficient CUDA kernel pass that addresses all the above issues. + +1. **Full Kernel Fusion**: All element-wise operations are fused into a single calculation within the kernel. Each thread computes the final loss value for an element `i` directly from `x_i` and `y_i`, eliminating all intermediate tensors. + +2. **Numerical Stability**: The numerically unstable `log(1 + exp(x))` is replaced with a stable equivalent form inside the kernel: + - If `x > 0`, compute `x + log(1 + exp(-x))`. + - Otherwise, compute `log(1 + exp(x))`. + This ensures correct results even for large inputs and makes the custom implementation more robust than a naive one. + +3. **Fused Parallel Reduction (for 'mean'/'sum')**: For reduction modes, the element-wise calculation is fused with a classic parallel reduction algorithm: + * **Stage 1**: A kernel is launched where each thread block calculates the sum of losses for a subset of the data. This intra-block reduction is performed efficiently using **shared memory**. Each block then writes its partial sum to a small temporary buffer. + * **Stage 2**: The small buffer of partial sums is summed to get the final total loss. This can be done with a second, tiny kernel or a simple `sum()` call on the GPU, as the buffer is small. For 'mean' reduction, a final division by the element count is performed. + +This strategy transforms a multi-stage, memory-bound operation into a single-pass, compute-efficient, and numerically stable kernel, yielding substantial performance improvements. + +You are given the following architecture: +import torch +import torch.nn as nn + +BATCH_SIZE = 512 +DIM = 4096 +SHAPE = (BATCH_SIZE, DIM) +REDUCTION = 'mean' + +class Model(nn.Module): + """ + 使用 PyTorch 内置的 torch.nn.SoftMarginLoss 作为基准模型。 + """ + def __init__(self, reduction='mean'): + super(Model, self).__init__() + self.loss_fn = nn.SoftMarginLoss(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(SHAPE, dtype=torch.float32) + # target 张量必须只包含 1 和 -1 + # 使用 randint 生成 0 或 1，然后映射到 -1 或 1 + target_tensor = torch.randint(0, 2, SHAPE, dtype=torch.float32) * 2 - 1 + + return [input_tensor.contiguous(), target_tensor.contiguous()] + +def get_init_inputs(): + """ + 提供模型初始化所需的参数。 + """ + return [REDUCTION] \ No newline at end of file diff --git a/S1/35/run_code.py b/S1/35/run_code.py new file mode 100644 index 0000000..aff0428 --- /dev/null +++ b/S1/35/run_code.py @@ -0,0 +1,80 @@ +import torch +import time +from softmarginloss_torch import Model, get_inputs, get_init_inputs +from softmarginloss_cuda import ModelNew + +def run_benchmark(): + if not torch.cuda.is_available(): + print("CUDA 不可用") + return + + device = torch.device("cuda") + + # 准备输入数据 + inputs = [x.cuda(device=device) for x in get_inputs()] + init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()] + + # 初始化模型 + torch_model = Model(*init_inputs).cuda() + cuda_model = ModelNew(*init_inputs).cuda() + + torch_model.eval() + cuda_model.eval() + + print("-------------------- 精度对齐验证 --------------------") + with torch.no_grad(): + # 预热GPU + _ = torch_model(*inputs) + _ = cuda_model(*inputs) + + # 正式测试 + 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 = 100 + + # 预热GPU + for _ in range(10): + _ = 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内置SoftMarginLoss平均执行时间: {torch_time:.6f}秒") + print(f"自定义CUDA SoftMarginLoss平均执行时间: {cuda_time:.6f}秒") + speedup = torch_time / cuda_time if cuda_time > 0 else 0 + print(f"加速比 (Speedup): {speedup:.2f}x") + + return precision_flag, speedup + +if __name__ == "__main__": + precision_flag, speedup = run_benchmark() \ No newline at end of file diff --git a/S1/35/softmarginloss_cuda.py b/S1/35/softmarginloss_cuda.py new file mode 100644 index 0000000..9c3605b --- /dev/null +++ b/S1/35/softmarginloss_cuda.py @@ -0,0 +1,161 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +cpp_source = """ +#include +#include + +torch::Tensor soft_margin_loss_cuda_forward( + const torch::Tensor& input, + const torch::Tensor& target, + const std::string& reduction +); +""" + +cuda_source = """ +#include +#include +#include + +#define BLOCK_SIZE 256 + +// ---------------------------------------------------------------------------- +// __device__ function for stable SoftMarginLoss calculation +// ---------------------------------------------------------------------------- +template +__device__ __forceinline__ T stable_soft_margin_loss(T input, T target) { + T val = -target * input; + if (val > 0) { + return val + logf(1.0f + expf(-val)); + } else { + return logf(1.0f + expf(val)); + } +} + +// ---------------------------------------------------------------------------- +// Element-wise Kernel for reduction='none' +// ---------------------------------------------------------------------------- +template +__global__ void soft_margin_loss_elementwise_kernel( + T* output, + const T* input, + const T* target, + int64_t n_elements) +{ + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= n_elements) return; + output[idx] = stable_soft_margin_loss(input[idx], target[idx]); +} + +// ---------------------------------------------------------------------------- +// Reduction Kernel: Stage 1 (calculate loss and reduce within blocks) +// ---------------------------------------------------------------------------- +template +__global__ void soft_margin_loss_reduce_kernel_stage1( + T* block_results, + const T* input, + const T* target, + int64_t n_elements) +{ + __shared__ T sdata[BLOCK_SIZE]; + + int64_t tid = threadIdx.x; + int64_t i = blockIdx.x * blockDim.x + tid; + + T my_sum = 0.0f; + // Grid-stride loop to process all elements + while (i < n_elements) { + my_sum += stable_soft_margin_loss(input[i], target[i]); + i += gridDim.x * blockDim.x; + } + sdata[tid] = my_sum; + __syncthreads(); + + // Intra-block parallel reduction + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + sdata[tid] += sdata[tid + s]; + } + __syncthreads(); + } + + if (tid == 0) { + block_results[blockIdx.x] = sdata[0]; + } +} + + +torch::Tensor soft_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.sizes() == target.sizes(), "Input and target shapes must match"); + TORCH_CHECK(input.is_contiguous() && target.is_contiguous(), "Tensors must be contiguous"); + + const int64_t n_elements = input.numel(); + const auto scalar_type = input.scalar_type(); + + if (reduction == "none") { + auto output = torch::empty_like(input); + const int num_blocks = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE; + AT_DISPATCH_FLOATING_TYPES(scalar_type, "soft_margin_loss_elementwise", ([&] { + soft_margin_loss_elementwise_kernel<<>>( + output.data_ptr(), + input.data_ptr(), + target.data_ptr(), + n_elements); + })); + return output; + } + else // 'sum' or 'mean' + { + // Limit grid size to avoid creating a massive intermediate tensor + int max_grid_size = 4096; + int num_blocks = std::min(max_grid_size, (int)((n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE)); + + auto options = torch::TensorOptions().device(input.device()).dtype(input.dtype()); + auto block_results = torch::empty({num_blocks}, options); + + AT_DISPATCH_FLOATING_TYPES(scalar_type, "soft_margin_loss_reduce_stage1", ([&] { + soft_margin_loss_reduce_kernel_stage1<<>>( + block_results.data_ptr(), + input.data_ptr(), + target.data_ptr(), + n_elements); + })); + + torch::Tensor total_sum = block_results.sum(); + + if (reduction == "mean") { + return total_sum / n_elements; + } + return total_sum; + } +} +""" + +class ModelNew(nn.Module): + """ + 使用自定义 CUDA 内核进行优化的 SoftMarginLoss 模型。 + """ + def __init__(self, reduction='mean'): + super(ModelNew, self).__init__() + self.reduction = reduction + + self.soft_margin_loss_op = load_inline( + name='soft_margin_loss_op', + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=['soft_margin_loss_cuda_forward'], + verbose=False + ) + + def forward(self, input_tensor: torch.Tensor, target_tensor: torch.Tensor) -> torch.Tensor: + return self.soft_margin_loss_op.soft_margin_loss_cuda_forward( + input_tensor, + target_tensor, + self.reduction + ) \ No newline at end of file diff --git a/S1/35/softmarginloss_torch.py b/S1/35/softmarginloss_torch.py new file mode 100644 index 0000000..160c5f9 --- /dev/null +++ b/S1/35/softmarginloss_torch.py @@ -0,0 +1,35 @@ +import torch +import torch.nn as nn + +BATCH_SIZE = 512 +DIM = 4096 +SHAPE = (BATCH_SIZE, DIM) +REDUCTION = 'mean' + +class Model(nn.Module): + """ + 使用 PyTorch 内置的 torch.nn.SoftMarginLoss 作为基准模型。 + """ + def __init__(self, reduction='mean'): + super(Model, self).__init__() + self.loss_fn = nn.SoftMarginLoss(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(SHAPE, dtype=torch.float32) + # target 张量必须只包含 1 和 -1 + # 使用 randint 生成 0 或 1，然后映射到 -1 或 1 + target_tensor = torch.randint(0, 2, SHAPE, dtype=torch.float32) * 2 - 1 + + return [input_tensor.contiguous(), target_tensor.contiguous()] + +def get_init_inputs(): + """ + 提供模型初始化所需的参数。 + """ + return [REDUCTION] \ No newline at end of file