forked from ccf-ai-infra/GPUCodeForces
107 lines
3.8 KiB
Plaintext
107 lines
3.8 KiB
Plaintext
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.
|
|
|
|
Technologies Used :
|
|
|
|
PyTorch: Deep learning framework
|
|
|
|
CUDA: GPU acceleration for parallel computing
|
|
|
|
C++/CUDA C++: High-performance kernel programming
|
|
|
|
Inline C++/CUDA Extension (torch.utils.cpp_extension.load_inline): Just-In-Time (JIT) compilation of custom operators
|
|
|
|
Tanimoto Coefficient (Jaccard Index): Similarity measure between sets or vectors
|
|
|
|
Fused Kernel Design: Combines dot product, squared sums, and division in single kernel
|
|
|
|
Vectorized Memory Access (float4): Uses 128-bit wide loads (4 floats) for improved memory bandwidth
|
|
|
|
One Block Per Sample: Each CUDA block processes one complete instance (N blocks for N samples)
|
|
|
|
Shared Memory Reduction: Uses __shared__ arrays for efficient block-level parallel reduction
|
|
|
|
Tree Reduction Pattern: Binary tree reduction within thread blocks using __syncthreads()
|
|
|
|
Precision Control: Uses volatile keyword to disable FMA (Fused Multiply-Add) for exact numerical matching
|
|
|
|
Compiler Precision Flags: Uses -ftz=false, -prec-div=true, -prec-sqrt=true for high precision
|
|
|
|
Dynamic Shared Memory Allocation: Allocates shared memory for three reduction buffers at kernel launch
|
|
|
|
Strided Memory Access: Threads process elements with stride equal to block size for coalesced access
|
|
|
|
Numerical Stability: Adds epsilon (eps) to prevent division by zero
|
|
|
|
Memory Coalescing: Ensures contiguous tensor layout for optimal memory access patterns
|
|
|
|
Automatic Type Conversion: Converts inputs to float32 and CUDA device if needed
|
|
|
|
Boundary Checking: Validates block indices to prevent out-of-bounds access
|
|
|
|
Three-Accumulator Design: Maintains separate accumulators for dot product, x squared, and y squared
|
|
|
|
Exact Numerical Matching: Designed to 100% match PyTorch's step-by-step floating point rounding
|
|
|
|
Optimized Memory Layout: Uses vectorized pointer arithmetic with float4 type casting
|
|
|
|
Efficient Resource Utilization: Maximizes memory bandwidth while maintaining numerical precision
|
|
|
|
|
|
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
|
|
|
|
# 定义维度常量
|
|
N, C, H, W = 32, 64, 56, 56
|
|
EPS = 1e-6
|
|
|
|
|
|
class Tanimoto(nn.Module):
|
|
|
|
|
|
def __init__(self, eps=1e-6):
|
|
super().__init__()
|
|
self.eps = eps
|
|
# 我们将在 C, H, W 维度上进行归约
|
|
self.reduction_dims = (1, 2, 3)
|
|
|
|
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
|
|
x_dot_y = torch.sum(x * y, dim=self.reduction_dims)
|
|
|
|
|
|
x_norm_sq = torch.sum(x * x, dim=self.reduction_dims)
|
|
y_norm_sq = torch.sum(y * y, dim=self.reduction_dims)
|
|
|
|
|
|
denominator = x_norm_sq + y_norm_sq - x_dot_y
|
|
|
|
|
|
similarity = (x_dot_y + self.eps) / (denominator + self.eps)
|
|
|
|
return similarity
|
|
|
|
|
|
class Model(nn.Module):
|
|
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.op = Tanimoto(EPS)
|
|
|
|
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return self.op(x, y)
|
|
|
|
|
|
def get_inputs():
|
|
|
|
x = torch.randn(N, C, H, W, dtype=torch.float32)
|
|
y = torch.randn(N, C, H, W, dtype=torch.float32)
|
|
return [x, y]
|
|
|
|
|
|
def get_init_inputs():
|
|
|
|
return [] |