From aff8b2bf1913bf1860f940a49420950a88ceb71d Mon Sep 17 00:00:00 2001 From: Ljy123 Date: Wed, 3 Dec 2025 23:30:21 +0800 Subject: [PATCH] optimized softmax-affine (use group) #37 --- S1/Ljy123_#37/cudacode.py | 110 +++++++++++++++++++++++++++++++++++++ S1/Ljy123_#37/prompt.txt | 2 + S1/Ljy123_#37/run_code.py | 52 ++++++++++++++++++ S1/Ljy123_#37/torchcode.py | 36 ++++++++++++ 4 files changed, 200 insertions(+) create mode 100644 S1/Ljy123_#37/cudacode.py create mode 100644 S1/Ljy123_#37/prompt.txt create mode 100644 S1/Ljy123_#37/run_code.py create mode 100644 S1/Ljy123_#37/torchcode.py diff --git a/S1/Ljy123_#37/cudacode.py b/S1/Ljy123_#37/cudacode.py new file mode 100644 index 00000000..7c136f1f --- /dev/null +++ b/S1/Ljy123_#37/cudacode.py @@ -0,0 +1,110 @@ +import torch +from torch.utils.cpp_extension import load_inline + +source = """ +#include + +__global__ void group_softmax_affine_kernel(const float* x, const float* scale, const float* bias, float* y, int N, int C, int H, int W, int group) { + 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; + int G = (C + group - 1) / group; + float* sX = sh; + float* sZ = sh + C; + float* sExp = sh + 2 * C; + float* sM = sh + 3 * C; + float* sDen = sM + G; + 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; + sZ[c] = xv * scale[c] + bias[c]; + } + __syncthreads(); + if (tid == 0) { + for (int gi = 0; gi < G; ++gi) { + int s = gi * group; + int e = s + group; if (e > C) e = C; + float m = -INFINITY; + for (int c = s; c < e; ++c) { + float z = sZ[c]; + if (z > m) m = z; + } + sM[gi] = m; + } + } + __syncthreads(); + for (int c = tid; c < C; c += stride) { + int gi = c / group; + float e = expf(sZ[c] - sM[gi]); + sExp[c] = e; + } + __syncthreads(); + if (tid == 0) { + for (int gi = 0; gi < G; ++gi) { + int s = gi * group; + int e = s + group; if (e > C) e = C; + float den = 0.0f; + for (int c = s; c < e; ++c) den += sExp[c]; + sDen[gi] = den; + } + } + __syncthreads(); + for (int c = tid; c < C; c += stride) { + int gi = c / group; + long long idx = ((long long)n * C + c) * (long long)H * (long long)W + (long long)h * W + w; + float g = sExp[c] / sDen[gi]; + y[idx] = g * sX[c]; + } + __syncthreads(); + } +} + +torch::Tensor group_softmax_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor group) { + 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); + int g = group.item(); + int block = C <= 1024 ? C : 1024; + int grid = (N * H * W); + if (grid > 65535) grid = 65535; + int G = (C + g - 1) / g; + size_t shmem = (size_t)(3 * C + G) * sizeof(float) + (size_t)block * sizeof(float); + group_softmax_affine_kernel<<>>(xc.data_ptr(), sc.data_ptr(), bc.data_ptr(), y.data_ptr(), N, C, H, W, g); + return y; +} +""" + +cpp_source = """ +torch::Tensor group_softmax_affine_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor group); +""" + +ops = load_inline( + name="group_softmax_affine", + cpp_sources=cpp_source, + cuda_sources=source, + functions=["group_softmax_affine_cuda"], + verbose=True +) + +class ModelNew(torch.nn.Module): + def __init__(self, scale: torch.Tensor, bias: torch.Tensor, group: int): + super(ModelNew, self).__init__() + self.ops = ops + self.register_buffer("scale", scale) + self.register_buffer("bias", bias) + self.register_buffer("group", torch.tensor(int(group), dtype=torch.int32)) + + def forward(self, x: torch.Tensor): + return self.ops.group_softmax_affine_cuda(x, self.scale, self.bias, self.group) diff --git a/S1/Ljy123_#37/prompt.txt b/S1/Ljy123_#37/prompt.txt new file mode 100644 index 00000000..2164441e --- /dev/null +++ b/S1/Ljy123_#37/prompt.txt @@ -0,0 +1,2 @@ +You write custom CUDA kernels to replace PyTorch operators for speedups. +Implement Group-Softmax-Affine gating on NCHW: Given x[N,C,H,W], per-channel scale[C], bias[C], and an integer group size g, compute z = x*scale[c] + bias[c], then for each (n,h,w) apply softmax within channel groups of size g: g_k = softmax(z over channels in group k), and output y = g_k * x for those channels. The CUDA kernel should process one spatial location per block, cache x and z in shared memory, compute per-group max/sub-exp/sum, and write gated outputs in one pass. Provide a PyTorch reference with nn.Parameter scale, bias, and group as Parameter. Accuracy must match within rtol=1e-3. diff --git a/S1/Ljy123_#37/run_code.py b/S1/Ljy123_#37/run_code.py new file mode 100644 index 00000000..a76b0920 --- /dev/null +++ b/S1/Ljy123_#37/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 Group-Softmax-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() diff --git a/S1/Ljy123_#37/torchcode.py b/S1/Ljy123_#37/torchcode.py new file mode 100644 index 00000000..0c0b96b7 --- /dev/null +++ b/S1/Ljy123_#37/torchcode.py @@ -0,0 +1,36 @@ +import torch +import torch.nn as nn + +class Model(nn.Module): + def __init__(self, scale: torch.Tensor, bias: torch.Tensor, group: int): + super(Model, self).__init__() + self.scale = nn.Parameter(scale) + self.bias = nn.Parameter(bias) + self.register_buffer("group", torch.tensor(int(group), dtype=torch.int32)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + g = int(self.group.item()) + z = x * self.scale.view(1,-1,1,1) + self.bias.view(1,-1,1,1) + C = z.size(1) + G = (C + g - 1) // g + y = torch.empty_like(x) + for gi in range(G): + s = gi * g + e = min(s + g, C) + m = torch.amax(z[:, s:e], dim=1, keepdim=True) + ex = torch.exp(z[:, s:e] - m) + den = torch.sum(ex, dim=1, keepdim=True) + y[:, s:e] = (ex / den) * x[:, s:e] + return y + +N, C, H, W = 4, 96, 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) + group = 16 + return [scale, bias, group]