From 8c67c2d687540aad8ac2b14923595dbdd330d545 Mon Sep 17 00:00:00 2001 From: Ljy123 Date: Wed, 3 Dec 2025 23:41:39 +0800 Subject: [PATCH] finish exp-square #42 --- S1/Ljy123_#42/cudacode.py | 95 ++++++++++++++++++++++++++++++++++++++ S1/Ljy123_#42/prompt.txt | 2 + S1/Ljy123_#42/run_code.py | 52 +++++++++++++++++++++ S1/Ljy123_#42/torchcode.py | 29 ++++++++++++ 4 files changed, 178 insertions(+) create mode 100644 S1/Ljy123_#42/cudacode.py create mode 100644 S1/Ljy123_#42/prompt.txt create mode 100644 S1/Ljy123_#42/run_code.py create mode 100644 S1/Ljy123_#42/torchcode.py diff --git a/S1/Ljy123_#42/cudacode.py b/S1/Ljy123_#42/cudacode.py new file mode 100644 index 00000000..6be06659 --- /dev/null +++ b/S1/Ljy123_#42/cudacode.py @@ -0,0 +1,95 @@ +import torch +from torch.utils.cpp_extension import load_inline + +source = """ +#include + +__global__ void exp_square_affine_gate_kernel(const float* x, const float* scale, const float* bias, float* y, int B, int D, float alpha, float beta){ + int b = blockIdx.x; + int tid = threadIdx.x; + int stride = blockDim.x; + int row_start = b * D; + int aligned = ((row_start & 3) == 0); + if(aligned){ + int D4 = (D / 4) * 4; + #pragma unroll 4 + for(int i = tid * 4; i < D4; i += stride * 4){ + int base = row_start + i; + float4 xv = reinterpret_cast(x)[base / 4]; + float4 sv = reinterpret_cast(scale)[i / 4]; + float4 bv = reinterpret_cast(bias)[i / 4]; + float z0 = fmaf(xv.x, sv.x, bv.x); + float z1 = fmaf(xv.y, sv.y, bv.y); + float z2 = fmaf(xv.z, sv.z, bv.z); + float z3 = fmaf(xv.w, sv.w, bv.w); + float s0 = expf(z0 * z0); + float s1 = expf(z1 * z1); + float s2 = expf(z2 * z2); + float s3 = expf(z3 * z3); + float4 yv; + yv.x = xv.x * (1.0f / (1.0f + expf(-(alpha * s0 + beta)))); + yv.y = xv.y * (1.0f / (1.0f + expf(-(alpha * s1 + beta)))); + yv.z = xv.z * (1.0f / (1.0f + expf(-(alpha * s2 + beta)))); + yv.w = xv.w * (1.0f / (1.0f + expf(-(alpha * s3 + beta)))); + reinterpret_cast(y)[base / 4] = yv; + } + #pragma unroll 4 + for(int i = D4 + tid; i < D; i += stride){ + int base = row_start + i; + float z = fmaf(x[base], scale[i], bias[i]); + float s = expf(z * z); + float g = 1.0f / (1.0f + expf(-(alpha * s + beta))); + y[base] = x[base] * g; + } + } else { + #pragma unroll 4 + for(int i = tid; i < D; i += stride){ + int base = row_start + i; + float z = fmaf(x[base], scale[i], bias[i]); + float s = expf(z * z); + float g = 1.0f / (1.0f + expf(-(alpha * s + beta))); + y[base] = x[base] * g; + } + } +} + +torch::Tensor exp_square_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, 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 a = alpha.item(); + float be = beta.item(); + int block = 1024; + int grid = B; + exp_square_affine_gate_kernel<<>>(xc.data_ptr(), sc.data_ptr(), bc.data_ptr(), y.data_ptr(), B, D, a, be); + return y; +} +""" + +cpp_source = """ +torch::Tensor exp_square_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta); +""" + +ops = load_inline( + name="exp_square_affine_gate", + cpp_sources=cpp_source, + cuda_sources=source, + functions=["exp_square_affine_gate_cuda"], + extra_cuda_cflags=["-O3","--use_fast_math"], + verbose=True +) + +class ModelNew(torch.nn.Module): + def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float): + super(ModelNew, self).__init__() + self.ops = ops + self.register_buffer("scale", scale) + self.register_buffer("bias", bias) + self.register_buffer("alpha", torch.tensor(float(alpha), dtype=torch.float32)) + self.register_buffer("beta", torch.tensor(float(beta), dtype=torch.float32)) + + def forward(self, x: torch.Tensor): + return self.ops.exp_square_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta) diff --git a/S1/Ljy123_#42/prompt.txt b/S1/Ljy123_#42/prompt.txt new file mode 100644 index 00000000..8b754476 --- /dev/null +++ b/S1/Ljy123_#42/prompt.txt @@ -0,0 +1,2 @@ +You write custom CUDA kernels to replace PyTorch operators for speedups. +Implement Exp-Square Affine Gate on x[B,D]: Compute z = x*scale + bias, s = exp(z^2), gate g = sigmoid(alpha*s + beta), output y = x * g. Fuse affine, square, exp, sigmoid, and multiplication in a single grid-stride kernel. Provide a PyTorch reference with nn.Parameters for scale, bias, alpha, beta. Accuracy rtol=1e-3. diff --git a/S1/Ljy123_#42/run_code.py b/S1/Ljy123_#42/run_code.py new file mode 100644 index 00000000..9f0f3986 --- /dev/null +++ b/S1/Ljy123_#42/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 = 100 + 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 Exp-Square-Affine-Gate 平均执行时间: {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_#42/torchcode.py b/S1/Ljy123_#42/torchcode.py new file mode 100644 index 00000000..bea883d3 --- /dev/null +++ b/S1/Ljy123_#42/torchcode.py @@ -0,0 +1,29 @@ +import torch +import torch.nn as nn + +class Model(nn.Module): + def __init__(self, scale: torch.Tensor, bias: torch.Tensor, alpha: float, beta: float): + super(Model, self).__init__() + self.scale = nn.Parameter(scale) + self.bias = nn.Parameter(bias) + self.alpha = nn.Parameter(torch.tensor(float(alpha), 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) + s = torch.exp(z * z) + g = torch.sigmoid(self.alpha * s + self.beta) + return x * 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) + alpha = 1.0 + beta = 0.0 + return [scale, bias, alpha, beta]