optimized channel-affine(temp-softmax) #40

S1/Ljy123_#40/cudacode.py

@@ -0,0 +1,103 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+source = """
+#include <torch/extension.h>
+
+__global__ void temp_softmax_channel_affine_kernel(const float* x, const float* scale, const float* bias, float* y, int N, int C, int H, int W, float temp) {
+    int tid = threadIdx.x;
+    int stride = blockDim.x;
+    int total = N * H * W;
+    extern __shared__ float sh[];
+    for (int pos = blockIdx.x; pos < total; pos += gridDim.x) {
+        int n = pos / (H * W);
+        int hw = pos % (H * W);
+        int h = hw / W;
+        int w = hw % W;
+        float* sX = sh;
+        float* sZ = sh + C;
+        float* sExp = sh + 2 * C;
+        float* sRed = sh + 3 * C;
+        for (int c = tid; c < C; c += stride) {
+            long long idx = ((long long)n * C + c) * (long long)H * (long long)W + (long long)h * W + w;
+            float xv = x[idx];
+            sX[c] = xv;
+            float z = xv * scale[c] + bias[c];
+            sZ[c] = z / temp;
+        }
+        __syncthreads();
+        float local_max = -INFINITY;
+        for (int c = tid; c < C; c += stride) {
+            float z = sZ[c];
+            if (z > local_max) local_max = z;
+        }
+        sRed[tid] = local_max;
+        __syncthreads();
+        for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
+            if (tid < s && sRed[tid + s] > sRed[tid]) sRed[tid] = sRed[tid + s];
+            __syncthreads();
+        }
+        float m = sRed[0];
+        float local_sum = 0.0f;
+        for (int c = tid; c < C; c += stride) {
+            float e = expf(sZ[c] - m);
+            sExp[c] = e;
+            local_sum += e;
+        }
+        sRed[tid] = local_sum;
+        __syncthreads();
+        for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
+            if (tid < s) sRed[tid] += sRed[tid + s];
+            __syncthreads();
+        }
+        float denom = sRed[0];
+        for (int c = tid; c < C; c += stride) {
+            long long idx = ((long long)n * C + c) * (long long)H * (long long)W + (long long)h * W + w;
+            float g = sExp[c] / denom;
+            y[idx] = g * sX[c];
+        }
+        __syncthreads();
+    }
+}
+
+torch::Tensor temp_softmax_channel_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor temp) {
+    auto xc = x.contiguous();
+    auto sc = scale.contiguous();
+    auto bc = bias.contiguous();
+    auto y = torch::empty_like(xc);
+    int N = (int)xc.size(0);
+    int C = (int)xc.size(1);
+    int H = (int)xc.size(2);
+    int W = (int)xc.size(3);
+    float t = temp.item<float>();
+    int block = C <= 1024 ? C : 1024;
+    int grid = (N * H * W);
+    if (grid > 65535) grid = 65535;
+    size_t shmem = (size_t)(3 * C + block) * sizeof(float);
+    temp_softmax_channel_affine_kernel<<<grid, block, shmem>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), N, C, H, W, t);
+    return y;
+}
+"""
+
+cpp_source = """
+torch::Tensor temp_softmax_channel_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor temp);
+"""
+
+ops = load_inline(
+    name="temp_softmax_channel_affine",
+    cpp_sources=cpp_source,
+    cuda_sources=source,
+    functions=["temp_softmax_channel_affine_cuda"],
+    verbose=True
+)
+
+class ModelNew(torch.nn.Module):
+    def __init__(self, scale: torch.Tensor, bias: torch.Tensor, temp: float):
+        super(ModelNew, self).__init__()
+        self.ops = ops
+        self.register_buffer("scale", scale)
+        self.register_buffer("bias", bias)
+        self.register_buffer("temp", torch.tensor(float(temp), dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor):
+        return self.ops.temp_softmax_channel_affine_cuda(x, self.scale, self.bias, self.temp)

S1/Ljy123_#40/prompt.txt

@@ -0,0 +1,2 @@
+You write custom CUDA kernels to replace PyTorch operators for speedups.
+Implement Temperature Softmax Channel Affine: For x[N,C,H,W], scale[C], bias[C], and temperature t, compute z = x*scale + bias, then g = softmax(z/t) along channel axis for each (n,h,w), output y = g * x. Use a single CUDA block per spatial location, caching x and z/t in shared memory, with max-subtraction and reduction for the denominator, writing results in one pass. Provide a PyTorch reference with nn.Parameter scale, bias, temp. Accuracy rtol=1e-3.

S1/Ljy123_#40/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)
+    if flag:
+        print("✅ 精度对齐：两个模型的输出结果非常接近。")
+    else:
+        print("❌ 精度不一致！")
+        print(f"最大绝对误差: {(out_torch - out_cuda).abs().max().item()}" )
+
+    print("\n-------------------- 性能加速比测试 --------------------")
+    iters = 30
+    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 Temp-Softmax-Channel-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_#40/torchcode.py

@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    def __init__(self, scale: torch.Tensor, bias: torch.Tensor, temp: float):
+        super(Model, self).__init__()
+        self.scale = nn.Parameter(scale)
+        self.bias = nn.Parameter(bias)
+        self.temp = nn.Parameter(torch.tensor(float(temp), dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        z = x * self.scale.view(1,-1,1,1) + self.bias.view(1,-1,1,1)
+        t = self.temp
+        zt = z / t
+        m = torch.amax(zt, dim=1, keepdim=True)
+        ex = torch.exp(zt - m)
+        den = torch.sum(ex, dim=1, keepdim=True)
+        g = ex / den
+        return g * x
+
+N, C, H, W = 4, 128, 32, 32
+
+def get_inputs():
+    x = torch.randn(N, C, H, W)
+    return [x]
+
+def get_init_inputs():
+    scale = torch.randn(C)
+    bias = torch.randn(C)
+    temp = 0.8
+    return [scale, bias, temp]