finish softmarginloss

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]

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()

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 <torch/extension.h>
+#include <string>
+
+torch::Tensor soft_margin_loss_cuda_forward(
+    const torch::Tensor& input, 
+    const torch::Tensor& target,
+    const std::string& reduction
+);
+"""
+
+cuda_source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+#include <cmath>
+
+#define BLOCK_SIZE 256
+
+// ----------------------------------------------------------------------------
+// __device__ function for stable SoftMarginLoss calculation
+// ----------------------------------------------------------------------------
+template <typename T>
+__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 <typename T>
+__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 <typename T>
+__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<scalar_t><<<num_blocks, BLOCK_SIZE>>>(
+                output.data_ptr<scalar_t>(),
+                input.data_ptr<scalar_t>(),
+                target.data_ptr<scalar_t>(),
+                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<scalar_t><<<num_blocks, BLOCK_SIZE>>>(
+                block_results.data_ptr<scalar_t>(),
+                input.data_ptr<scalar_t>(),
+                target.data_ptr<scalar_t>(),
+                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
+        )

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]