forked from ccf-ai-infra/GPUCodeForces
finish DoubleGLU #44
This commit is contained in:
parent
f62d76a9f7
commit
221bc114bf
|
|
@ -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/extension.h>
|
||||
torch::Tensor double_glu_cuda(torch::Tensor input);
|
||||
"""
|
||||
|
||||
cuda_source = """
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
__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<<<blocks, threads>>>(
|
||||
reinterpret_cast<const float4*>(x_c.data_ptr<float>()),
|
||||
reinterpret_cast<float4*>(output.data_ptr<float>()),
|
||||
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)
|
||||
|
|
@ -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 []
|
||||
|
|
@ -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 <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
|
||||
__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<const float4*>(x);
|
||||
float4* out_vec = reinterpret_cast<float4*>(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<<<blocks, threads>>>(
|
||||
x_c.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
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)
|
||||
|
|
@ -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 []
|
||||
|
|
@ -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():
|
||||
|
|
|
|||
Loading…
Reference in New Issue