From c9ee753fad9de41e0a2afd93c122eea7a0f9327a Mon Sep 17 00:00:00 2001 From: Ljy123 Date: Mon, 1 Dec 2025 20:41:09 +0800 Subject: [PATCH] optimized affine-relu #13 --- S1/Ljy123_#13/cudacode.py | 55 +++++++++++++++++++++++++++ S1/Ljy123_#13/prompt.txt | 5 +++ S1/Ljy123_#13/run_code.py | 76 ++++++++++++++++++++++++++++++++++++++ S1/Ljy123_#13/torchcode.py | 24 ++++++++++++ 4 files changed, 160 insertions(+) create mode 100644 S1/Ljy123_#13/cudacode.py create mode 100644 S1/Ljy123_#13/prompt.txt create mode 100644 S1/Ljy123_#13/run_code.py create mode 100644 S1/Ljy123_#13/torchcode.py diff --git a/S1/Ljy123_#13/cudacode.py b/S1/Ljy123_#13/cudacode.py new file mode 100644 index 00000000..a96c9bdd --- /dev/null +++ b/S1/Ljy123_#13/cudacode.py @@ -0,0 +1,55 @@ +import torch +from torch.utils.cpp_extension import load_inline + +source = """ +#include +#include + +__global__ void affine_relu6_kernel(const float* x, const float* scale, const float* bias, float* y, int dim, long long total) { + long long idx = blockIdx.x * blockDim.x + threadIdx.x; + long long stride = blockDim.x * gridDim.x; + for (long long i = idx; i < total; i += stride) { + int j = (int)(i % dim); + float v = __fmul_rn(x[i], scale[j]); + v = __fadd_rn(v, bias[j]); + v = fminf(fmaxf(v, 0.0f), 6.0f); + y[i] = v; + } +} + +torch::Tensor affine_relu6_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias) { + auto x_contig = x.contiguous(); + auto s_contig = scale.contiguous(); + auto b_contig = bias.contiguous(); + auto y = torch::empty_like(x_contig); + long long total = x_contig.numel(); + int dim = (int)x_contig.size(-1); + int block = 512; + long long grid = (total + block - 1) / block; + grid = grid > 65535 ? 65535 : grid; + affine_relu6_kernel<<<(int)grid, block>>>(x_contig.data_ptr(), s_contig.data_ptr(), b_contig.data_ptr(), y.data_ptr(), dim, total); + return y; +} +""" + +cpp_source = """ +torch::Tensor affine_relu6_cuda(torch::Tensor x, torch::Tensor scale, torch::Tensor bias); +""" + +ops = load_inline( + name="affine_relu6", + cpp_sources=cpp_source, + cuda_sources=source, + functions=["affine_relu6_cuda"], + verbose=True +) + +class ModelNew(torch.nn.Module): + def __init__(self, scale: torch.Tensor, bias: torch.Tensor): + super(ModelNew, self).__init__() + self.ops = ops + self.register_buffer("scale", scale) + self.register_buffer("bias", bias) + + def forward(self, x): + return self.ops.affine_relu6_cuda(x, self.scale, self.bias) diff --git a/S1/Ljy123_#13/prompt.txt b/S1/Ljy123_#13/prompt.txt new file mode 100644 index 00000000..abdbc9e6 --- /dev/null +++ b/S1/Ljy123_#13/prompt.txt @@ -0,0 +1,5 @@ +Affine+ReLU6 融合：一次内核完成仿射与 ReLU6（0..6）裁剪，减少内核与显存往返。 + +torchcode.py：参考实现 `y = clamp(x * scale + bias, 0, 6)`。 +cudacode.py：`__global__ void affine_relu6_kernel(...)` 完成融合计算。 +run_code.py：比较精度与性能（100 次迭代，`rtol=1e-03`）。 diff --git a/S1/Ljy123_#13/run_code.py b/S1/Ljy123_#13/run_code.py new file mode 100644 index 00000000..8619b30d --- /dev/null +++ b/S1/Ljy123_#13/run_code.py @@ -0,0 +1,76 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +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 + else: + device = torch.device("cuda") + + init_inputs = get_init_inputs() + init_inputs = [ + x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in init_inputs + ] + inputs = get_inputs() + inputs = [ + x.cuda(device=device) if isinstance(x, torch.Tensor) else x for x in inputs + ] + + torch_model = Model(*init_inputs).cuda() + cuda_model = ModelNew(*init_inputs).cuda() + + torch_model.eval() + cuda_model.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("❌ 精度不一致！") + diff = (output_torch - output_cuda).abs().max().item() + print(f"最大绝对误差: {diff}") + print(f"输出张量形状: torch={tuple(output_torch.shape)}, cuda={tuple(output_cuda.shape)}") + print(f"数据类型: torch={output_torch.dtype}, cuda={output_cuda.dtype}") + print(f"设备: torch={output_torch.device}, cuda={output_cuda.device}") + + 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 Affine+ReLU6 平均执行时间: {torch_time:.6f} 秒") + print(f"自定义 CUDA 融合内核 平均执行时间: {cuda_time:.6f} 秒") + speedup = 0 + if cuda_time > 0: + speedup = torch_time / cuda_time + print(f"加速比 (Speedup): {speedup:.2f}x") + else: + print("CUDA 内核执行时间为0，无法计算加速比。") + return precision_flag,speedup + +if __name__ == "__main__": + precision_flag,speedup = run_benchmark() diff --git a/S1/Ljy123_#13/torchcode.py b/S1/Ljy123_#13/torchcode.py new file mode 100644 index 00000000..e3f46b50 --- /dev/null +++ b/S1/Ljy123_#13/torchcode.py @@ -0,0 +1,24 @@ +import torch +import torch.nn as nn + +class Model(nn.Module): + def __init__(self, scale: torch.Tensor, bias: torch.Tensor): + super(Model, self).__init__() + self.register_buffer("scale", scale) + self.register_buffer("bias", bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = x * self.scale + self.bias + return torch.clamp(y, 0.0, 6.0) + +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]