From 0ef1bea766c7f20160da7323a982a1681ddc2a74 Mon Sep 17 00:00:00 2001
From: Ljy123 <Ljy123@example.org>
Date: Sat, 6 Dec 2025 17:07:20 +0800
Subject: [PATCH] finish square-sigmoid (affine-gate) #47

---
 S1/LJy123_#47/cudacode.py  | 94 ++++++++++++++++++++++++++++++++++++++
 S1/LJy123_#47/prompt.txt   | 28 ++++++++++++
 S1/LJy123_#47/run_code.py  | 57 +++++++++++++++++++++++
 S1/LJy123_#47/torchcode.py | 29 ++++++++++++
 4 files changed, 208 insertions(+)
 create mode 100644 S1/LJy123_#47/cudacode.py
 create mode 100644 S1/LJy123_#47/prompt.txt
 create mode 100644 S1/LJy123_#47/run_code.py
 create mode 100644 S1/LJy123_#47/torchcode.py

diff --git a/S1/LJy123_#47/cudacode.py b/S1/LJy123_#47/cudacode.py
new file mode 100644
index 00000000..36973645
--- /dev/null
+++ b/S1/LJy123_#47/cudacode.py
@@ -0,0 +1,94 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+
+__global__ void square_sigmoid_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){
+    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 4
+        for(int i = lane * 4; i < D4; i += stride * 4){
+            float4 xv = reinterpret_cast<const float4*>(xr)[i / 4];
+            float4 sv = reinterpret_cast<const float4*>(scale)[i / 4];
+            float4 bv = reinterpret_cast<const float4*>(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 g0 = 1.0f / (1.0f + expf(-(alpha * (z0 * z0) + beta)));
+            float g1 = 1.0f / (1.0f + expf(-(alpha * (z1 * z1) + beta)));
+            float g2 = 1.0f / (1.0f + expf(-(alpha * (z2 * z2) + beta)));
+            float g3 = 1.0f / (1.0f + expf(-(alpha * (z3 * z3) + beta)));
+            yv.x = xv.x * g0;
+            yv.y = xv.y * g1;
+            yv.z = xv.z * g2;
+            yv.w = xv.w * g3;
+            reinterpret_cast<float4*>(yr)[i / 4] = yv;
+        }
+        #pragma unroll 4
+        for(int i = D4 + lane; i < D; i += stride){
+            float z = fmaf(xr[i], scale[i], bias[i]);
+            float g = 1.0f / (1.0f + expf(-(alpha * (z * z) + beta)));
+            yr[i] = xr[i] * g;
+        }
+    } else {
+        #pragma unroll 4
+        for(int i = lane; i < D; i += stride){
+            float z = fmaf(xr[i], scale[i], bias[i]);
+            float g = 1.0f / (1.0f + expf(-(alpha * (z * z) + beta)));
+            yr[i] = xr[i] * g;
+        }
+    }
+}
+
+torch::Tensor square_sigmoid_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>();
+    float be = beta.item<float>();
+    int block = 256;
+    int gy = max(1, min((D + 4095) / 4096, 8));
+    dim3 grid(B, gy);
+    square_sigmoid_affine_gate_kernel<<<grid, block>>>(xc.data_ptr<float>(), sc.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, be);
+    return y;
+}
+"""
+
+cpp_source = """
+torch::Tensor square_sigmoid_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, torch::Tensor alpha, torch::Tensor beta);
+"""
+
+ops = load_inline(
+    name="square_sigmoid_affine_gate",
+    cpp_sources=cpp_source,
+    cuda_sources=source,
+    functions=["square_sigmoid_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):
+        return self.ops.square_sigmoid_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)
diff --git a/S1/LJy123_#47/prompt.txt b/S1/LJy123_#47/prompt.txt
new file mode 100644
index 00000000..01fa6e1c
--- /dev/null
+++ b/S1/LJy123_#47/prompt.txt
@@ -0,0 +1,28 @@
+融合算子：Square-Sigmoid-Affine-Gate（一次核内完成仿射、平方与 Sigmoid 门控，返回 y = x * σ(α * z^2 + β)，其中 z = x*scale + bias）。平方增强幅值差异并通过 Sigmoid 控制门控强度。
+
+目标与定义
+- 输入张量：`x[B, D]`
+- 逐维参数：`scale[D]`、`bias[D]`
+- 标量超参：`alpha`、`beta`
+- 计算流程：`z = x*scale + bias`，`v = z*z`，`g = sigmoid(alpha*v + beta)`，`y = x * g`
+
+参考实现（文件要求）
+- `torchcode.py`：PyTorch 参考 `Model`；统一的 `get_inputs()`/`get_init_inputs()`
+- `cudacode.py`：单核融合（仿射+平方+sigmoid+乘法）；`-O3 --use_fast_math`
+- `run_code.py`：迭代 100 次；`rtol=1e-03, atol=1e-06` 精度；打印加速比
+
+CUDA 实现要点
+- 并行布局：`grid = B`；块内沿 D 合并访存
+- 对齐向量化：16 字节对齐且 `D%4==0` 时走 `float4`；否则标量回退
+- 指令优化：仿射用 `fmaf`；sigmoid 用 `expf`；循环 `#pragma unroll 4`
+- 溢出注意：`z^2` 对大幅值会放大，sigmoid 可缓和，但仍需避免中间溢出；使用 `float` 常规范围下问题不大
+- 线程配置：推荐 `block=1024`，按设备与规模微调
+
+评估与目标
+- 精度：对齐 `rtol=1e-03, atol=1e-06`
+- 性能：≥1.0x 加速；对齐触发向量化时更佳
+
+加分项（可选）
+- 尾元素处理与分支收敛优化
+- 每线程批量步长以提高吞吐与占用
+
diff --git a/S1/LJy123_#47/run_code.py b/S1/LJy123_#47/run_code.py
new file mode 100644
index 00000000..f221f2e3
--- /dev/null
+++ b/S1/LJy123_#47/run_code.py
@@ -0,0 +1,57 @@
+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():
+        output_torch = torch_model(*inputs)
+        output_cuda = cuda_model(*inputs)
+    precision_flag = torch.allclose(output_torch, output_cuda, rtol=1e-03, atol=1e-06)
+    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 Square-Sigmoid-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_#47/torchcode.py b/S1/LJy123_#47/torchcode.py
new file mode 100644
index 00000000..30a486c7
--- /dev/null
+++ b/S1/LJy123_#47/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.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 = z * 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]
+