From a305af6609629c4db4e98d71d2b0bd75c1d97a21 Mon Sep 17 00:00:00 2001 From: Ljy123 Date: Sat, 6 Dec 2025 16:30:02 +0800 Subject: [PATCH] re upload gtu-fuse operator #19 --- S1/Ljy123_#19/cudacode.py | 82 ++++++++++++++++++++++++++++++++++++++ S1/Ljy123_#19/prompt.txt | 5 +++ S1/Ljy123_#19/run_code.py | 76 +++++++++++++++++++++++++++++++++++ S1/Ljy123_#19/torchcode.py | 24 +++++++++++ 4 files changed, 187 insertions(+) create mode 100644 S1/Ljy123_#19/cudacode.py create mode 100644 S1/Ljy123_#19/prompt.txt create mode 100644 S1/Ljy123_#19/run_code.py create mode 100644 S1/Ljy123_#19/torchcode.py diff --git a/S1/Ljy123_#19/cudacode.py b/S1/Ljy123_#19/cudacode.py new file mode 100644 index 00000000..57b90d30 --- /dev/null +++ b/S1/Ljy123_#19/cudacode.py @@ -0,0 +1,82 @@ +import torch +from torch.utils.cpp_extension import load_inline + +source = """ +#include +#include + +__global__ void gtu_kernel(const float* x, float* y, int D, long long rows) { + int b = blockIdx.x; + int tid = threadIdx.x; + int stride = blockDim.x; + long long row_start = b * (long long)(2 * D); + const float* a_row = x + row_start; + const float* b_row = x + row_start + D; + float* y_row = y + b * D; + int aligned = ((((long long)a_row & 15LL) == 0) && (((long long)b_row & 15LL) == 0) && (((long long)y_row & 15LL) == 0) && ((D & 3) == 0)); + if (aligned) { + int D4 = (D / 4) * 4; + #pragma unroll 4 + for (int i = tid * 4; i < D4; i += stride * 4) { + float4 av = reinterpret_cast(a_row)[i / 4]; + float4 bv = reinterpret_cast(b_row)[i / 4]; + float4 yv; + yv.x = tanhf(av.x) * (1.0f / (1.0f + expf(-bv.x))); + yv.y = tanhf(av.y) * (1.0f / (1.0f + expf(-bv.y))); + yv.z = tanhf(av.z) * (1.0f / (1.0f + expf(-bv.z))); + yv.w = tanhf(av.w) * (1.0f / (1.0f + expf(-bv.w))); + reinterpret_cast(y_row)[i / 4] = yv; + } + #pragma unroll 4 + for (int i = D4 + tid; i < D; i += stride) { + float a = a_row[i]; + float b = b_row[i]; + float ta = tanhf(a); + float sb = 1.0f / (1.0f + expf(-b)); + y_row[i] = ta * sb; + } + } else { + #pragma unroll 4 + for (int i = tid; i < D; i += stride) { + float a = a_row[i]; + float b = b_row[i]; + float ta = tanhf(a); + float sb = 1.0f / (1.0f + expf(-b)); + y_row[i] = ta * sb; + } + } +} + +torch::Tensor gtu_cuda(torch::Tensor x) { + auto x_contig = x.contiguous(); + long long rows = 1; + for (int i = 0; i < x_contig.dim() - 1; ++i) rows *= x_contig.size(i); + int D = (int)(x_contig.size(-1) / 2); + auto y = torch::empty({rows, D}, x_contig.options()); + int block = 1024; + int grid = (int)rows; + gtu_kernel<<>>(x_contig.data_ptr(), y.data_ptr(), D, rows); + return y; +} +""" + +cpp_source = """ +torch::Tensor gtu_cuda(torch::Tensor x); +""" + +ops = load_inline( + name="gtu", + cpp_sources=cpp_source, + cuda_sources=source, + functions=["gtu_cuda"], + extra_cuda_cflags=["-O3","--use_fast_math"], + verbose=True +) + +class ModelNew(torch.nn.Module): + def __init__(self): + super(ModelNew, self).__init__() + self.ops = ops + + def forward(self, x): + return self.ops.gtu_cuda(x) diff --git a/S1/Ljy123_#19/prompt.txt b/S1/Ljy123_#19/prompt.txt new file mode 100644 index 00000000..f646b2ce --- /dev/null +++ b/S1/Ljy123_#19/prompt.txt @@ -0,0 +1,5 @@ +独特融合算子:GTU(tanh(a) * sigmoid(b))。输入最后维度为 2D,按行拆分为 a 与 b,一次内核完成融合。 + +torchcode.py:参考实现,按最后维切分并组合。 +cudacode.py:`__global__ void gtu_kernel(...)` 完成融合计算。 +run_code.py:比较精度与性能(100 次迭代,`rtol=1e-03`)。 diff --git a/S1/Ljy123_#19/run_code.py b/S1/Ljy123_#19/run_code.py new file mode 100644 index 00000000..acaa5d5e --- /dev/null +++ b/S1/Ljy123_#19/run_code.py @@ -0,0 +1,76 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +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 + else: + device = torch.device("cuda") + + init_inputs = get_init_inputs() + init_inputs = [ + x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs + ] + inputs = get_inputs() + inputs = [ + x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs + ] + + torch_model = Model(*init_inputs).cuda() + cuda_model = ModelNew(*init_inputs).cuda() + + torch_model.eval() + cuda_model.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("❌ 精度不一致!") + diff = (output_torch - output_cuda).abs().max().item() + print(f"最大绝对误差: {diff}") + print(f"输出张量形状: torch={tuple(output_torch.shape)}, cuda={tuple(output_cuda.shape)}") + print(f"数据类型: torch={output_torch.dtype}, cuda={output_cuda.dtype}") + print(f"设备: torch={output_torch.device}, cuda={output_cuda.device}") + + 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 GTU 平均执行时间: {torch_time:.6f} 秒") + print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f} 秒") + speedup = 0 + if cuda_time > 0: + speedup = torch_time / cuda_time + print(f"加速比 (Speedup): {speedup:.2f}x") + else: + print("CUDA 内核执行时间为0,无法计算加速比。") + return precision_flag,speedup + +if __name__ == "__main__": + precision_flag,speedup = run_benchmark() diff --git a/S1/Ljy123_#19/torchcode.py b/S1/Ljy123_#19/torchcode.py new file mode 100644 index 00000000..6ae46a1d --- /dev/null +++ b/S1/Ljy123_#19/torchcode.py @@ -0,0 +1,24 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + D = x.size(-1) // 2 + a = x[..., :D] + b = x[..., D:] + return torch.tanh(a) * torch.sigmoid(b) + +batch_size = 16 +dim_half = 16384 +dim = dim_half * 2 + +def get_inputs(): + x = torch.randn(batch_size, dim) + return [x] + +def get_init_inputs(): + return []