# Building a GPU Compiler from Scratch: Lowering PyTorch Graphs to CUDA via Multi-Level IRs

## Executive 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: 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: 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): 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: 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: 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: 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: 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): 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.

## Practical implications

- The compiler design principles presented are highly applicable to optimizing complex deep learning models for deployment on specialized hardware (GPUs).
- Understanding the IR stack (especially the transition from Loop IR to Tile IR) is crucial for building next-generation, high-performance ML compilers.
- The ability to decouple the IR from hardware specifics (keeping it pure up to the Tile IR) allows for maximum portability and future-proofing the compiler stack.

## Topics

ML Compilers, GPU Architecture, Intermediate Representation (IR), Compiler Optimization, PyTorch, CUDA, cloudrift.ai / riftstack.ai, kernelstack.com

Source: https://www.youtube.com/watch?v=EKf0cPYq4e4