finish prototype-cosine-gate #36

S1/Ljy123_#36/cudacode.py

@@ -0,0 +1,103 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+source = """
+#include <torch/extension.h>
+
+__inline__ __device__ float warpReduceSum(float val){
+    for(int offset=16; offset>0; offset>>=1){
+        val += __shfl_down_sync(0xffffffff, val, offset);
+    }
+    return val;
+}
+
+__global__ void proto_cos_gate_kernel(const float* x, const float* p, float* y, int B, int D, float alpha, float beta){
+    int b = blockIdx.x;
+    int tid = threadIdx.x;
+    int lane = tid & 31;
+    int warpId = tid >> 5;
+    extern __shared__ float sh[];
+    float s_dot = 0.0f;
+    float s_x2 = 0.0f;
+    float s_p2 = 0.0f;
+    for(int i = tid; i < D; i += blockDim.x){
+        float xv = x[b * D + i];
+        float pv = p[i];
+        s_dot += xv * pv;
+        s_x2 += xv * xv;
+        s_p2 += pv * pv;
+    }
+    s_dot = warpReduceSum(s_dot);
+    s_x2 = warpReduceSum(s_x2);
+    s_p2 = warpReduceSum(s_p2);
+    if(lane == 0){
+        sh[warpId] = s_dot;
+        sh[warpId + ((blockDim.x + 31)>>5)] = s_x2;
+        sh[warpId + 2*((blockDim.x + 31)>>5)] = s_p2;
+    }
+    __syncthreads();
+    float dot = 0.0f, x2 = 0.0f, p2 = 0.0f;
+    if(warpId == 0){
+        int WN = ((blockDim.x + 31)>>5);
+        float vd = (lane < WN) ? sh[lane] : 0.0f;
+        float vx = (lane < WN) ? sh[WN + lane] : 0.0f;
+        float vp = (lane < WN) ? sh[2*WN + lane] : 0.0f;
+        vd = warpReduceSum(vd);
+        vx = warpReduceSum(vx);
+        vp = warpReduceSum(vp);
+        if(lane == 0){
+            sh[0] = vd;
+            sh[1] = vx;
+            sh[2] = vp;
+        }
+    }
+    __syncthreads();
+    dot = sh[0]; x2 = sh[1]; p2 = sh[2];
+    float nx = sqrtf(x2 + 1e-8f);
+    float np = sqrtf(p2 + 1e-8f);
+    float cosv = dot / (nx * np);
+    float g = 1.0f / (1.0f + expf(-(alpha * cosv + beta)));
+    for(int i = tid; i < D; i += blockDim.x){
+        y[b * D + i] = x[b * D + i] * g;
+    }
+}
+
+torch::Tensor proto_cos_gate_cuda(torch::Tensor x, torch::Tensor proto, torch::Tensor alpha, torch::Tensor beta){
+    auto xc = x.contiguous();
+    auto pc = proto.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 b = beta.item<float>();
+    int block = 256;
+    int grid = B;
+    size_t WN = (size_t)((block + 31)>>5);
+    size_t shmem = WN * 3 * sizeof(float);
+    proto_cos_gate_kernel<<<grid, block, shmem>>>(xc.data_ptr<float>(), pc.data_ptr<float>(), y.data_ptr<float>(), B, D, a, b);
+    return y;
+}
+"""
+
+cpp_source = """
+torch::Tensor proto_cos_gate_cuda(torch::Tensor x, torch::Tensor proto, torch::Tensor alpha, torch::Tensor beta);
+"""
+
+ops = load_inline(
+    name="prototype_cosine_gate",
+    cpp_sources=cpp_source,
+    cuda_sources=source,
+    functions=["proto_cos_gate_cuda"],
+    verbose=True
+)
+
+class ModelNew(torch.nn.Module):
+    def __init__(self, proto: torch.Tensor, alpha: float, beta: float):
+        super(ModelNew, self).__init__()
+        self.ops = ops
+        self.register_buffer("proto", proto)
+        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):
+        return self.ops.proto_cos_gate_cuda(x, self.proto, self.alpha, self.beta)

S1/Ljy123_#36/prompt.txt

@@ -0,0 +1,2 @@
+You write custom CUDA kernels to replace PyTorch operators for speedups.
+Implement Prototype-Cosine Gate: For x[B,D] and a learnable prototype p[D], compute cos similarity per row cos = (x·p) / (||x||·||p||), gate g = sigmoid(alpha*cos + beta), and output y = x * g. Use one CUDA block per row with warp-level reductions for dot and norms, then apply the scalar gate across the row within the same kernel. Provide a PyTorch reference using nn.Parameter for p, alpha, beta. Ensure accuracy within rtol=1e-3.

S1/Ljy123_#36/run_code.py

@@ -0,0 +1,52 @@
+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():
+        out_torch = torch_model(*inputs)
+        out_cuda = cuda_model(*inputs)
+    flag = torch.allclose(out_torch, out_cuda, rtol=1e-03)
+    if flag:
+        print("✅ 精度对齐：两个模型的输出结果非常接近。")
+    else:
+        print("❌ 精度不一致！")
+        print(f"最大绝对误差: {(out_torch - out_cuda).abs().max().item()}" )
+
+    print("\n-------------------- 性能加速比测试 --------------------")
+    iters = 100
+    torch.cuda.synchronize(); t0 = time.time()
+    for _ in range(iters):
+        _ = torch_model(*inputs)
+    torch.cuda.synchronize(); t_torch = (time.time() - t0) / iters
+
+    torch.cuda.synchronize(); t0 = time.time()
+    for _ in range(iters):
+        _ = cuda_model(*inputs)
+    torch.cuda.synchronize(); t_cuda = (time.time() - t0) / iters
+
+    print(f"PyTorch Prototype-Cosine-Gate 平均执行时间: {t_torch:.6f} 秒")
+    print(f"自定义 CUDA 融合内核 平均执行时间: {t_cuda:.6f} 秒")
+    sp = t_torch / t_cuda if t_cuda > 0 else 0
+    if t_cuda > 0:
+        print(f"加速比 (Speedup): {sp:.2f}x")
+    else:
+        print("CUDA 内核执行时间为0，无法计算加速比。")
+    return flag, sp
+
+if __name__ == "__main__":
+    run_benchmark()

S1/Ljy123_#36/torchcode.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    def __init__(self, proto: torch.Tensor, alpha: float, beta: float):
+        super(Model, self).__init__()
+        self.proto = nn.Parameter(proto)
+        self.alpha = nn.Parameter(torch.tensor(float(alpha), dtype=torch.float32))
+        self.beta = nn.Parameter(torch.tensor(float(beta), dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dot = torch.sum(x * self.proto.view(1,-1), dim=1, keepdim=True)
+        nx = torch.sqrt(torch.sum(x * x, dim=1, keepdim=True) + 1e-8)
+        np = torch.sqrt(torch.sum(self.proto * self.proto) + 1e-8)
+        cosv = dot / (nx * np)
+        g = torch.sigmoid(self.alpha * cosv + self.beta)
+        return x * g
+
+batch_size = 32
+dim = 8192
+
+def get_inputs():
+    x = torch.randn(batch_size, dim)
+    return [x]
+
+def get_init_inputs():
+    proto = torch.randn(dim)
+    alpha = 1.0
+    beta = 0.0
+    return [proto, alpha, beta]