Building a GPU Compiler from Scratch: Lowering PyTorch Graphs to CUDA via Multi-Level IRs
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.