feat:add elu+batchnorm #73

S1/wut0n_#73/elu_batchnorm_cudacode.py

@@ -0,0 +1,200 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+elu_batchnorm_source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+#include <math_constants.h>
+
+// ELU + BatchNorm融合内核
+__global__ void elu_batchnorm_kernel(
+    const float* __restrict__ x,
+    const float* __restrict__ weight,
+    const float* __restrict__ bias,
+    const float* __restrict__ running_mean,
+    const float* __restrict__ running_var,
+    float* __restrict__ output,
+    float alpha,
+    float eps,
+    int total_elements,
+    int channels,
+    int spatial_size
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    
+    if (idx >= total_elements) return;
+    
+    // 计算通道索引
+    int channel_idx = (idx / spatial_size) % channels;
+    
+    // 读取输入值
+    float x_val = x[idx];
+    
+    // 应用BatchNorm
+    float normalized = (x_val - running_mean[channel_idx]) / sqrtf(running_var[channel_idx] + eps);
+    float batchnorm_output = normalized * weight[channel_idx] + bias[channel_idx];
+    
+    // 应用ELU
+    if (batchnorm_output > 0.0f) {
+        output[idx] = batchnorm_output;
+    } else {
+        output[idx] = alpha * (expf(batchnorm_output) - 1.0f);
+    }
+}
+
+// 向量化ELU + BatchNorm融合内核
+__global__ void elu_batchnorm_vectorized_kernel(
+    const float* __restrict__ x,
+    const float* __restrict__ weight,
+    const float* __restrict__ bias,
+    const float* __restrict__ running_mean,
+    const float* __restrict__ running_var,
+    float* __restrict__ output,
+    float alpha,
+    float eps,
+    int total_elements,
+    int channels,
+    int spatial_size
+) {
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int vec_idx = idx * 4;
+    
+    if (vec_idx + 3 >= total_elements) return;
+    
+    // 读取4个元素
+    float4 x_vec = *reinterpret_cast<const float4*>(&x[vec_idx]);
+    
+    // 处理4个元素
+    float4 result;
+    
+    // 元素1
+    int channel_idx1 = (vec_idx / spatial_size) % channels;
+    float normalized1 = (x_vec.x - running_mean[channel_idx1]) / sqrtf(running_var[channel_idx1] + eps);
+    float batchnorm1 = normalized1 * weight[channel_idx1] + bias[channel_idx1];
+    result.x = (batchnorm1 > 0.0f) ? batchnorm1 : alpha * (expf(batchnorm1) - 1.0f);
+    
+    // 元素2
+    int channel_idx2 = ((vec_idx + 1) / spatial_size) % channels;
+    float normalized2 = (x_vec.y - running_mean[channel_idx2]) / sqrtf(running_var[channel_idx2] + eps);
+    float batchnorm2 = normalized2 * weight[channel_idx2] + bias[channel_idx2];
+    result.y = (batchnorm2 > 0.0f) ? batchnorm2 : alpha * (expf(batchnorm2) - 1.0f);
+    
+    // 元素3
+    int channel_idx3 = ((vec_idx + 2) / spatial_size) % channels;
+    float normalized3 = (x_vec.z - running_mean[channel_idx3]) / sqrtf(running_var[channel_idx3] + eps);
+    float batchnorm3 = normalized3 * weight[channel_idx3] + bias[channel_idx3];
+    result.z = (batchnorm3 > 0.0f) ? batchnorm3 : alpha * (expf(batchnorm3) - 1.0f);
+    
+    // 元素4
+    int channel_idx4 = ((vec_idx + 3) / spatial_size) % channels;
+    float normalized4 = (x_vec.w - running_mean[channel_idx4]) / sqrtf(running_var[channel_idx4] + eps);
+    float batchnorm4 = normalized4 * weight[channel_idx4] + bias[channel_idx4];
+    result.w = (batchnorm4 > 0.0f) ? batchnorm4 : alpha * (expf(batchnorm4) - 1.0f);
+    
+    // 写入结果
+    *reinterpret_cast<float4*>(&output[vec_idx]) = result;
+}
+
+torch::Tensor elu_batchnorm_cuda(
+    torch::Tensor x,
+    torch::Tensor weight,
+    torch::Tensor bias,
+    torch::Tensor running_mean,
+    torch::Tensor running_var,
+    float alpha = 1.0f,
+    float eps = 1e-5f
+) {
+    auto x_contig = x.contiguous();
+    
+    int batch_size = x_contig.size(0);
+    int channels = x_contig.size(1);
+    int height = x_contig.size(2);
+    int width = x_contig.size(3);
+    int spatial_size = height * width;
+    int total_elements = batch_size * channels * spatial_size;
+    
+    auto output = torch::zeros_like(x_contig);
+    
+    // 向量化处理：最佳平衡点
+    int vec_elements = total_elements & ~3;  // 向下取整到4的倍数
+    int scalar_elements = total_elements & 3;  // 剩余元素
+    
+    if (vec_elements > 0) {
+        const int block_size = 256;
+        const int grid_size = (vec_elements / 4 + block_size - 1) / block_size;
+        
+        elu_batchnorm_vectorized_kernel<<<grid_size, block_size>>>(
+            x_contig.data_ptr<float>(),
+            weight.data_ptr<float>(),
+            bias.data_ptr<float>(),
+            running_mean.data_ptr<float>(),
+            running_var.data_ptr<float>(),
+            output.data_ptr<float>(),
+            alpha,
+            eps,
+            total_elements,
+            channels,
+            spatial_size
+        );
+    }
+    
+    // 处理剩余元素
+    if (scalar_elements > 0) {
+        const int block_size = 256;
+        const int grid_size = (scalar_elements + block_size - 1) / block_size;
+        
+        elu_batchnorm_kernel<<<grid_size, block_size>>>(
+            x_contig.data_ptr<float>() + vec_elements,
+            weight.data_ptr<float>(),
+            bias.data_ptr<float>(),
+            running_mean.data_ptr<float>(),
+            running_var.data_ptr<float>(),
+            output.data_ptr<float>() + vec_elements,
+            alpha,
+            eps,
+            total_elements,
+            channels,
+            spatial_size
+        );
+    }
+    
+    return output;
+}
+"""
+
+elu_batchnorm_cpp_source = """
+torch::Tensor elu_batchnorm_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor running_mean, torch::Tensor running_var, float alpha, float eps);
+"""
+
+# 编译CUDA代码
+elu_batchnorm = load_inline(
+    name="elu_batchnorm",
+    cpp_sources=elu_batchnorm_cpp_source,
+    cuda_sources=elu_batchnorm_source,
+    functions=["elu_batchnorm_cuda"],
+    extra_cuda_cflags=[
+        "-O3", 
+        "--use_fast_math",
+        "-std=c++17"
+    ],
+    verbose=True
+)
+
+class ModelNew(torch.nn.Module):
+    def __init__(self, num_features, alpha=1.0, eps=1e-5):
+        super(ModelNew, self).__init__()
+        self.num_features = num_features
+        self.alpha = alpha
+        self.eps = eps
+        self.elu_batchnorm = elu_batchnorm
+        
+        # BatchNorm参数
+        self.weight = torch.nn.Parameter(torch.ones(num_features))
+        self.bias = torch.nn.Parameter(torch.zeros(num_features))
+        self.register_buffer('running_mean', torch.zeros(num_features))
+        self.register_buffer('running_var', torch.ones(num_features))
+
+    def forward(self, x):
+        return self.elu_batchnorm.elu_batchnorm_cuda(
+            x, self.weight, self.bias, self.running_mean, self.running_var, self.alpha, self.eps
+        )

S1/wut0n_#73/elu_batchnorm_torchcode.py

@@ -0,0 +1,47 @@
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    """
+    ELU + BatchNorm融合实现。
+    先应用BatchNorm，然后应用ELU激活。
+    """
+    def __init__(self, num_features, alpha=1.0, eps=1e-5):
+        super(Model, self).__init__()
+        self.num_features = num_features
+        self.alpha = alpha
+        self.eps = eps
+        
+        # BatchNorm层
+        self.batchnorm = nn.BatchNorm2d(num_features, eps=eps)
+        
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Apply BatchNorm + ELU fusion.
+
+        Args:
+            x (torch.Tensor): Input tensor [batch_size, channels, height, width]
+
+        Returns:
+            torch.Tensor: BatchNorm + ELU output [batch_size, channels, height, width]
+        """
+        # 先应用BatchNorm
+        batchnorm_output = self.batchnorm(x)
+        
+        # 再应用ELU
+        elu_output = torch.where(batchnorm_output > 0, batchnorm_output, self.alpha * (torch.exp(batchnorm_output) - 1))
+        
+        return elu_output
+
+batch_size = 256
+channels = 64
+height = 32
+width = 32
+
+def get_inputs():
+    # Generate input tensor with mixed positive and negative values
+    x = torch.randn(batch_size, channels, height, width) * 2.0
+    return [x]
+
+def get_init_inputs():
+    return [channels]  # num_features

