From cb3102003e76847cbb43083db2691efe1a8686be Mon Sep 17 00:00:00 2001 From: Ljy123 Date: Wed, 10 Dec 2025 22:43:24 +0800 Subject: [PATCH] finish logcosh-affinegate #58 --- S1/Ljy123_#58/cudacode.py | 165 +++++++++++++++++++++++++++++++++++++ S1/Ljy123_#58/prompt.txt | 19 +++++ S1/Ljy123_#58/run_code.py | 50 +++++++++++ S1/Ljy123_#58/torchcode.py | 28 +++++++ 4 files changed, 262 insertions(+) create mode 100644 S1/Ljy123_#58/cudacode.py create mode 100644 S1/Ljy123_#58/prompt.txt create mode 100644 S1/Ljy123_#58/run_code.py create mode 100644 S1/Ljy123_#58/torchcode.py diff --git a/S1/Ljy123_#58/cudacode.py b/S1/Ljy123_#58/cudacode.py new file mode 100644 index 00000000..bd580848 --- /dev/null +++ b/S1/Ljy123_#58/cudacode.py @@ -0,0 +1,165 @@ +import torch +from torch.utils.cpp_extension import load_inline + +source = """ +#include +#include +#include + +__device__ __forceinline__ float logcosh(float x){ + float ax = fabsf(x); + if(ax < 1e-3f){ + float x2 = x * x; + float x4 = x2 * x2; + return 0.5f * x2 - (1.0f/12.0f) * x4; + } + float t = __expf(-2.0f * ax); + if(t < 1e-5f){ + return ax + t - 0.6931471805599453f; + } + return ax + __logf(1.0f + t) - 0.6931471805599453f; +} + +__device__ __forceinline__ float sigmoid_stable(float x){ + if(x >= 0.0f){ + float e = __expf(-x); + return 1.0f / (1.0f + e); + } else { + float e = __expf(x); + return e / (1.0f + e); + } +} + +__global__ __launch_bounds__(1024) void logcosh_affine_gate_kernel(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){ + long long tid = blockIdx.x * blockDim.x + threadIdx.x; + long long stride = blockDim.x * gridDim.x; + long long total = (long long)B * D; + for(long long i = tid; i < total; i += stride){ + int d = (int)(i % D); + float xv = __ldg(x + i); + float z = fmaf(xv, __ldg(scale + d), __ldg(bias + d)); + float t = fmaf(logcosh(z), alpha, beta); + float g = 1.0f / (1.0f + __expf(-t)); + y[i] = xv * g; + } +} + +torch::Tensor logcosh_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; + long long total = (long long)B * D; + int grid = (int)std::min(65535LL, (total + block - 1) / block); + logcosh_affine_gate_kernel<<>>(xc.data_ptr(), sc.data_ptr(), bc.data_ptr(), y.data_ptr(), B, D, a, be); + return y; +} + +__global__ __launch_bounds__(512) void logcosh_affine_gate_kernel_v2(const float* __restrict__ x, const float* __restrict__ scale, const float* __restrict__ bias, float* __restrict__ y, int B, int D, float alpha, float beta){ + int b = blockIdx.x; + int lane = blockIdx.y * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.y; + int row_start = b * D; + const float* xr = x + row_start; + float* yr = y + row_start; + int aligned = ((((long long)xr & 15LL) == 0) && (((long long)yr & 15LL) == 0) && (((long long)scale & 15LL) == 0) && (((long long)bias & 15LL) == 0) && ((D & 3) == 0)); + if(aligned){ + int D4 = (D / 4) * 4; + #pragma unroll 12 + for(int i = lane * 4; i < D4; i += stride * 4){ + float4 xv = reinterpret_cast(xr)[i / 4]; + float4 sv = reinterpret_cast(scale)[i / 4]; + float4 bv = reinterpret_cast(bias)[i / 4]; + float4 yv; + 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 t0 = fmaf(logcosh(z0), alpha, beta); + float t1 = fmaf(logcosh(z1), alpha, beta); + float t2 = fmaf(logcosh(z2), alpha, beta); + float t3 = fmaf(logcosh(z3), alpha, beta); + float g0 = __fdividef(1.0f, 1.0f + __expf(-t0)); + float g1 = __fdividef(1.0f, 1.0f + __expf(-t1)); + float g2 = __fdividef(1.0f, 1.0f + __expf(-t2)); + float g3 = __fdividef(1.0f, 1.0f + __expf(-t3)); + yv.x = xv.x * g0; + yv.y = xv.y * g1; + yv.z = xv.z * g2; + yv.w = xv.w * g3; + reinterpret_cast(yr)[i / 4] = yv; + } + #pragma unroll 12 + for(int i = D4 + lane; i < D; i += stride){ + float z = fmaf(__ldg(xr + i), __ldg(scale + i), __ldg(bias + i)); + float t = fmaf(logcosh(z), alpha, beta); + float g = __fdividef(1.0f, 1.0f + __expf(-t)); + yr[i] = xr[i] * g; + } + } else { + #pragma unroll 12 + for(int i = lane; i < D; i += stride){ + float z = fmaf(__ldg(xr + i), __ldg(scale + i), __ldg(bias + i)); + float t = fmaf(logcosh(z), alpha, beta); + float g = __fdividef(1.0f, 1.0f + __expf(-t)); + yr[i] = xr[i] * g; + } + } +} + +torch::Tensor logcosh_affine_gate_cuda_v2(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 = 512; + int gy = std::min(32, std::max(1, (D + 1023) / 1024)); + dim3 grid(B, gy); + logcosh_affine_gate_kernel_v2<<>>(xc.data_ptr(), sc.data_ptr(), bc.data_ptr(), y.data_ptr(), B, D, a, be); + return y; +} + +torch::Tensor logcosh_affine_gate_cuda_opt(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta){ + int D = (int)x.size(1); + if((D & 3) == 0){ + return logcosh_affine_gate_cuda_v2(x, scale, bias, alpha, beta); + } + return logcosh_affine_gate_cuda(x, scale, bias, alpha, beta); +} +""" + +cpp_source = """ +torch::Tensor logcosh_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta); +torch::Tensor logcosh_affine_gate_cuda_v2(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta); +torch::Tensor logcosh_affine_gate_cuda_opt(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta); +""" + +ops = load_inline( + name="logcosh_affine_gate", + cpp_sources=cpp_source, + cuda_sources=source, + functions=["logcosh_affine_gate_cuda","logcosh_affine_gate_cuda_v2","logcosh_affine_gate_cuda_opt"], + extra_cuda_cflags=["-O3","--use_fast_math","-gencode=arch=compute_80,code=sm_80","-Xptxas","-O3,-dlcm=ca","-maxrregcount=64"], + 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): + return self.ops.logcosh_affine_gate_cuda_opt(x, self.scale, self.bias, self.alpha, self.beta) diff --git a/S1/Ljy123_#58/prompt.txt b/S1/Ljy123_#58/prompt.txt new file mode 100644 index 00000000..e430b2f1 --- /dev/null +++ b/S1/Ljy123_#58/prompt.txt @@ -0,0 +1,19 @@ +Objective: Replace PyTorch operations with a single fused CUDA kernel for LogCosh-Affine-Gate to achieve ≥1.3x speedup while matching outputs within rtol=1e-3. + +Constraints: +- Keep the 4-file structure identical to `example` folder: `cudacode.py`, `torchcode.py`, `run_code.py`, `prompt.txt`. +- The fused kernel must compute: `z = x*scale + bias; v = log(cosh(z)); g = sigmoid(alpha*v + beta); y = x*g`. +- Use numerically stable math and avoid NaNs for typical random inputs. +- Exploit GPU multi-threading aggressively and minimize memory traffic via fusion. + +Design Guidelines: +- Launch one block per row with 256–512 threads, and set `grid.y` to split long rows for higher SM occupancy. +- Prefer vectorized loads/stores (`float4`) on aligned paths, and fallback to scalar path otherwise. +- Use FMA for affine and gate preparation to reduce instruction count and rounding. +- Use fast intrinsic `__expf` for the sigmoid path, verify accuracy under rtol=1e-3. +- Keep reductions out of the hot path; this kernel is purely elementwise and memory-bound. + +Benchmark Setup: +- Batch size: 16, Dimension: 16384. +- Measure average latency over 100 iterations for both PyTorch and the fused CUDA kernel. +- Report precision alignment and speedup. Target speedup ≥1.3x. diff --git a/S1/Ljy123_#58/run_code.py b/S1/Ljy123_#58/run_code.py new file mode 100644 index 00000000..3a768a92 --- /dev/null +++ b/S1/Ljy123_#58/run_code.py @@ -0,0 +1,50 @@ +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().eval() + cuda_model = ModelNew(*init_inputs).cuda().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 + 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 + + 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 LogCosh-Affine-Gate 平均执行时间: {torch_time:.6f} 秒") + print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f} 秒") + speedup = torch_time / cuda_time if cuda_time > 0 else 0 + if cuda_time > 0: + print(f"加速比 (Speedup): {speedup:.2f}x") + else: + print("CUDA 内核执行时间为0，无法计算加速比。") + return precision_flag, speedup + +if __name__ == "__main__": + run_benchmark() diff --git a/S1/Ljy123_#58/torchcode.py b/S1/Ljy123_#58/torchcode.py new file mode 100644 index 00000000..b376cf3f --- /dev/null +++ b/S1/Ljy123_#58/torchcode.py @@ -0,0 +1,28 @@ +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.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) -> torch.Tensor: + z = x * self.scale + self.bias + v = torch.log(torch.cosh(z)) + g = torch.sigmoid(self.alpha * v + self.beta) + return x * g + +batch_size = 16 +dim = 16384 + +def get_inputs(): + x = torch.randn(batch_size, dim) + return [x] + +def get_init_inputs(): + scale = torch.randn(dim) + bias = torch.randn(dim) + return [scale, bias, 1.0, 0.0]