finish range-gate-affine #39

S1/Ljy123_#39/cudacode.py

@@ -0,0 +1,99 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+source = """
+#include <torch/extension.h>
+
+__inline__ __device__ float warpReduceMax(float val){
+    for(int offset=16; offset>0; offset>>=1){ float v = __shfl_down_sync(0xffffffff, val, offset); val = val > v ? val : v; }
+    return val;
+}
+__inline__ __device__ float warpReduceMin(float val){
+    for(int offset=16; offset>0; offset>>=1){ float v = __shfl_down_sync(0xffffffff, val, offset); val = val < v ? val : v; }
+    return val;
+}
+
+__global__ void range_gate_affine_kernel(const float* x, const float* scale, const float* bias, float* y, int B, int D, float gamma, float beta){
+    int b = blockIdx.x;
+    int tid = threadIdx.x;
+    int lane = tid & 31;
+    int warpId = tid >> 5;
+    extern __shared__ float sh[];
+    float lmax = -INFINITY;
+    float lmin = INFINITY;
+    for(int i = tid * 4; i < D; i += blockDim.x * 4){
+        int base = b * D + i;
+        for(int k = 0; k < 4 && (i + k) < D; ++k){
+            float z = fmaf(x[base + k], scale[i + k], bias[i + k]);
+            lmax = z > lmax ? z : lmax;
+            lmin = z < lmin ? z : lmin;
+        }
+    }
+    lmax = warpReduceMax(lmax);
+    lmin = warpReduceMin(lmin);
+    if(lane == 0){
+        sh[warpId] = lmax;
+        sh[warpId + ((blockDim.x + 31)>>5)] = lmin;
+    }
+    __syncthreads();
+    float rmax = -INFINITY, rmin = INFINITY;
+    if(warpId == 0){
+        int WN = ((blockDim.x + 31)>>5);
+        float vmax = (lane < WN) ? sh[lane] : -INFINITY;
+        float vmin = (lane < WN) ? sh[WN + lane] : INFINITY;
+        vmax = warpReduceMax(vmax);
+        vmin = warpReduceMin(vmin);
+        if(lane == 0){ sh[0] = vmax; sh[1] = vmin; }
+    }
+    __syncthreads();
+    rmax = sh[0]; rmin = sh[1];
+    float g = 1.0f / (1.0f + expf(-(gamma * (rmax - rmin) + beta)));
+    for(int i = tid * 4; i < D; i += blockDim.x * 4){
+        int base = b * D + i;
+        for(int k = 0; k < 4 && (i + k) < D; ++k){
+            float z = fmaf(x[base + k], scale[i + k], bias[i + k]);
+            y[base + k] = z * g;
+        }
+    }
+}
+
+torch::Tensor range_gate_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor gamma, torch::Tensor beta){
+    auto xc = x.contiguous();
+    auto sc = scale.contiguous();
+    auto bc = bias.contiguous();
+    auto y = torch::empty_like(xc);
+    int B = (int)xc.size(0);
+    int D = (int)xc.size(1);
+    float g = gamma.item<float>();
+    float be = beta.item<float>();
+    int block = 1024;
+    int grid = B;
+    size_t shmem = (size_t)((block + 31)>>5) * 2 * sizeof(float);
+    range_gate_affine_kernel<<<grid, block, shmem>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, g, be);
+    return y;
+}
+"""
+
+cpp_source = """
+torch::Tensor range_gate_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor gamma, torch::Tensor beta);
+"""
+
+ops = load_inline(
+    name="range_gate_affine",
+    cpp_sources=cpp_source,
+    cuda_sources=source,
+    functions=["range_gate_affine_cuda"],
+    verbose=True
+)
+
+class ModelNew(torch.nn.Module):
+    def __init__(self, scale: torch.Tensor, bias: torch.Tensor, gamma: float, beta: float):
+        super(ModelNew, self).__init__()
+        self.ops = ops
+        self.register_buffer("scale", scale)
+        self.register_buffer("bias", bias)
+        self.register_buffer("gamma", torch.tensor(float(gamma), dtype=torch.float32))
+        self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor):
+        return self.ops.range_gate_affine_cuda(x, self.scale, self.bias, self.gamma, self.beta)

