From 221bc114bf592a1fd72e23b8a4e2b2b92ff36cf4 Mon Sep 17 00:00:00 2001 From: uucoco Date: Tue, 9 Dec 2025 17:39:49 +0800 Subject: [PATCH] finish DoubleGLU #44 --- S1/uucoco_#44/DoubleGLU_cuda.py | 117 ++++++++++++++++++ .../{Relu6_torch.py => DoubleGLU_torch.py} | 10 +- S1/uucoco_#44/Relu6_cuda.py | 84 ------------- S1/uucoco_#44/prompt.txt | 55 +++++--- S1/uucoco_#44/run_code.py | 4 +- 5 files changed, 163 insertions(+), 107 deletions(-) create mode 100644 S1/uucoco_#44/DoubleGLU_cuda.py rename S1/uucoco_#44/{Relu6_torch.py => DoubleGLU_torch.py} (68%) delete mode 100644 S1/uucoco_#44/Relu6_cuda.py diff --git a/S1/uucoco_#44/DoubleGLU_cuda.py b/S1/uucoco_#44/DoubleGLU_cuda.py new file mode 100644 index 00000000..258b7652 --- /dev/null +++ b/S1/uucoco_#44/DoubleGLU_cuda.py @@ -0,0 +1,117 @@ +import torch +import torch.nn as nn +from torch.utils.cpp_extension import load_inline + + +class ModelNew(nn.Module): + def __init__(self): + super().__init__() + self._compile_cuda_kernel() + + def _compile_cuda_kernel(self): + cpp_source = """ + #include + torch::Tensor double_glu_cuda(torch::Tensor input); + """ + + cuda_source = """ + #include + + __device__ __forceinline__ float sigmoid_f(float x) { + return 1.0f / (1.0f + expf(-x)); + } + + __global__ void double_glu_vec4_kernel( + const float4* __restrict__ x, + float4* __restrict__ y, + int chunk_vec_dim, + int total_chunk_vecs) + { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + + for (int i = idx; i < total_chunk_vecs; i += stride) { + int row = i / chunk_vec_dim; + int col = i % chunk_vec_dim; + + // Input width = 4 * chunk + // Output width = 2 * chunk + + int row_offset_in = row * 4 * chunk_vec_dim; + int row_offset_out = row * 2 * chunk_vec_dim; + + // Process Pair 1 (G1, X1) + int g1_idx = row_offset_in + col; + int x1_idx = row_offset_in + chunk_vec_dim + col; + + float4 g1 = x[g1_idx]; + float4 x1 = x[x1_idx]; + float4 out1; + + out1.x = sigmoid_f(g1.x) * x1.x; + out1.y = sigmoid_f(g1.y) * x1.y; + out1.z = sigmoid_f(g1.z) * x1.z; + out1.w = sigmoid_f(g1.w) * x1.w; + + y[row_offset_out + col] = out1; + + // Process Pair 2 (G2, X2) + int g2_idx = row_offset_in + 2 * chunk_vec_dim + col; + int x2_idx = row_offset_in + 3 * chunk_vec_dim + col; + + float4 g2 = x[g2_idx]; + float4 x2 = x[x2_idx]; + float4 out2; + + out2.x = sigmoid_f(g2.x) * x2.x; + out2.y = sigmoid_f(g2.y) * x2.y; + out2.z = sigmoid_f(g2.z) * x2.z; + out2.w = sigmoid_f(g2.w) * x2.w; + + y[row_offset_out + chunk_vec_dim + col] = out2; + } + } + + torch::Tensor double_glu_cuda(torch::Tensor input) { + auto x_c = input.contiguous(); + + int last_dim = x_c.size(-1); + TORCH_CHECK(last_dim % 16 == 0, "Feature dim must be divisible by 16 (4 chunks * float4) for optimization"); + + auto out_sizes = x_c.sizes().vec(); + out_sizes.back() /= 2; + auto output = torch::empty(out_sizes, x_c.options()); + + int chunk_dim = last_dim / 4; + int chunk_vec_dim = chunk_dim / 4; + + int batch_size = x_c.numel() / last_dim; + int total_chunk_vecs = batch_size * chunk_vec_dim; + + int threads = 256; + int blocks = (total_chunk_vecs + threads - 1) / threads; + if (blocks > 65535) blocks = 65535; + if (blocks == 0) blocks = 1; + + double_glu_vec4_kernel<<>>( + reinterpret_cast(x_c.data_ptr()), + reinterpret_cast(output.data_ptr()), + chunk_vec_dim, + total_chunk_vecs + ); + + return output; + } + """ + + self.op = load_inline( + name="double_glu_opt_vec4", + cpp_sources=cpp_source, + cuda_sources=cuda_source, + functions=["double_glu_cuda"], + extra_cuda_cflags=["-O3", "--use_fast_math"], + verbose=False + ) + + def forward(self, x): + return self.op.double_glu_cuda(x) \ No newline at end of file diff --git a/S1/uucoco_#44/Relu6_torch.py b/S1/uucoco_#44/DoubleGLU_torch.py similarity index 68% rename from S1/uucoco_#44/Relu6_torch.py rename to S1/uucoco_#44/DoubleGLU_torch.py index e56fcb8c..f33ef867 100644 --- a/S1/uucoco_#44/Relu6_torch.py +++ b/S1/uucoco_#44/DoubleGLU_torch.py @@ -2,24 +2,20 @@ import torch import torch.nn as nn import torch.nn.functional as F - class Model(nn.Module): def __init__(self): super().__init__() def forward(self, x: torch.Tensor) -> torch.Tensor: - # 直接使用 PyTorch 內建的優化函式 F.relu6 - return F.relu6(x) - + g1, x1, g2, x2 = x.chunk(4, dim=-1) + return torch.cat([torch.sigmoid(g1) * x1, torch.sigmoid(g2) * x2], dim=-1) batch_size = 128 -feature_dim = 512 - +feature_dim = 4096 def get_inputs(): x = torch.randn(batch_size, feature_dim, dtype=torch.float32) return [x] - def get_init_inputs(): return [] \ No newline at end of file diff --git a/S1/uucoco_#44/Relu6_cuda.py b/S1/uucoco_#44/Relu6_cuda.py deleted file mode 100644 index f03ed90f..00000000 --- a/S1/uucoco_#44/Relu6_cuda.py +++ /dev/null @@ -1,84 +0,0 @@ -import torch -import torch.nn as nn -from torch.utils.cpp_extension import load_inline - -class ModelNew(nn.Module): - def __init__(self): - super().__init__() - self._compile_cuda_kernel() - - def _compile_cuda_kernel(self): - cpp_source = """ - torch::Tensor relu6_cuda(torch::Tensor x); - """ - - cuda_source = """ - #include - #include - #include - - __device__ __forceinline__ float relu6_op(float x) { - // f(x) = min(6, max(0, x)) - return fminf(6.0f, fmaxf(0.0f, x)); - } - - __global__ void relu6_kernel( - const float* __restrict__ x, - float* __restrict__ output, - const int n_elements) - { - const int tid = blockIdx.x * blockDim.x + threadIdx.x; - const int stride = blockDim.x * gridDim.x; - - const int vec_loops = n_elements >> 2; - const float4* x_vec = reinterpret_cast(x); - float4* out_vec = reinterpret_cast(output); - - for (int i = tid; i < vec_loops; i += stride) { - float4 v = __ldg(&x_vec[i]); - float4 r; - - r.x = relu6_op(v.x); - r.y = relu6_op(v.y); - r.z = relu6_op(v.z); - r.w = relu6_op(v.w); - - out_vec[i] = r; - } - - const int tail_start = vec_loops << 2; - for (int i = tail_start + tid; i < n_elements; i += stride) { - output[i] = relu6_op(x[i]); - } - } - - torch::Tensor relu6_cuda(torch::Tensor x) { - auto x_c = x.contiguous(); - const int n_elements = x_c.numel(); - auto output = torch::empty_like(x_c); - - const int threads = 256; - const int max_blocks = 65535; - const int blocks = std::min((n_elements + threads * 4 - 1) / (threads * 4), max_blocks); - - relu6_kernel<<>>( - x_c.data_ptr(), - output.data_ptr(), - n_elements - ); - - return output; - } - """ - - self.op = load_inline( - name="relu6_v1", - cpp_sources=cpp_source, - cuda_sources=cuda_source, - functions=["relu6_cuda"], - extra_cuda_cflags=["-O3", "--use_fast_math"], - verbose=False - ) - - def forward(self, x): - return self.op.relu6_cuda(x) \ No newline at end of file diff --git a/S1/uucoco_#44/prompt.txt b/S1/uucoco_#44/prompt.txt index 496a80df..310e3481 100644 --- a/S1/uucoco_#44/prompt.txt +++ b/S1/uucoco_#44/prompt.txt @@ -2,38 +2,65 @@ You write custom CUDA kernels to replace the pytorch operators in the given GeGL 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 chunk+gelu+elementwise_mul), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination. -This CUDA kernel implements several key optimizations for the ReLU6 activation function: -Vectorization: Uses float4memory operations to process 4 elements per thread, increasing memory throughput. -Memory Coalescing: Accesses contiguous memory blocks through vector loads/stores, optimizing GPU memory bandwidth. -Grid-Stride Loop: Handles arbitrary-sized tensors efficiently by having threads process multiple elements with strided indexing. -Fast Math: Uses CUDA's fminf/fmaxfintrinsics and --use_fast_mathcompiler flag for optimized mathematical operations. -Occupancy Optimization: Employs 256 threads per block and calculates optimal grid size to maximize GPU utilization. -Tail Processing: Handles non-multiple-of-4 elements separately after vectorized operations. +This CUDA kernel implements optimized Double Gated Linear Unit (Double GLU) with: +Memory Optimization: + +Vectorized memory access using float4 for 4x bandwidth + +Contiguous tensor inputs for coalesced memory access + +Processes two GLU pairs simultaneously per thread + +Parallelization Strategy: + +Grid-stride loop for efficient workload distribution + +256 threads per block optimal configuration + +Automatic grid size calculation with 65535 block limit + +Computational Optimization: + +Double GLU: Two parallel GLU operations sigmoid(gate) * activation + +Optimized sigmoid: 1.0f / (1.0f + expf(-x)) + +Fast math compilation flags for optimized exponential + +Efficient indexing for four input chunks (G1, X1, G2, X2) + +Work Distribution: + +Each thread processes 8 total elements (4 per GLU pair) via float4 + +Processes two independent GLU operations simultaneously + +Input divided into four equal chunks, output into two chunks + +Requires input feature dimension divisible by 16 for optimal performance + +The implementation maximizes throughput by processing two GLU operations in parallel through vectorization and efficient memory access patterns. 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 - class Model(nn.Module): def __init__(self): super().__init__() def forward(self, x: torch.Tensor) -> torch.Tensor: - # 直接使用 PyTorch 內建的優化函式 F.relu6 - return F.relu6(x) - + g1, x1, g2, x2 = x.chunk(4, dim=-1) + return torch.cat([torch.sigmoid(g1) * x1, torch.sigmoid(g2) * x2], dim=-1) batch_size = 128 -feature_dim = 512 - +feature_dim = 4096 def get_inputs(): x = torch.randn(batch_size, feature_dim, dtype=torch.float32) return [x] - def get_init_inputs(): return [] \ No newline at end of file diff --git a/S1/uucoco_#44/run_code.py b/S1/uucoco_#44/run_code.py index 185561fe..5f0f3b6b 100644 --- a/S1/uucoco_#44/run_code.py +++ b/S1/uucoco_#44/run_code.py @@ -4,8 +4,8 @@ import torch import torch.nn as nn import time -from Relu6_torch import Model, get_inputs, get_init_inputs -from Relu6_cuda import ModelNew +from DoubleGLU_torch import Model, get_inputs, get_init_inputs +from DoubleGLU_cuda import ModelNew def run_benchmark():