From fb8a6833740fb3af6441c7d981ce9c4afafefab9 Mon Sep 17 00:00:00 2001 From: hli28146 Date: Tue, 2 Dec 2025 21:03:39 +0800 Subject: [PATCH] finish ASU #30 --- S1/hli28146_#30/asu_cuda.py | 102 +++++++++++++++++++++++++++++++++++ S1/hli28146_#30/asu_torch.py | 35 ++++++++++++ S1/hli28146_#30/prompt.txt | 57 ++++++++++++++++++++ S1/hli28146_#30/run_code.py | 74 +++++++++++++++++++++++++ 4 files changed, 268 insertions(+) create mode 100644 S1/hli28146_#30/asu_cuda.py create mode 100644 S1/hli28146_#30/asu_torch.py create mode 100644 S1/hli28146_#30/prompt.txt create mode 100644 S1/hli28146_#30/run_code.py diff --git a/S1/hli28146_#30/asu_cuda.py b/S1/hli28146_#30/asu_cuda.py new file mode 100644 index 00000000..f9c7c258 --- /dev/null +++ b/S1/hli28146_#30/asu_cuda.py @@ -0,0 +1,102 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +cpp_source = """ +#include + +torch::Tensor asu_cuda_forward(const torch::Tensor& input); +""" + +cuda_source = """ +#include +#include +#include + +// Vectorized type for 128-bit access with doubles +struct __align__(16) Double2 { + double x, y; +}; + +// Core computation +// Formula: x * sin(x) +__device__ __forceinline__ double asu_op(double x) { + return x * sin(x); +} + +__global__ void asu_kernel_double( + const double* __restrict__ input, + double* __restrict__ output, + const int n_elements) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + // 1. Vectorized Loop + int vec_loops = n_elements / 2; + const Double2* vec_input = reinterpret_cast(input); + Double2* vec_output = reinterpret_cast(output); + + for (int i = idx; i < vec_loops; i += stride) { + Double2 in_val = vec_input[i]; + Double2 out_val; + + out_val.x = asu_op(in_val.x); + out_val.y = asu_op(in_val.y); + + vec_output[i] = out_val; + } + + // 2. Scalar Loop + int tail_start = vec_loops * 2; + for (int i = tail_start + idx; i < n_elements; i += stride) { + output[i] = asu_op(input[i]); + } +} + +torch::Tensor asu_cuda_forward(const torch::Tensor& input) { + TORCH_CHECK(input.is_cuda(), "Input tensor must be a CUDA tensor"); + TORCH_CHECK(input.scalar_type() == torch::kDouble, "Input tensor must be float64"); + TORCH_CHECK(input.is_contiguous(), "Input tensor must be contiguous"); + + auto output = torch::empty_like(input); + const int n_elements = input.numel(); + + const int block_size = 256; + // Grid size for Double2 (2 elements per thread) + int grid_size = (n_elements + block_size * 2 - 1) / (block_size * 2); + if (grid_size > 65535) grid_size = 65535; + + asu_kernel_double<<>>( + input.data_ptr(), + output.data_ptr(), + n_elements + ); + + return output; +} +""" + +asu_op = load_inline( + name='asu_op', + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=['asu_cuda_forward'], + verbose=False, + extra_cuda_cflags=['-O3'] +) + +class ASUNew(nn.Module): + def __init__(self): + super(ASUNew, self).__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return asu_op.asu_cuda_forward(x) + +class ModelNew(nn.Module): + def __init__(self): + super(ModelNew, self).__init__() + self.act = ASUNew() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.act(x) \ No newline at end of file diff --git a/S1/hli28146_#30/asu_torch.py b/S1/hli28146_#30/asu_torch.py new file mode 100644 index 00000000..04a164d0 --- /dev/null +++ b/S1/hli28146_#30/asu_torch.py @@ -0,0 +1,35 @@ +import torch +import torch.nn as nn + +BATCH_SIZE = 4096 +DIM = 4096 +SHAPE = (BATCH_SIZE, DIM) + +DTYPE = torch.float64 + +class ASU(nn.Module): + """ + Amplifying Sine Unit: An Oscillatory Activation Function for Deep Neural Networks to Recover Nonlinear Oscillations Efficiently + https://arxiv.org/pdf/2304.09759 + Formula: f(x) = x * sin(x) + """ + def __init__(self): + super(ASU, self).__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.sin(x) + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.act = ASU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.act(x) + +def get_inputs(): + x = torch.randn(SHAPE, dtype=DTYPE) + return [x.contiguous()] + +def get_init_inputs(): + return [] \ No newline at end of file diff --git a/S1/hli28146_#30/prompt.txt b/S1/hli28146_#30/prompt.txt new file mode 100644 index 00000000..2e8f900b --- /dev/null +++ b/S1/hli28146_#30/prompt.txt @@ -0,0 +1,57 @@ +Write a custom CUDA kernel to optimize the ASU activation function as defined in the provided table. + +The mathematical definition is: +f(x) = x * sin(x) + +Problem Analysis: +1. Memory Bandwidth: The operation is element-wise and strictly memory-bound. The arithmetic intensity is low (one sin, one mul). Standard PyTorch implementation executes `sin(x)` followed by `x * result`, involving intermediate memory traffic. +2. Precision: Trigonometric functions are sensitive to precision. Double precision (float64) is required for strict accuracy alignment with the reference. + +Optimization Strategy: Fused Vectorized Kernel in Double Precision + +1. Data Type: Use `double` for all computations to guarantee numerical stability and accuracy. + +2. Vectorized Memory Access: Use `double2` types to load/store 128 bits (2 doubles) per instruction. This is the optimal transaction size for float64 data on GPUs, significantly reducing instruction overhead and maximizing bandwidth. + +3. Fused Computation: Compute `val * sin(val)` entirely in registers. This fuses the two element-wise operations into a single kernel pass (1 read, 1 write). + +4. Grid-Stride Loop: Implement a robust grid-stride loop to handle arbitrary input tensor sizes efficiently. + +Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: + +```python +import torch +import torch.nn as nn + +BATCH_SIZE = 4096 +DIM = 4096 +SHAPE = (BATCH_SIZE, DIM) + +DTYPE = torch.float64 + +class ASU(nn.Module): + """ + Amplifying Sine Unit: An Oscillatory Activation Function for Deep Neural Networks to Recover Nonlinear Oscillations Efficiently + https://arxiv.org/pdf/2304.09759 + Formula: f(x) = x * sin(x) + """ + def __init__(self): + super(ASU, self).__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.sin(x) + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.act = ASU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.act(x) + +def get_inputs(): + x = torch.randn(SHAPE, dtype=DTYPE) + return [x.contiguous()] + +def get_init_inputs(): + return [] \ No newline at end of file diff --git a/S1/hli28146_#30/run_code.py b/S1/hli28146_#30/run_code.py new file mode 100644 index 00000000..a6f45133 --- /dev/null +++ b/S1/hli28146_#30/run_code.py @@ -0,0 +1,74 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from asu_torch import Model,get_inputs,get_init_inputs +from asu_cuda 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 torch.relu 平均执行时间: {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