From 232314973ea7c409def7d1ba86137f00eee9d53a Mon Sep 17 00:00:00 2001 From: Ljy123 Date: Wed, 3 Dec 2025 18:26:57 +0800 Subject: [PATCH] finish contrastive-gate #34 --- S1/Ljy123_#34/cudacode.py | 71 ++++++++++++++++++++++++++++++++++++++ S1/Ljy123_#34/prompt.txt | 2 ++ S1/Ljy123_#34/run_code.py | 52 ++++++++++++++++++++++++++++ S1/Ljy123_#34/torchcode.py | 24 +++++++++++++ 4 files changed, 149 insertions(+) create mode 100644 S1/Ljy123_#34/cudacode.py create mode 100644 S1/Ljy123_#34/prompt.txt create mode 100644 S1/Ljy123_#34/run_code.py create mode 100644 S1/Ljy123_#34/torchcode.py diff --git a/S1/Ljy123_#34/cudacode.py b/S1/Ljy123_#34/cudacode.py new file mode 100644 index 00000000..54ff8642 --- /dev/null +++ b/S1/Ljy123_#34/cudacode.py @@ -0,0 +1,71 @@ +import torch +from torch.utils.cpp_extension import load_inline + +source = """ +#include +#include + +__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(a)[i / 4]; + float4 bv = reinterpret_cast(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(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 be = beta.item(); + int block = 1024; + long long grid = (total + block - 1) / block; + if (grid > 65535) grid = 65535; + contrastive_gate_kernel<<<(int)grid, block>>>(ac.data_ptr(), bc.data_ptr(), y.data_ptr(), 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) diff --git a/S1/Ljy123_#34/prompt.txt b/S1/Ljy123_#34/prompt.txt new file mode 100644 index 00000000..b0dfe6ef --- /dev/null +++ b/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. diff --git a/S1/Ljy123_#34/run_code.py b/S1/Ljy123_#34/run_code.py new file mode 100644 index 00000000..2407e4bc --- /dev/null +++ b/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() diff --git a/S1/Ljy123_#34/torchcode.py b/S1/Ljy123_#34/torchcode.py new file mode 100644 index 00000000..a19250b8 --- /dev/null +++ b/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]