Merge pull request 'finish ReflectionPad1d #24' (#48) from gsd123/GPUCodeForces:ReflectionPad1d into main

S1/24/ReflectionPad1d_cuda.py

@@ -0,0 +1,145 @@
+import torch
+import torch.nn as nn
+from torch.utils.cpp_extension import load_inline
+
+# -------------------------------------------------------------
+# 常量定义
+# -------------------------------------------------------------
+BATCH_SIZE = 32
+CHANNELS = 64
+WIDTH = 128  # W_in
+PADDING = (3, 1)  # (padding_left, padding_right)
+BLOCK_SIZE = 256  # CUDA Block 维度
+
+
+# -------------------------------------------------------------
+
+class ModelNew(nn.Module):
+    """
+    ReflectionPad1d 的高性能 CUDA 融合核函数实现
+    (修复了编译错误)
+    """
+
+    def __init__(self, padding):
+        super().__init__()
+
+        if isinstance(padding, int):
+            self.pad_L = padding
+            self.pad_R = padding
+        else:
+            self.pad_L = padding[0]
+            self.pad_R = padding[1]
+
+        self.block_size = BLOCK_SIZE
+        self._compile_cuda_kernel()
+
+    def _compile_cuda_kernel(self):
+
+        cpp_header = f"""
+        #include <torch/extension.h>
+
+        // C++ 接口
+        torch::Tensor reflection_pad1d_forward_cuda(
+            torch::Tensor input,
+            int pad_L,
+            int pad_R
+        );
+        """
+
+        cuda_source = f"""
+        #include <torch/extension.h> 
+        #include <cuda_runtime.h>
+
+        // [修复] 将 #define 移至此处
+        #define BLOCK_SIZE {self.block_size}
+
+        /*
+         * ReflectionPad1d 融合核函数
+         */
+        __global__ void reflection_pad1d_fused_kernel(
+            const float* __restrict__ input_data,
+            float* __restrict__ output_data,
+            int N, int C, int W_in, int W_out,
+            int pad_L, int pad_R
+        ) {{ // <-- f-string 转义
+            const int n_idx = blockIdx.x; 
+            const int c_idx = blockIdx.y; 
+            const int tid = threadIdx.x;
+
+            const float* p_in = input_data + (n_idx * C + c_idx) * W_in;
+            float* p_out = output_data + (n_idx * C + c_idx) * W_out;
+
+            // [修复] BLOCK_SIZE 现在可见
+            for (int j = tid; j < W_out; j += BLOCK_SIZE) {{ // <-- f-string 转义
+                int in_idx = 0;
+
+                if (j < pad_L) {{
+                    in_idx = pad_L - j;
+                }} else if (j < (pad_L + W_in)) {{
+                    in_idx = j - pad_L;
+                }} else {{
+                    int j_rel = j - (pad_L + W_in);
+                    in_idx = W_in - 2 - j_rel;
+                }}
+
+                p_out[j] = p_in[in_idx];
+            }}
+        }} 
+
+        // C++ 封装函数
+        // [修复] torch.Tensor -> torch::Tensor
+        torch::Tensor reflection_pad1d_forward_cuda(
+            torch::Tensor input,
+            int pad_L,
+            int pad_R
+        ) {{ 
+            TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
+            TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
+            TORCH_CHECK(input.dim() == 3, "input must be 3D (N, C, W)");
+
+            const int64_t N_64 = input.size(0);
+            const int64_t C_64 = input.size(1);
+            const int64_t W_in_64 = input.size(2);
+
+            TORCH_CHECK(pad_L < W_in_64, "padding_left should be less than input width");
+            TORCH_CHECK(pad_R < W_in_64, "padding_right should be less than input width");
+
+            const int64_t W_out_64 = W_in_64 + pad_L + pad_R;
+
+            auto output = torch::empty({{N_64, C_64, W_out_64}}, input.options());
+
+            dim3 grid_dim(N_64, C_64); 
+            dim3 block_dim(BLOCK_SIZE);
+
+            reflection_pad1d_fused_kernel<<<grid_dim, block_dim>>>(
+                input.data_ptr<float>(),
+                output.data_ptr<float>(),
+                static_cast<int>(N_64),
+                static_cast<int>(C_64),
+                static_cast<int>(W_in_64),
+                static_cast<int>(W_out_64),
+                pad_L,
+                pad_R
+            );
+
+            return output;
+        }} 
+        """
+
+        # JIT (Just-In-Time) 编译
+        self.pad_op = load_inline(
+            name="reflection_pad1d_op_v3_fixed",  # 更改名称以避免缓存
+            cpp_sources=cpp_header,
+            cuda_sources=cuda_source,
+            functions=["reflection_pad1d_forward_cuda"],
+            verbose=False  # 如果还报错，请设为 True
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        # 调用我们编译好的 CUDA C++ 函数
+        return self.pad_op.reflection_pad1d_forward_cuda(
+            x,
+            self.pad_L,
+            self.pad_R
+        )

S1/24/ReflectionPad1d_torch.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# -------------------------------------------------------------
+# 常量定义
+# -------------------------------------------------------------
+BATCH_SIZE = 32
+CHANNELS = 64
+WIDTH = 128  # W_in
+PADDING = (3, 1)  # (padding_left, padding_right)
+
+
+# -------------------------------------------------------------
+
+class Model(nn.Module):
+    """
+    nn.ReflectionPad1d 的纯 PyTorch 基准实现
+    (使用 F.pad)
+    """
+
+    def __init__(self, padding):
+        super().__init__()
+
+        if isinstance(padding, int):
+            # F.pad 需要 (left, right) 格式
+            self.padding_tuple = (padding, padding)
+        else:
+            self.padding_tuple = padding
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # F.pad 的 padding 格式是 (pad_dim_0_left, pad_dim_0_right, pad_dim_1_left, ...)
+        # 因为我们只 pad 最后一个维度 (dim -1)，所以元组是 (pad_L, pad_R)
+        return F.pad(x, self.padding_tuple, mode='reflect')
+
+
+def get_inputs():
+    """
+    生成一个 (N, C, W) 形状的输入
+    """
+    x = torch.randn(BATCH_SIZE, CHANNELS, WIDTH, dtype=torch.float32)
+    return [x]
+
+
+def get_init_inputs():
+    return [PADDING]

S1/24/prompt.txt

@@ -0,0 +1,105 @@
+You write custom CUDA kernels to replace the PyTorch operators in the given EvoNorm 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 normalization+affine_transform+nonlinear_gating), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.
+
+Key Technologies Used:
+
+Inline CUDA Extension in PyTorch: Uses torch.utils.cpp_extension.load_inline() to compile and load CUDA code directly within Python, providing seamless integration without external compilation steps.
+
+Simplified 1D Kernel Design: Implements a streamlined kernel optimized for 1D padding operations, focusing on width dimension processing only.
+
+Direct Global Memory Access: Unlike the 2D/3D versions, this implementation accesses global memory directly without shared memory caching, suitable for the simpler 1D case.
+
+1D Thread Blocking: Employs 1D thread blocks (BLOCK_SIZE = 256) for efficient parallelization across the width dimension.
+
+Grid-Strided Loop Pattern: Uses a strided loop (for (int j = tid; j < W_out; j += BLOCK_SIZE)) to distribute work across threads and handle arbitrary output sizes.
+
+Inline Reflection Logic: Implements reflection indexing directly within the kernel using conditional statements, avoiding separate device function calls.
+
+Batched Channel Processing: Processes multiple batches and channels concurrently through 2D grid dimensions (grid_dim(N, C)).
+
+Memory Layout Optimization: Leverages the natural memory layout of 3D tensors (N, C, W) with straightforward stride calculations.
+
+Comprehensive Error Checking: Includes validation for tensor dimensions, CUDA requirements, contiguity, and padding bounds.
+
+Performance Optimizations:
+
+Minimal kernel design with no synchronization overhead
+
+Coalesced memory access patterns for 1D data
+
+Grid-strided loops for optimal load balancing
+
+Restricted pointers for compiler optimization
+
+Direct indexing calculations
+
+Architecture Features:
+
+Separate C++ interface declaration and CUDA implementation
+
+Template-style parameter passing for padding values
+
+Automatic output tensor allocation with correct dimensions
+
+Efficient handling of 3D tensor layout (N, C, W)
+
+Key Differences from 2D/3D Versions:
+
+No shared memory usage (simpler access pattern)
+
+Direct reflection calculation in kernel
+
+1D thread blocks instead of 2D/3D
+
+Simpler memory addressing
+
+Reduced computational complexity
+
+
+Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# -------------------------------------------------------------
+# 常量定义
+# -------------------------------------------------------------
+BATCH_SIZE = 32
+CHANNELS = 64
+WIDTH = 128  # W_in
+PADDING = (3, 1)  # (padding_left, padding_right)
+
+
+# -------------------------------------------------------------
+
+class Model(nn.Module):
+    """
+    nn.ReflectionPad1d 的纯 PyTorch 基准实现
+    (使用 F.pad)
+    """
+
+    def __init__(self, padding):
+        super().__init__()
+
+        if isinstance(padding, int):
+            # F.pad 需要 (left, right) 格式
+            self.padding_tuple = (padding, padding)
+        else:
+            self.padding_tuple = padding
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # F.pad 的 padding 格式是 (pad_dim_0_left, pad_dim_0_right, pad_dim_1_left, ...)
+        # 因为我们只 pad 最后一个维度 (dim -1)，所以元组是 (pad_L, pad_R)
+        return F.pad(x, self.padding_tuple, mode='reflect')
+
+
+def get_inputs():
+    """
+    生成一个 (N, C, W) 形状的输入
+    """
+    x = torch.randn(BATCH_SIZE, CHANNELS, WIDTH, dtype=torch.float32)
+    return [x]
+
+
+def get_init_inputs():
+    return [PADDING]

S1/24/run_code.py

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