Building a GPU Compiler from Scratch: Lowering PyTorch Graphs to CUDA via Multi-Level IRs
Summary
This talk provides a deep dive into modern ML compiler technology, demonstrating how high-level PyTorch graph operations are progressively lowered into highly optimized GPU kernels (CUDA). The core concept is progressive IR lowering, which transforms operations through multiple Intermediate Representations (IRs)—from PyTorch to Tensor AR, Loop IR, Tile IR, and finally Kernel IR—to enable crucial optimizations like operation fusion, normalization, and tiling. This process ensures that the final kernel is efficient, portable, and achieves near-hardware-limit performance.
Key takeaways
-
Progressive IR Lowering Stack
2:00
The compilation process moves through a structured stack of IRs: PyTorch (high-level ops) $\rightarrow$ Tensor AR (atomic, universal ops) $\rightarrow$ Loop IR (enables fusion) $\rightarrow$ Tile IR (algebraic scheduling) $\rightarrow$ Kernel IR (abstract CUDA) $\rightarrow$ CUDA (final output). This layered approach ensures portability and optimization at each stage.
-
Operation Fusion via Loop IR
4:10
Loop IR is critical because it allows for operation fusion, minimizing costly round trips to global memory. By aligning the iteration space of multiple operations, the compiler can combine them into a single, efficient kernel (e.g., fusing $x-x$ into one loop).
-
Tiling and Scheduling (Tile IR)
5:50
The Tile IR represents the algebra of the operation (using 'fold' operations) rather than the exact computation. This allows the compiler to enumerate scheduling choices, such as shared memory reduction or register tiling, which are hardware-aware decisions made late in the pipeline.
-
Compiler Optimization Philosophy
6:40
The system is designed to fuse everything possible at the Loop IR level, deferring hardware-specific decisions (like cost modeling or scheduling) until the Tile IR or Kernel IR stages. This keeps the lower IRs portable.
Technical details
-
IR Layers and Transformation
170s
The process begins by mirroring PyTorch operations to the Tensor AR level, which uses a primitive, NumPy-like representation to handle different framework weight layouts (e.g., TensorFlow's row major vs. PyTorch's column major). The RMS norm example demonstrates this transformation.
-
SSA Form and Data Flow
280s
The entire program is represented in Static Single Assignment (SSA) form—a list of statements where each statement has a unique output name. Fusion relies on analyzing the Directed Acyclic Graph (DAG) to identify producer and consumer operations.
-
Normalization and Kernel Identity
310s
Normalization is used both to optimize execution (e.g., eliminating common subexpressions) and to ensure that different internal representations of the same mathematical operation (e.g., $2x+y$ vs. $y+2x$) are lowered to the identical Loop IR.
-
Tiling Operations (Fold)
350s
The Tile IR uses 'fold' operations to represent the algebra of the computation. Key components include the reduction (e.g., sum over an axis), the combine operation (e.g., addition), and the initial value. This structure allows for advanced scheduling like register tiling and defining input pipelines.
Mentioned resources
- cloudrift.ai / riftstack.ai
- kernelstack.com
Channel & topics
Watch on YouTube · Back to latest
This independent, AI-assisted summary is provided for commentary and informational purposes. It may contain errors or omit important context. Please watch the original video for the creator's complete presentation. Video, thumbnail, and related copyrights belong to their respective owners.