S1/Ljy123_#39/prompt.txt

@@ -0,0 +1,2 @@
+You write custom CUDA kernels to replace PyTorch operators for speedups.
+Implement Range-Gate Affine on x[B,D]: Compute z = x*scale + bias, then per-row g = sigmoid(gamma*(max(z)-min(z)) + beta), output y = z * g. Use one CUDA block per row, reduce max and min with warp shuffles, and apply the scalar gate across the row in the same kernel. Provide a PyTorch reference with nn.Parameter scale, bias, gamma, beta. Accuracy rtol=1e-3.

S1/Ljy123_#39/run_code.py

@@ -0,0 +1,52 @@
+import torch
+import time
+from torchcode import Model, get_inputs, get_init_inputs
+from cudacode import ModelNew
+
+def run_benchmark():
+    if not torch.cuda.is_available():
+        print("CUDA 不可用，请确保您有可用的 NVIDIA GPU 并已正确安装 PyTorch CUDA 版本。")
+        return
+    device = torch.device("cuda")
+
+    init_inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_init_inputs()]
+    inputs = [x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in get_inputs()]
+
+    torch_model = Model(*init_inputs).cuda()
+    cuda_model = ModelNew(*init_inputs).cuda()
+    torch_model.eval(); cuda_model.eval()
+
+    print("-------------------- 精度对齐验证 --------------------")
+    with torch.no_grad():
+        out_torch = torch_model(*inputs)
+        out_cuda = cuda_model(*inputs)
+    flag = torch.allclose(out_torch, out_cuda, rtol=1e-03, atol=1e-06)
+    if flag:
+        print("✅ 精度对齐：两个模型的输出结果非常接近。")
+    else:
+        print("❌ 精度不一致！")
+        print(f"最大绝对误差: {(out_torch - out_cuda).abs().max().item()}" )
+
+    print("\n-------------------- 性能加速比测试 --------------------")
+    iters = 50
+    torch.cuda.synchronize(); t0 = time.time()
+    for _ in range(iters):
+        _ = torch_model(*inputs)
+    torch.cuda.synchronize(); t_torch = (time.time() - t0) / iters
+
+    torch.cuda.synchronize(); t0 = time.time()
+    for _ in range(iters):
+        _ = cuda_model(*inputs)
+    torch.cuda.synchronize(); t_cuda = (time.time() - t0) / iters
+
+    print(f"PyTorch Range-Gate-Affine 平均执行时间: {t_torch:.6f} 秒")
+    print(f"自定义 CUDA 融合内核 平均执行时间: {t_cuda:.6f} 秒")
+    sp = t_torch / t_cuda if t_cuda > 0 else 0
+    if t_cuda > 0:
+        print(f"加速比 (Speedup): {sp:.2f}x")
+    else:
+        print("CUDA 内核执行时间为0，无法计算加速比。")
+    return flag, sp
+
+if __name__ == "__main__":
+    run_benchmark()

S1/Ljy123_#39/torchcode.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    def __init__(self, scale: torch.Tensor, bias: torch.Tensor, gamma: float, beta: float):
+        super(Model, self).__init__()
+        self.scale = nn.Parameter(scale)
+        self.bias = nn.Parameter(bias)
+        self.gamma = nn.Parameter(torch.tensor(float(gamma), dtype=torch.float32))
+        self.beta = nn.Parameter(torch.tensor(float(beta), dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        z = x * self.scale.view(1,-1) + self.bias.view(1,-1)
+        zmax = torch.amax(z, dim=1, keepdim=True)
+        zmin = torch.amin(z, dim=1, keepdim=True)
+        g = torch.sigmoid(self.gamma * (zmax - zmin) + self.beta)
+        return z * g
+
+B, D = 64, 8192
+
+def get_inputs():
+    x = torch.randn(B, D)
+    return [x]
+
+def get_init_inputs():
+    scale = torch.randn(D)
+    bias = torch.randn(D)
+    gamma = 1.0
+    beta = 0.0
+    return [scale, bias, gamma, beta]