diff --git a/S1/22/ReflectionPad3d_cuda.py b/S1/22/ReflectionPad3d_cuda.py new file mode 100644 index 0000000..a75bc92 --- /dev/null +++ b/S1/22/ReflectionPad3d_cuda.py @@ -0,0 +1,213 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + +BATCH_SIZE = 8 +CHANNELS = 16 +DEPTH = 16 +HEIGHT = 16 +WIDTH = 16 + +PADDING = (1, 1, 2, 2, 1, 0) + +BLOCK_DIM_X = 8 +BLOCK_DIM_Y = 8 +BLOCK_DIM_Z = 8 + + +class ModelNew(nn.Module): + + def __init__(self, padding): + super().__init__() + + if isinstance(padding, int): + self.pad_L, self.pad_R, self.pad_T, self.pad_B, self.pad_F, self.pad_K = (padding,) * 6 + else: + self.pad_L, self.pad_R, self.pad_T, self.pad_B, self.pad_F, self.pad_K = padding + + self.block_dim_x = BLOCK_DIM_X + self.block_dim_y = BLOCK_DIM_Y + self.block_dim_z = BLOCK_DIM_Z + self._compile_cuda_kernel() + + def _compile_cuda_kernel(self): + + cpp_header = f""" + #include + + // C++ 接口 + torch::Tensor reflection_pad3d_forward_cuda( + torch::Tensor input, + int pad_L, int pad_R, + int pad_T, int pad_B, + int pad_F, int pad_K + ); + """ + + cuda_source = f""" + #include + #include + + #define BLOCK_DIM_X {self.block_dim_x} + #define BLOCK_DIM_Y {self.block_dim_y} + #define BLOCK_DIM_Z {self.block_dim_z} + + + __device__ inline int reflect_idx( + int j, int pad_before, int W_in + ) {{ + if (j < pad_before) {{ + return pad_before - j; + }} else if (j < (pad_before + W_in)) {{ + return j - pad_before; + }} else {{ + int j_rel = j - (pad_before + W_in); + return W_in - 2 - j_rel; + }} + }} + + + __global__ void reflection_pad3d_fused_kernel( + const float* __restrict__ input_data, + float* __restrict__ output_data, + int N, int C, + int D_in, int H_in, int W_in, + int D_out, int H_out, int W_out, + int pad_L, int pad_R, + int pad_T, int pad_B, + int pad_F, int pad_K + ) {{ + extern __shared__ float s_in[]; + + const int n_idx = blockIdx.x; + const int c_idx = blockIdx.y; + const int tid_x = threadIdx.x; + const int tid_y = threadIdx.y; + const int tid_z = threadIdx.z; + + + const int64_t H_in_stride = W_in; + const int64_t D_in_stride = H_in * W_in; + + const int64_t C_in_stride = D_in * D_in_stride; + + + const int64_t H_out_stride = W_out; + const int64_t D_out_stride = H_out * W_out; + + const int64_t C_out_stride = D_out * D_out_stride; + + + const float* p_in_base = input_data + (n_idx * C + c_idx) * C_in_stride; + float* p_out_base = output_data + (n_idx * C + c_idx) * C_out_stride; + + + for (int k = tid_z; k < D_in; k += BLOCK_DIM_Z) {{ + for (int i = tid_y; i < H_in; i += BLOCK_DIM_Y) {{ + for (int j = tid_x; j < W_in; j += BLOCK_DIM_X) {{ + // (k, i, j) -> 1D index + int64_t in_idx = k*D_in_stride + i*H_in_stride + j; + s_in[in_idx] = p_in_base[in_idx]; + }} + }} + }} + __syncthreads(); // 确保 s_in 加载完成 + + + for (int k = tid_z; k < D_out; k += BLOCK_DIM_Z) {{ + int in_k = reflect_idx(k, pad_F, D_in); + + for (int i = tid_y; i < H_out; i += BLOCK_DIM_Y) {{ + int in_i = reflect_idx(i, pad_T, H_in); + + for (int j = tid_x; j < W_out; j += BLOCK_DIM_X) {{ + int in_j = reflect_idx(j, pad_L, W_in); + + // 从共享内存读取 (使用 IN 步长) + int64_t s_in_idx = in_k*D_in_stride + in_i*H_in_stride + in_j; + + // 写入全局内存 (使用 OUT 步长) + int64_t p_out_idx = k*D_out_stride + i*H_out_stride + j; + + p_out_base[p_out_idx] = s_in[s_in_idx]; + }} + }} + }} + }} + + // C++ 封装函数 + torch::Tensor reflection_pad3d_forward_cuda( + torch::Tensor input, + int pad_L, int pad_R, + int pad_T, int pad_B, + int pad_F, int pad_K + ) {{ + TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor"); + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + TORCH_CHECK(input.dim() == 5, "input must be 5D (N, C, D, H, W)"); + + const int64_t N_64 = input.size(0); + const int64_t C_64 = input.size(1); + const int64_t D_in_64 = input.size(2); + const int64_t H_in_64 = input.size(3); + const int64_t W_in_64 = input.size(4); + + TORCH_CHECK(pad_L < W_in_64, "pad_L error"); + TORCH_CHECK(pad_R < W_in_64, "pad_R error"); + TORCH_CHECK(pad_T < H_in_64, "pad_T error"); + TORCH_CHECK(pad_B < H_in_64, "pad_B error"); + TORCH_CHECK(pad_F < D_in_64, "pad_F error"); + TORCH_CHECK(pad_K < D_in_64, "pad_K error"); + + const int64_t D_out_64 = D_in_64 + pad_F + pad_K; + const int64_t H_out_64 = H_in_64 + pad_T + pad_B; + const int64_t W_out_64 = W_in_64 + pad_L + pad_R; + + auto output = torch::empty({{N_64, C_64, D_out_64, H_out_64, W_out_64}}, input.options()); + + dim3 grid_dim(N_64, C_64); + dim3 block_dim(BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z); + + + const int shared_mem_size = D_in_64 * H_in_64 * W_in_64 * sizeof(float); + + reflection_pad3d_fused_kernel<<>>( + input.data_ptr(), + output.data_ptr(), + static_cast(N_64), static_cast(C_64), + static_cast(D_in_64), static_cast(H_in_64), static_cast(W_in_64), + static_cast(D_out_64), static_cast(H_out_64), static_cast(W_out_64), + pad_L, pad_R, + pad_T, pad_B, + pad_F, pad_K + ); + + return output; + }} + """ + + nvcc_flags = [ + '-O3', + '--use_fast_math', + '--expt-relaxed-constexpr' + ] + + self.pad_op = load_inline( + name="reflection_pad3d_op_v2_fixed", + cpp_sources=cpp_header, + cuda_sources=cuda_source, + functions=["reflection_pad3d_forward_cuda"], + extra_cuda_cflags=nvcc_flags, + verbose=False + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + + x_cont = x.contiguous() + + return self.pad_op.reflection_pad3d_forward_cuda( + x_cont, + self.pad_L, self.pad_R, + self.pad_T, self.pad_B, + self.pad_F, self.pad_K + ) \ No newline at end of file diff --git a/S1/22/ReflectionPad3d_torch.py b/S1/22/ReflectionPad3d_torch.py new file mode 100644 index 0000000..a836516 --- /dev/null +++ b/S1/22/ReflectionPad3d_torch.py @@ -0,0 +1,40 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +BATCH_SIZE = 8 +CHANNELS = 16 +DEPTH = 16 # D_in +HEIGHT = 16 # H_in +WIDTH = 16 # W_in + +PADDING = (1, 1, 2, 2, 1, 0) + + +# ------------------------------------------------------------- + +class Model(nn.Module): + + def __init__(self, padding): + super().__init__() + + if isinstance(padding, int): + # F.pad 需要 6-tuple + self.padding_tuple = (padding,) * 6 + else: + self.padding_tuple = padding + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # F.pad 5D 张量 (N, C, D, H, W) + # 填充顺序: (pad_W_L, pad_W_R, pad_H_T, pad_H_B, pad_D_F, pad_D_K) + # 这与 nn.ReflectionPad3d 的构造函数顺序一致 + return F.pad(x, self.padding_tuple, mode='reflect') + + +def get_inputs(): + x = torch.randn(BATCH_SIZE, CHANNELS, DEPTH, HEIGHT, WIDTH, dtype=torch.float32) + return [x] + + +def get_init_inputs(): + return [PADDING] diff --git a/S1/22/prompt.txt b/S1/22/prompt.txt new file mode 100644 index 0000000..b6d514f --- /dev/null +++ b/S1/22/prompt.txt @@ -0,0 +1,82 @@ +You write custom CUDA kernels to replace the PyTorch operators in the given EvoNorm 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 some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining normalization+affine_transform+nonlinear_gating), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination. + + +The provided code implements a custom CUDA kernel for 3D reflection padding in PyTorch using several advanced techniques: + +Key Technologies Used: + +Inline CUDA Extension in PyTorch: Uses torch.utils.cpp_extension.load_inline to compile and load CUDA code directly within Python, avoiding separate compilation steps. + +Fused GPU Kernel Design: Implements a single kernel that handles both data loading and padding operations, reducing kernel launch overhead. + +Shared Memory Optimization: Leverages CUDA shared memory (s_in[]) to cache input data, enabling faster data access patterns compared to global memory. + +Multi-dimensional Thread Blocking: Employs 3D thread blocks (BLOCK_DIM_X/Y/Z) for efficient parallelization across depth, height, and width dimensions. + +Strided Memory Access Patterns: Calculates explicit strides for both input and output tensors to optimize memory access. + +Reflective Index Calculation: Implements a device-side reflect_idx function that handles boundary reflection using mathematical calculations rather than conditional branching. + +Grid-Strided Loops: Uses grid-strided loops in the kernel to handle arbitrary output sizes while maintaining coalesced memory access. + +Batched Channel Processing: Processes multiple batches and channels concurrently through grid dimensions (grid_dim(N, C)). + +Runtime Bounds Checking: Includes comprehensive error checking for tensor dimensions and padding values. + +Memory Contiguity Enforcement: Ensures input tensor is contiguous for optimal memory access patterns. + +Performance Optimizations: + +Shared memory caching of input data + +Coalesced global memory accesses + +Minimal synchronization points (single __syncthreads()) + +Compiler optimizations (-O3, --use_fast_math) + +Grid-strided loops for load balancing + + +Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: +import torch +import torch.nn as nn +import torch.nn.functional as F + +BATCH_SIZE = 8 +CHANNELS = 16 +DEPTH = 16 # D_in +HEIGHT = 16 # H_in +WIDTH = 16 # W_in + +PADDING = (1, 1, 2, 2, 1, 0) + + +# ------------------------------------------------------------- + +class Model(nn.Module): + + def __init__(self, padding): + super().__init__() + + if isinstance(padding, int): + # F.pad 需要 6-tuple + self.padding_tuple = (padding,) * 6 + else: + self.padding_tuple = padding + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # F.pad 5D 张量 (N, C, D, H, W) + # 填充顺序: (pad_W_L, pad_W_R, pad_H_T, pad_H_B, pad_D_F, pad_D_K) + # 这与 nn.ReflectionPad3d 的构造函数顺序一致 + return F.pad(x, self.padding_tuple, mode='reflect') + + +def get_inputs(): + x = torch.randn(BATCH_SIZE, CHANNELS, DEPTH, HEIGHT, WIDTH, dtype=torch.float32) + return [x] + + +def get_init_inputs(): + return [PADDING] \ No newline at end of file diff --git a/S1/22/run_code.py b/S1/22/run_code.py new file mode 100644 index 0000000..1b8eb0e --- /dev/null +++ b/S1/22/run_code.py @@ -0,0 +1,77 @@ +########################################################### +# 性能和精度验证程序 +########################################################### +import torch +import torch.nn as nn +import time +from ReflectionPad3d_torch import Model, get_inputs, get_init_inputs +from ReflectionPad3d_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