S1/wut0n_#73/prompt.txt

@@ -0,0 +1,120 @@
+You write custom CUDA kernels to replace pytorch operators in the given architecture to get speedups.
+
+You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.
+
+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
+
+class Model(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, a, b):
+        return a + b
+
+def get_inputs():
+    # randomly generate input tensors based on the model architecture
+    a = torch.randn(1, 128).cuda()
+    b = torch.randn(1, 128).cuda()
+    return [a, b]
+
+def get_init_inputs():
+    # randomly generate tensors required for initialization based on the model architecture
+    return []
+
+The example new arch with custom CUDA kernels looks like this:
+
+python
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ModelNew(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, a, b):
+        return a + b
+
+def get_inputs():
+    # randomly generate input tensors based on the model architecture
+    a = torch.randn(1, 128).cuda()
+    b = torch.randn(1, 128).cuda()
+    return [a, b]
+
+def get_init_inputs():
+    # randomly generate tensors required for initialization based on the model architecture
+    return []
+
+You are given the following architecture:
+
+python
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    """
+    ELU (Exponential Linear Unit) implementation.
+    Applies the ELU activation function element-wise.
+    """
+    def __init__(self, alpha=1.0):
+        super(Model, self).__init__()
+        self.alpha = alpha
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Apply ELU activation to input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor [batch_size, channels, height, width]
+
+        Returns:
+            torch.Tensor: ELU-activated tensor [batch_size, channels, height, width]
+        """
+        # Apply ELU: x if x > 0 else alpha * (exp(x) - 1)
+        return torch.where(x > 0, x, self.alpha * (torch.exp(x) - 1))
+
+batch_size = 128
+channels = 64
+height = 32
+width = 32
+
+def get_inputs():
+    # Generate input tensor with mixed positive and negative values
+    x = torch.randn(batch_size, channels, height, width) * 2.0
+    return [x]
+
+def get_init_inputs():
+    return []  # No special initialization inputs needed
+
+IMPORTANT: The ELU activation function involves conditional branching and exponential computation that can be significantly optimized through CUDA vectorization and branch-free techniques. Consider implementing multiple optimization strategies: 1) Standard implementation with proper branching, 2) Vectorized processing using float4 for memory efficiency, 3) Branch-free implementations to avoid warp divergence, and 4) Fast exponential approximations for performance-critical scenarios. Focus on achieving both numerical stability and significant speedups while maintaining precision.
+
+The implementation should include:
+1. A vectorized kernel using float4 for processing 4 elements simultaneously
+2. A scalar kernel for handling remaining elements
+3. Proper memory coalescing and efficient memory access patterns
+4. Branch optimization techniques to minimize warp divergence
+5. Fast math optimizations where appropriate while maintaining numerical accuracy
+6. Comprehensive error checking and input validation
+7. Proper CUDA kernel launch configuration with optimal block and grid sizes
+8. Support for different alpha values and tensor shapes
+9. Integration with PyTorch's tensor system and automatic differentiation
+10. Performance benchmarking against the original PyTorch implementation
+
+The CUDA implementation should achieve significant speedups (3-5x or more) while maintaining numerical precision and stability. Consider using techniques like:
+- Memory coalescing with float4 vectorization
+- Shared memory usage when beneficial
+- Warp-level primitives for reduction operations
+- Fast exponential approximations (e.g., polynomial approximations)
+- Branch-free implementations using mathematical tricks
+- Proper handling of edge cases and numerical stability
+- Optimized kernel launch parameters
+- Compilation flags for maximum performance
+
+Generate the complete CUDA implementation with both cudacode.py and torchcode.py files that can be directly used for benchmarking and integration.

S1/wut0n_#73/run_code.py

@@ -0,0 +1,74 @@
+###########################################################
+# 性能和精度验证程序
+###########################################################
+import torch
+import torch.nn as nn
+import time
+from elu_batchnorm_torchcode import Model, get_inputs, get_init_inputs
+from elu_batchnorm_cudacode 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 elu_batchnorm 平均执行时间: {torch_time:.6f} 秒")
+    print(f"自定义 CUDA elu_batchnorm 平均执行时间: {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()