finish SoftShrinkValve #103

S1/Ljy123_#103/cudacode.py

@@ -0,0 +1,127 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+#include <vector_types.h>
+
+__device__ __forceinline__ float softshrinkf(float x, float lambda_){
+    if (x > lambda_) return x - lambda_;
+    if (x < -lambda_) return x + lambda_;
+    return 0.0f;
+}
+
+__global__ __launch_bounds__(256) void softshrink_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 row = blockIdx.x;
+    int tid = threadIdx.x;
+    int col_block = blockIdx.y;
+    int stride4 = blockDim.x * 4;
+    int col_idx = (col_block * stride4) + tid * 4;
+    const float* xr = x + row * D;
+    float* yr = y + row * D;
+    const float lambda_ = 0.5f;
+    for(int base = col_idx; base < D; base += stride4 * 2){
+        if (base < D){
+            float4 xv = reinterpret_cast<const float4*>(xr + base)[0];
+            float4 sv = reinterpret_cast<const float4*>(scale + base)[0];
+            float4 bv = reinterpret_cast<const float4*>(bias + base)[0];
+            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 m0 = softshrinkf(z0, lambda_);
+            float m1 = softshrinkf(z1, lambda_);
+            float m2 = softshrinkf(z2, lambda_);
+            float m3 = softshrinkf(z3, lambda_);
+            float t0 = fmaf(alpha, m0, beta);
+            float t1 = fmaf(alpha, m1, beta);
+            float t2 = fmaf(alpha, m2, beta);
+            float t3 = fmaf(alpha, m3, 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<float4*>(yr + base)[0] = yv;
+        }
+        int base2 = base + stride4;
+        if (base2 < D){
+            float4 xv2 = reinterpret_cast<const float4*>(xr + base2)[0];
+            float4 sv2 = reinterpret_cast<const float4*>(scale + base2)[0];
+            float4 bv2 = reinterpret_cast<const float4*>(bias + base2)[0];
+            float4 yv2;
+            float z0b = fmaf(xv2.x, sv2.x, bv2.x);
+            float z1b = fmaf(xv2.y, sv2.y, bv2.y);
+            float z2b = fmaf(xv2.z, sv2.z, bv2.z);
+            float z3b = fmaf(xv2.w, sv2.w, bv2.w);
+            float mb0 = softshrinkf(z0b, lambda_);
+            float mb1 = softshrinkf(z1b, lambda_);
+            float mb2 = softshrinkf(z2b, lambda_);
+            float mb3 = softshrinkf(z3b, lambda_);
+            float tb0 = fmaf(alpha, mb0, beta);
+            float tb1 = fmaf(alpha, mb1, beta);
+            float tb2 = fmaf(alpha, mb2, beta);
+            float tb3 = fmaf(alpha, mb3, beta);
+            float gb0 = __fdividef(1.0f, 1.0f + __expf(-tb0));
+            float gb1 = __fdividef(1.0f, 1.0f + __expf(-tb1));
+            float gb2 = __fdividef(1.0f, 1.0f + __expf(-tb2));
+            float gb3 = __fdividef(1.0f, 1.0f + __expf(-tb3));
+            yv2.x = xv2.x * gb0;
+            yv2.y = xv2.y * gb1;
+            yv2.z = xv2.z * gb2;
+            yv2.w = xv2.w * gb3;
+            reinterpret_cast<float4*>(yr + base2)[0] = yv2;
+        }
+    }
+}
+
+torch::Tensor softshrink_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double 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 = (float)alpha;
+    float be = (float)beta;
+    int block = 256;
+    int elements_per_thread = 4;
+    int elements_per_block = block * elements_per_thread;
+    int gy = (D + elements_per_block - 1) / elements_per_block;
+    dim3 grid(B, gy);
+    softshrink_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 = """
+#include <torch/extension.h>
+torch::Tensor softshrink_affine_gate_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias, double alpha, double beta);
+"""
+
+ops = load_inline(
+    name="softshrink_affine_gate",
+    cpp_sources=cpp_source,
+    cuda_sources=source,
+    functions=["softshrink_affine_gate_cuda"],
+    extra_cflags=["-O3","-std=c++17"],
+    extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-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.alpha = float(alpha)
+        self.beta = float(beta)
+
+    def forward(self, x):
+        return self.ops.softshrink_affine_gate_cuda(x, self.scale, self.bias, self.alpha, self.beta)

S1/Ljy123_#103/prompt.txt

@@ -0,0 +1,10 @@
+Operator: SoftShrink-Affine-Gate (Fused CUDA Kernel)
+
+Definition
+- z = x * scale + bias
+- m = SoftShrink(z, lambda=0.5)
+- g = sigmoid(alpha * m + beta)
+- y = x * g
+
+Goal
+- Fuse ops to reduce bandwidth and kernel launches; target ≥1.30x speedup.

S1/Ljy123_#103/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 SoftShrink-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()

S1/Ljy123_#103/torchcode.py

@@ -0,0 +1,29 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+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
+        m = F.softshrink(z, lambd=0.5)
+        g = torch.sigmoid(self.alpha * m + 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]