re upload gtu-fuse operator #19

S1/Ljy123_#19/cudacode.py

@@ -0,0 +1,82 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+
+__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<const float4*>(a_row)[i / 4];
+            float4 bv = reinterpret_cast<const float4*>(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<float4*>(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<<<grid, block>>>(x_contig.data_ptr<float>(), y.data_ptr<float>(), 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)

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`）。

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()

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 []