From c5dafff0aaaca92bb34e3c865ae167133225d02e Mon Sep 17 00:00:00 2001 From: Ljy123 Date: Wed, 3 Dec 2025 23:28:58 +0800 Subject: [PATCH] finish prototype-cosine-gate #36 --- S1/Ljy123_#36/cudacode.py | 103 +++++++++++++++++++++++++++++++++++++ S1/Ljy123_#36/prompt.txt | 2 + S1/Ljy123_#36/run_code.py | 52 +++++++++++++++++++ S1/Ljy123_#36/torchcode.py | 30 +++++++++++ 4 files changed, 187 insertions(+) create mode 100644 S1/Ljy123_#36/cudacode.py create mode 100644 S1/Ljy123_#36/prompt.txt create mode 100644 S1/Ljy123_#36/run_code.py create mode 100644 S1/Ljy123_#36/torchcode.py diff --git a/S1/Ljy123_#36/cudacode.py b/S1/Ljy123_#36/cudacode.py new file mode 100644 index 00000000..f7dc652e --- /dev/null +++ b/S1/Ljy123_#36/cudacode.py @@ -0,0 +1,103 @@ +import torch +from torch.utils.cpp_extension import load_inline + +source = """ +#include + +__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 b = beta.item(); + 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<<>>(xc.data_ptr(), pc.data_ptr(), y.data_ptr(), 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) diff --git a/S1/Ljy123_#36/prompt.txt b/S1/Ljy123_#36/prompt.txt new file mode 100644 index 00000000..029da458 --- /dev/null +++ b/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. diff --git a/S1/Ljy123_#36/run_code.py b/S1/Ljy123_#36/run_code.py new file mode 100644 index 00000000..14c6d1f4 --- /dev/null +++ b/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() diff --git a/S1/Ljy123_#36/torchcode.py b/S1/Ljy123_#36/torchcode.py new file mode 100644 index 00000000..97443c08 --- /dev/null +++ b/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]