forked from ccf-ai-infra/GPUCodeForces
139 lines
5.1 KiB
Plaintext
139 lines
5.1 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.
|
||
|
||
Poisson Negative Log-Likelihood Loss CUDA Optimization with Fused Kernel
|
||
CUDA Optimization Techniques
|
||
1. Parallel Reduction Architecture
|
||
Grid-Stride Loop Pattern: Each thread processes multiple elements with stride gridDim.x * blockDim.x
|
||
|
||
Block-Level Reduction: Partial sums computed in shared memory
|
||
|
||
Dynamic Grid Sizing: grid_size = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE
|
||
|
||
2. Fused Kernel Design
|
||
Single Kernel Execution: Combines all loss computations in one kernel launch
|
||
|
||
Branch Handling: Efficiently handles log_input and full mode conditions
|
||
|
||
Mathematical Fusion: Integrates exponential, logarithmic, and conditional operations
|
||
|
||
3. Memory Access Optimization
|
||
Coalesced Memory Access: Sequential reading of input and target tensors
|
||
|
||
Shared Memory Utilization: s_data[BLOCK_SIZE] for block-level sum reduction
|
||
|
||
Contiguous Tensors: Ensures input and target are contiguous in memory
|
||
|
||
4. Numerical Stability Features
|
||
Epsilon Protection: eps_val prevents log(0) in non-log mode
|
||
|
||
Stirling's Approximation: Conditional Stirling term for full mode when y > 1.0f
|
||
|
||
Float Safety: Proper handling of edge cases and special values
|
||
|
||
5. Mathematical Operations
|
||
7. Performance Optimizations
|
||
Minimal Global Memory Writes: Only thread 0 writes block sum to global memory
|
||
|
||
Efficient Thread Utilization: All threads participate in computation and reduction
|
||
|
||
Load Balancing: Grid-stride loops handle arbitrary tensor sizes
|
||
|
||
Constant Propagation: PI and configuration parameters as compile-time constants
|
||
|
||
8. Implementation Features
|
||
Configuration Flexibility: Supports log_input, full, and eps parameters
|
||
|
||
Comprehensive Validation: Tensor device, contiguity, and size checking
|
||
|
||
Edge Case Handling: Empty tensor detection and proper zero handling
|
||
|
||
PyTorch Integration: Seamless tensor passing and automatic differentiation support
|
||
|
||
Key CUDA Concepts Used
|
||
Grid-Stride Loops for workload distribution across all elements
|
||
|
||
Shared Memory Reduction for parallel sum computation
|
||
|
||
Conditional Execution for handling different mathematical modes
|
||
|
||
Memory Coalescing for efficient global memory access
|
||
|
||
Kernel Fusion combining multiple mathematical operations
|
||
|
||
Workflow Summary
|
||
Configuration Setup: Parse log_input, full, and eps parameters
|
||
|
||
Memory Preparation: Ensure contiguous tensor layouts
|
||
|
||
Kernel Launch: Execute fused Poisson NLL computation with parallel reduction
|
||
|
||
Final Reduction: Sum block partial sums and compute mean loss
|
||
|
||
Mathematical Components
|
||
Exponential Computation: expf() for log-input mode
|
||
|
||
Logarithmic Computation: logf() for non-log mode and Stirling term
|
||
|
||
Stirling's Approximation: Complete term for Poisson distribution normalization
|
||
|
||
Element-wise Operations: Parallel computation across all tensor elements
|
||
|
||
Expected Performance Benefits
|
||
2-4x speedup over PyTorch implementation for large tensors
|
||
|
||
Reduced kernel launches through operation fusion
|
||
|
||
Better memory efficiency through coalesced access patterns
|
||
|
||
Scalable performance with increasing tensor sizes
|
||
|
||
This implementation provides a production-ready Poisson NLL Loss with significant performance improvements through careful CUDA optimization, parallel reduction patterns, and numerical stability considerations.
|
||
|
||
|
||
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 = 4096
|
||
FEATURE_DIM = 512
|
||
|
||
# --- 损失函数的参数 ---
|
||
# log_input=True: loss = exp(input) - target * input
|
||
# log_input=False: loss = input - target * log(input + eps)
|
||
LOG_INPUT = True
|
||
# full=True: 添加 Stirling's approximation
|
||
FULL = False
|
||
EPS = 1e-8
|
||
|
||
|
||
class Model(nn.Module):
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.criterion = nn.PoissonNLLLoss(
|
||
log_input=LOG_INPUT,
|
||
full=FULL,
|
||
eps=EPS,
|
||
reduction='mean'
|
||
)
|
||
|
||
def forward(self, input_tensor: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||
# 在 PyTorch 中，input_tensor 是文档中的 'input'
|
||
return self.criterion(input_tensor, target)
|
||
|
||
|
||
def get_inputs():
|
||
# Input (log_input=True 时) 可以是任意实数
|
||
input_tensor = torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32)
|
||
|
||
# Target 在 Poisson 分布中代表计数，且在 'full' 模式下会计算 log(target)
|
||
# 因此 target 必须是 >= 0 的。我们使用 rand 来确保
|
||
target = torch.rand(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32) * 10 # 乘以 10 以便有一些 > 1
|
||
|
||
return [input_tensor, target]
|
||
|
||
|
||
def get_init_inputs():
|
||
return [] |