Merge pull request 'finish contrastive-gate #34' (#336) from Ljy123/GPUCodeForces:contrastive into main

S1/Ljy123_#34/cudacode.py

@@ -0,0 +1,71 @@
+import torch
+from torch.utils.cpp_extension import load_inline
+
+source = """
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+
+__global__ void contrastive_gate_kernel(const float* a, const float* b, float* y, long long total, float alpha, float beta) {
+    long long tid = blockIdx.x * blockDim.x + threadIdx.x;
+    long long stride = blockDim.x * gridDim.x;
+    long long total4 = (total / 4) * 4;
+    for (long long i = tid * 4; i < total4; i += stride * 4) {
+        float4 av = reinterpret_cast<const float4*>(a)[i / 4];
+        float4 bv = reinterpret_cast<const float4*>(b)[i / 4];
+        float4 dv;
+        dv.x = av.x - bv.x; dv.y = av.y - bv.y; dv.z = av.z - bv.z; dv.w = av.w - bv.w;
+        float4 sv;
+        sv.x = 1.0f / (1.0f + expf(-(alpha * (av.x + bv.x) + beta)));
+        sv.y = 1.0f / (1.0f + expf(-(alpha * (av.y + bv.y) + beta)));
+        sv.z = 1.0f / (1.0f + expf(-(alpha * (av.z + bv.z) + beta)));
+        sv.w = 1.0f / (1.0f + expf(-(alpha * (av.w + bv.w) + beta)));
+        float4 yv;
+        yv.x = tanhf(dv.x) * sv.x;
+        yv.y = tanhf(dv.y) * sv.y;
+        yv.z = tanhf(dv.z) * sv.z;
+        yv.w = tanhf(dv.w) * sv.w;
+        reinterpret_cast<float4*>(y)[i / 4] = yv;
+    }
+    for (long long i = total4 + tid; i < total; i += stride) {
+        float d = a[i] - b[i];
+        float s = 1.0f / (1.0f + expf(-(alpha * (a[i] + b[i]) + beta)));
+        y[i] = tanhf(d) * s;
+    }
+}
+
+torch::Tensor contrastive_gate_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor alpha, torch::Tensor beta) {
+    auto ac = a.contiguous();
+    auto bc = b.contiguous();
+    auto y = torch::empty_like(ac);
+    long long total = ac.numel();
+    float al = alpha.item<float>();
+    float be = beta.item<float>();
+    int block = 1024;
+    long long grid = (total + block - 1) / block;
+    if (grid > 65535) grid = 65535;
+    contrastive_gate_kernel<<<(int)grid, block>>>(ac.data_ptr<float>(), bc.data_ptr<float>(), y.data_ptr<float>(), total, al, be);
+    return y;
+}
+"""
+
+cpp_source = """
+torch::Tensor contrastive_gate_cuda(torch::Tensor a, torch::Tensor b, torch::Tensor alpha, torch::Tensor beta);
+"""
+
+ops = load_inline(
+    name="contrastive_gate",
+    cpp_sources=cpp_source,
+    cuda_sources=source,
+    functions=["contrastive_gate_cuda"],
+    verbose=True
+)
+
+class ModelNew(torch.nn.Module):
+    def __init__(self, alpha: float, beta: float):
+        super(ModelNew, self).__init__()
+        self.ops = ops
+        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, a: torch.Tensor, b: torch.Tensor):
+        return self.ops.contrastive_gate_cuda(a, b, self.alpha, self.beta)

S1/Ljy123_#34/prompt.txt

@@ -0,0 +1,2 @@
+You write custom CUDA kernels to replace PyTorch operators for speedups.
+Implement a Dual-Input Contrastive Gate: Given two tensors a and b of shape [B, D], compute y = tanh(a - b) * sigmoid(alpha * (a + b) + beta). The CUDA kernel must fuse both inputs in a single pass with grid-stride loops over total elements, using contiguous memory and minimizing intermediate reads/writes. Provide a PyTorch reference module using nn.Parameters for alpha and beta, and ensure outputs match within rtol=1e-3. This operator emphasizes pairwise contrast and gated aggregation in one kernel.

S1/Ljy123_#34/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 Contrastive-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_#34/torchcode.py

@@ -0,0 +1,24 @@
+import torch
+import torch.nn as nn
+
+class Model(nn.Module):
+    def __init__(self, alpha: float, beta: float):
+        super(Model, self).__init__()
+        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, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        diff = torch.tanh(a - b)
+        gate = torch.sigmoid(self.alpha * (a + b) + self.beta)
+        return diff * gate
+
+batch_size = 32
+dim = 8192
+
+def get_inputs():
+    a = torch.randn(batch_size, dim)
+    b = torch.randn(batch_size, dim)
+    return [a, b]
+
+def get_init_inputs():
+    return [1.0, 0.0]