Merge pull request 'FEAT:ADD chebyshev_hardswish #105' (#872) from wut0n/GPUCodeForces:chebyshev_hardswish into main

S1/wut0n_#105/chebyshev_hardswish_cudacode.py

@@ -0,0 +1,108 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+chebyshev_hardswish_source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+
+// 多线程版本 - 每个样本由多个线程并行处理，融合了Hard-Swish激活
+__global__ void chebyshev_hardswish_kernel(
+    const float* __restrict__ x,
+    const float* __restrict__ y,
+    float* __restrict__ activated_distances,
+    int batch_size,
+    int feature_dim
+) {
+    // 每个block处理一个样本
+    int sample_idx = blockIdx.x;
+    if (sample_idx >= batch_size) return;
+    
+    int tid = threadIdx.x;
+    int base = sample_idx * feature_dim;
+    
+    // 使用共享内存进行归约
+    extern __shared__ float shared_max[];
+    shared_max[tid] = 0.0f;
+    
+    // 每个线程处理多个元素
+    int stride = blockDim.x;
+    for (int dim = tid; dim < feature_dim; dim += stride) {
+        float diff = x[base + dim] - y[base + dim];
+        float abs_diff = fabsf(diff);
+        if (abs_diff > shared_max[tid]) {
+            shared_max[tid] = abs_diff;
+        }
+    }
+    
+    __syncthreads();
+    
+    // 块内归约找最大值
+    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+        if (tid < stride) {
+            if (shared_max[tid + stride] > shared_max[tid]) {
+                shared_max[tid] = shared_max[tid + stride];
+            }
+        }
+        __syncthreads();
+    }
+    
+    // 第一个线程进行最终归约并应用Hard-Swish激活
+    if (tid == 0) {
+        float chebyshev_dist = shared_max[0];
+        
+        // --- 融合计算Hard-Swish激活函数 ---
+        // 使用条件运算符避免分支分歧，提高效率
+        float hardswish_val = (chebyshev_dist <= -3.0f) ? 0.0f : 
+                              ((chebyshev_dist >= 3.0f) ? chebyshev_dist : 
+                               chebyshev_dist * (chebyshev_dist + 3.0f) / 6.0f);
+        
+        activated_distances[sample_idx] = hardswish_val;
+    }
+}
+
+torch::Tensor chebyshev_hardswish_cuda(torch::Tensor x, torch::Tensor y) {
+    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "X must be float32");
+    TORCH_CHECK(y.scalar_type() == torch::kFloat32, "Y must be float32");
+    
+    auto x_contig = x.contiguous();
+    auto y_contig = y.contiguous();
+    
+    int batch_size = x_contig.size(0);
+    int feature_dim = x_contig.size(1);
+    
+    auto activated_distances = torch::zeros({batch_size}, x.options());
+    
+    const int block_size = 256;
+    size_t shared_mem = block_size * sizeof(float);
+    
+    chebyshev_hardswish_kernel<<<batch_size, block_size, shared_mem>>>(
+        x_contig.data_ptr<float>(),
+        y_contig.data_ptr<float>(),
+        activated_distances.data_ptr<float>(),
+        batch_size,
+        feature_dim
+    );
+    
+    return activated_distances;
+}
+"""
+
+chebyshev_hardswish_cpp_source = """
+torch::Tensor chebyshev_hardswish_cuda(torch::Tensor x, torch::Tensor y);
+"""
+
+chebyshev_hardswish = load_inline(
+    name="chebyshev_hardswish",
+    cpp_sources=chebyshev_hardswish_cpp_source,
+    cuda_sources=chebyshev_hardswish_source,
+    functions=["chebyshev_hardswish_cuda"],
+    verbose=True
+)
+
+class ModelNew(torch.nn.Module):
+    def __init__(self):
+        super(ModelNew, self).__init__()
+        self.chebyshev_hardswish = chebyshev_hardswish  # The module containing the kernel
+
+    def forward(self, x, y):
+        return self.chebyshev_hardswish.chebyshev_hardswish_cuda(x, y)

S1/wut0n_#105/chebyshev_hardswish_torchcode.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    """
+    PyTorch基准实现：切比雪夫距离 + Hard-Swish激活
+    """
+    def __init__(self):
+        super(Model, self).__init__()
+    
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        """
+        Compute Chebyshev distance between x and y, then apply Hard-Swish activation.
+
+        Args:
+            x (torch.Tensor): First set of vectors [batch_size, feature_dim]
+            y (torch.Tensor): Second set of vectors [batch_size, feature_dim]
+
+        Returns:
+            torch.Tensor: Hard-Swish-activated Chebyshev distances [batch_size]
+        """
+        # Input validation
+        if x.shape != y.shape:
+            raise ValueError(f"Input tensors must have the same shape, got {x.shape} and {y.shape}")
+        
+        if x.dim() != 2:
+            raise ValueError(f"Input tensors must be 2D, got {x.dim()}D")
+        
+        # --- 第一步：计算切比雪夫距离 ---
+        # max(|x_i - y_i|)
+        chebyshev_dist = torch.max(torch.abs(x - y), dim=1)[0]
+        
+        # --- 第二步：应用Hard-Swish激活函数 ---
+        # 使用F.hardswish，这是PyTorch内置的高效实现
+        activated_distances = torch.nn.functional.hardswish(chebyshev_dist)
+        
+        return activated_distances
+
+batch_size = 256
+feature_dim = 512
+
+def get_inputs():
+    # Generate two sets of vectors
+    x = torch.randn(batch_size, feature_dim)
+    y = torch.randn(batch_size, feature_dim)
+    return [x, y]
+
+def get_init_inputs():
+    return []  # No special initialization inputs needed

S1/wut0n_#105/prompt.txt

@@ -0,0 +1,147 @@
+You write custom CUDA kernels to replace pytorch operators in 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 PyTorch. The example given architecture is a simple addition:
+
+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():
+a = torch.randn(1, 128).cuda()
+b = torch.randn(1, 128).cuda()
+return [a, b]
+
+def get_init_inputs():
+return []
+
+
+
+The example new architecture with a custom CUDA kernel looks like this:
+
+python
+import torch
+from torch.utils.cpp_extension import load_inline
+
+add_source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+
+global void add_kernel(const float* a, const float* b, float* out, int size) {
+int idx = blockIdx.x * blockDim.x + threadIdx.x;
+if (idx < size) {
+out[idx] = a[idx] + b[idx];
+}
+}
+
+torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b) {
+auto out = torch::empty_like(a);
+int size = a.numel();
+const int block_size = 256;
+int num_blocks = (size + block_size - 1) / block_size;
+add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);
+return out;
+}
+"""
+
+add_cpp_source = """
+torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b);
+"""
+
+Compile the inline CUDA code
+add = load_inline(
+name="add",
+cpp_sources=add_cpp_source,
+cuda_sources=add_source,
+functions=["add_cuda"],
+verbose=True
+)
+
+class ModelNew(torch.nn.Module):
+def init(self):
+super(ModelNew, self).init()
+self.add = add
+
+def forward(self, a, b):
+    return self.add.add_cuda(a, b)
+
+
+---
+
+Now, you are given the following PyTorch architecture to accelerate. The model computes the Chebyshev distance (maximum absolute difference) between two sets of vectors and then applies the Hard-Swish activation function to each distance. This baseline implementation is efficient and uses PyTorch's highly optimized built-in functions for correctness and performance.
+
+python
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    """
+    PyTorch基准实现：切比雪夫距离 + Hard-Swish激活
+    """
+    def __init__(self):
+        super(Model, self).__init__()
+    
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        """
+        Compute Chebyshev distance between x and y, then apply Hard-Swish activation.
+
+        Args:
+            x (torch.Tensor): First set of vectors [batch_size, feature_dim]
+            y (torch.Tensor): Second set of vectors [batch_size, feature_dim]
+
+        Returns:
+            torch.Tensor: Hard-Swish-activated Chebyshev distances [batch_size]
+        """
+        # Input validation
+        if x.shape != y.shape:
+            raise ValueError(f"Input tensors must have the same shape, got {x.shape} and {y.shape}")
+        
+        if x.dim() != 2:
+            raise ValueError(f"Input tensors must be 2D, got {x.dim()}D")
+        
+        # --- 第一步：计算切比雪夫距离 ---
+        # max(|x_i - y_i|)
+        chebyshev_dist = torch.max(torch.abs(x - y), dim=1)[0]
+        
+        # --- 第二步：应用Hard-Swish激活函数 ---
+        # 使用F.hardswish，这是PyTorch内置的高效实现
+        activated_distances = torch.nn.functional.hardswish(chebyshev_dist)
+        
+        return activated_distances
+
+batch_size = 256
+feature_dim = 512
+
+def get_inputs():
+    # Generate two sets of vectors
+    x = torch.randn(batch_size, feature_dim)
+    y = torch.randn(batch_size, feature_dim)
+    return [x, y]
+
+def get_init_inputs():
+    return []  # No special initialization inputs needed
+
+
+
+Your task is to generate the `ModelNew` architecture with a custom CUDA kernel that computes the Chebyshev distance and applies the Hard-Swish activation in a fused manner. The implementation must be highly optimized.
+
+**CRITICAL REQUIREMENTS:**
+
+1.  **Performance Optimization & Fusion:**
+    *   **Operator Fusion:** The entire calculation (computing the Chebyshev distance and then applying the Hard-Swish activation function for each sample) must be performed within a **single CUDA kernel**. This kernel should output a tensor of Hard-Swish-activated distances.
+    *   The kernel should use a **multi-threaded reduction** strategy within each block to find the maximum absolute difference. After the reduction, a single thread should apply the Hard-Swish activation and store the result.
+
+2.  **Kernel Logic:**
+    *   The kernel should launch a 1D grid where each block corresponds to one sample in the batch.
+    *   Use shared memory to store the local maximums found by each thread, and then perform a parallel reduction within the block to find the global maximum for that sample.
+    *   The Hard-Swish activation is a piecewise function: `0` if `x <= -3`, `x` if `x >= 3`, and `x * (x + 3) / 6` otherwise. **Implement this efficiently using the ternary conditional operator (`?:`) to avoid branch divergence.**
+
+3.  **Code Structure:** Follow the exact structure of the provided example, including `load_inline`, the CUDA source string, the C++ wrapper source string, and the `ModelNew` class. The `get_init_inputs` function must return `[]` to match the baseline.
+
+4.  **Compilation Flags:** Use `-O3` for optimization but **do not** use `--use_fast_math` to ensure numerical accuracy with the PyTorch baseline, as finding the exact maximum is crucial. Avoid hardcoding compute capabilities to ensure portability.

S1/wut0n_#105/run_code.py

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