diff --git a/S1/zizi05_#4/prompt.txt b/S1/zizi05_#4/prompt.txt new file mode 100644 index 00000000..150b22be --- /dev/null +++ b/S1/zizi05_#4/prompt.txt @@ -0,0 +1,113 @@ +You write custom CUDA kernels to replace the PyTorch operators in the given SELU-Clip activation architecture to get speedups. +You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace the combined SELU calculation and clamp operators with a custom CUDA kernel (considering operator fusion opportunities to combine the element-wise SELU computation and range clipping into a single kernel) or adjust algorithms for better performance. You are only limited by your imagination. + +Here's an example to show you the syntax of inline embedding custom CUDA operators in PyTorch (for reference of code structure, not functional alignment): +The example given architecture (sample structure): +import torch +import torch.nn as nn +import torch.nn.functional as F +class Model(nn.Module): + def __init__(self) -> None: + super().__init__() + def forward(self, a, b): + return a + b +def get_inputs(): + # randomly generate input tensors based on the model architecture + a = torch.randn(1, 128).cuda() + b = torch.randn(1, 128).cuda() + return [a, b] +def get_init_inputs(): + # randomly generate tensors required for initialization based on the model architecture + return [] + +The example new arch with custom CUDA kernels (sample structure): +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.cpp_extension import load_inline +# Define custom CUDA kernel and load it inline +custom_add_source = """ +#include +#include +__global__ void custom_add_kernel(const float* a, const float* b, float* out, int size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < size) { + out[idx] = a[idx] + b[idx]; + } +} +torch::Tensor custom_add_cuda(torch::Tensor a, torch::Tensor b) { + auto size = a.numel(); + auto out = torch::empty_like(a); + const int block_size = 256; + int num_blocks = (size + block_size - 1) / block_size; + custom_add_kernel<<>>(a.data_ptr(), b.data_ptr(), out.data_ptr(), size); + return out; +} +""" +custom_add_cpp_source = "torch::Tensor custom_add_cuda(torch::Tensor a, torch::Tensor b);" +custom_add = load_inline( + name="custom_add", + cpp_sources=custom_add_cpp_source, + cuda_sources=custom_add_source, + functions=["custom_add_cuda"], + verbose=True +) +class Model(nn.Module): + def __init__(self) -> None: + super().__init__() + self.custom_add = custom_add + def forward(self, a, b): + return self.custom_add.custom_add_cuda(a, b) +def get_inputs(): + # randomly generate input tensors based on the model architecture + a = torch.randn(1, 128).cuda() + b = torch.randn(1, 128).cuda() + return [a, b] +def get_init_inputs(): + # randomly generate tensors required for initialization based on the model architecture + return [] + +You are given the following SELU-Clip activation architecture (base PyTorch implementation): +import torch +import torch.nn as nn +class Model(nn.Module): + """ + SELU-Clip activation function: Mathematical formulation is clamp(scale * (x if x > 0 else alpha*(exp(x)-1)), clip_min, clip_max), + where alpha=1.6732632423543772 (SELU shape parameter), scale=1.0507009873554804 (SELU scaling parameter), + clip_min=-5.0 (minimum clipping value), clip_max=5.0 (maximum clipping value). + """ + def __init__(self, alpha=1.6732632423543772, scale=1.0507009873554804, clip_min=-5.0, clip_max=5.0): + super(Model, self).__init__() + self.alpha = alpha + self.scale = scale + self.clip_min = clip_min + self.clip_max = clip_max + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Applies SELU-Clip activation to the input tensor. + Args: + x (torch.Tensor): Input tensor with fixed shape (batch_size, dim) + where batch_size=1024 and dim=2048. + Returns: + torch.Tensor: Output tensor with SELU-Clip applied, same shape as input. + """ + # Calculate SELU + selu_out = self.scale * torch.where( + x > 0, + x, + self.alpha * (torch.exp(x) - 1) + ) + # Apply clipping + clipped_out = torch.clamp(selu_out, self.clip_min, self.clip_max) + return clipped_out + +batch_size = 1024 +dim = 2048 +def get_inputs(): + # Randomly generate input tensor matching the fixed shape (batch_size, dim) + x = torch.randn(batch_size, dim) + return [x] +def get_init_inputs(): + # Provide initialization parameters for the model (non-trainable hyperparameters) + return (1.6732632423543772, 1.0507009873554804, -5.0, 5.0) \ No newline at end of file diff --git a/S1/zizi05_#4/run_code.py b/S1/zizi05_#4/run_code.py new file mode 100644 index 00000000..ddb83b7a --- /dev/null +++ b/S1/zizi05_#4/run_code.py @@ -0,0 +1,74 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from selu_clip_torchcode import Model,get_inputs,get_init_inputs +from selu_clip_cudacode import ModelNew + +def run_benchmark(): + # 检查 CUDA 是否可用 + 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("❌ 精度不一致!") + + print("\n-------------------- 性能加速比测试 --------------------") + num_iterations = 100 + + # PyTorch 模型计时 + 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 + + # 自定义 CUDA 内核计时 + 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 selu_clip 平均执行时间: {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() \ No newline at end of file diff --git a/S1/zizi05_#4/selu_clip_cudacode.py b/S1/zizi05_#4/selu_clip_cudacode.py new file mode 100644 index 00000000..3347c427 --- /dev/null +++ b/S1/zizi05_#4/selu_clip_cudacode.py @@ -0,0 +1,153 @@ +import torch +import torch.nn as nn +from torch.autograd import Function +from torch.utils.cpp_extension import load_inline + +# 定义CUDA内核和C++绑定代码 +cuda_source = """ +#include +#include +#include +#include + +// SELU-Clip前向传播CUDA内核 +__global__ void selu_clip_forward_kernel( + const float* __restrict__ x, + float* __restrict__ output, + float alpha, + float scale, + float clip_min, + float clip_max, + int num_elements) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + float val = x[idx]; + // 计算SELU + float selu_val; + if (val > 0.0f) { + selu_val = scale * val; + } else { + selu_val = scale * alpha * (expf(val) - 1.0f); + } + // 裁剪操作 + if (selu_val < clip_min) { + output[idx] = clip_min; + } else if (selu_val > clip_max) { + output[idx] = clip_max; + } else { + output[idx] = selu_val; + } + } +} + +// SELU-Clip反向传播CUDA内核 +__global__ void selu_clip_backward_kernel( + const float* __restrict__ x, + const float* __restrict__ grad_output, + float* __restrict__ grad_input, + float alpha, + float scale, + int num_elements) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + float val = x[idx]; + float grad = grad_output[idx]; + // 计算梯度 + if (val > 0.0f) { + grad_input[idx] = scale * grad; + } else { + grad_input[idx] = scale * alpha * expf(val) * grad; + } + } +} + +// C++绑定前向函数 +torch::Tensor selu_clip_forward( + torch::Tensor x, + float alpha, + float scale, + float clip_min, + float clip_max) { + auto output = torch::empty_like(x); + int num_elements = x.numel(); + int block_size = 256; + int grid_size = (num_elements + block_size - 1) / block_size; + + selu_clip_forward_kernel<<>>( + x.data_ptr(), + output.data_ptr(), + alpha, + scale, + clip_min, + clip_max, + num_elements + ); + return output; +} + +// C++绑定反向函数 +torch::Tensor selu_clip_backward( + torch::Tensor x, + torch::Tensor grad_output, + float alpha, + float scale) { + auto grad_input = torch::empty_like(x); + int num_elements = x.numel(); + int block_size = 256; + int grid_size = (num_elements + block_size - 1) / block_size; + + selu_clip_backward_kernel<<>>( + x.data_ptr(), + grad_output.data_ptr(), + grad_input.data_ptr(), + alpha, + scale, + num_elements + ); + return grad_input; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", &selu_clip_forward, "SELU-Clip forward"); + m.def("backward", &selu_clip_backward, "SELU-Clip backward"); +} +""" + +# 动态编译CUDA内核 +selu_clip_cuda = load_inline( + name="selu_clip_cuda", + cpp_sources=[], + cuda_sources=[cuda_source], + extra_cuda_cflags=["-O2"], + with_cuda=True, + verbose=False +) + +class SELUClipFunction(Function): + """自定义CUDA实现的SELU-Clip算子函数""" + @staticmethod + def forward(ctx, x, alpha, scale, clip_min, clip_max): + ctx.alpha = alpha + ctx.scale = scale + ctx.save_for_backward(x) + return selu_clip_cuda.forward(x, alpha, scale, clip_min, clip_max) + + @staticmethod + def backward(ctx, grad_output): + x, = ctx.saved_tensors + alpha = ctx.alpha + scale = ctx.scale + grad_input = selu_clip_cuda.backward(x, grad_output, alpha, scale) + return grad_input, None, None, None, None + +class ModelNew(nn.Module): + """使用内置CUDA内核的SELU-Clip模型""" + def __init__(self, alpha=1.6732632423543772, scale=1.0507009873554804, clip_min=-5.0, clip_max=5.0): + super(ModelNew, self).__init__() + self.alpha = alpha + self.scale = scale + self.clip_min = clip_min + self.clip_max = clip_max + + def forward(self, x): + return SELUClipFunction.apply(x, self.alpha, self.scale, self.clip_min, self.clip_max) \ No newline at end of file diff --git a/S1/zizi05_#4/selu_clip_torchcode.py b/S1/zizi05_#4/selu_clip_torchcode.py new file mode 100644 index 00000000..740ee35d --- /dev/null +++ b/S1/zizi05_#4/selu_clip_torchcode.py @@ -0,0 +1,32 @@ +import torch +import torch.nn as nn + +class Model(nn.Module): + """PyTorch实现的SELU-Clip算子""" + def __init__(self, alpha=1.6732632423543772, scale=1.0507009873554804, clip_min=-5.0, clip_max=5.0): + super(Model, self).__init__() + self.alpha = alpha + self.scale = scale + self.clip_min = clip_min + self.clip_max = clip_max + + def forward(self, x): + # 计算SELU + selu_out = self.scale * torch.where( + x > 0, + x, + self.alpha * (torch.exp(x) - 1) + ) + # 裁剪操作 + clipped_out = torch.clamp(selu_out, self.clip_min, self.clip_max) + return clipped_out + +def get_init_inputs(): + """提供模型初始化参数""" + return (1.6732632423543772, 1.0507009873554804, -5.0, 5.0) + +def get_inputs(): + """提供测试输入数据""" + torch.manual_seed(42) + x = torch.randn(1024, 2048) # 典型特征张量形状 + return (x,) \ No newline at end of file