ml-compiler

How it is built

Architecture overview

An ML compiler takes a model written as a graph of tensor operations (matmul, convolution, softmax, add, reshape) and produces machine code for a specific chip that runs it faster than calling one pre-written library kernel per operation. It works in layers. A graph optimizer rewrites the operator graph (infers shapes, simplifies, fuses operators, picks tensor layouts, plans memory). A loop and tile layer turns each fused operator into a loop nest and decides how to tile, vectorise and map it to threads. A code generator lowers that to an SSA IR, selects instructions and allocates registers. A schedule search loop measures candidate implementations on the real hardware and keeps the best, and a runtime compiles on demand and caches the result.

The one idea that makes it work is the split between what is computed and how it is scheduled. The tensor expression stays fixed; tiling, ordering, fusion and vector width are separate decisions that can be searched, so performance work becomes search over schedules guided by a cost model and checked by measurement, not hand-tuning one kernel per operator and per chip. XLA, TVM and Triton are the best-known open systems; they differ in where they put the split (whole graph, tensor expressions, or tile programs written by the user).

Parts:

• Capture and JIT driver — traces or imports the model, compiles on first call, caches by shape and dtype — JIT compilation
• Graph IR and analyses — SSA operator graph, shapes, liveness — SSA form, Dataflow analysis, Shape inference (no page yet)
• Graph optimizer — fuses operators, rewrites algebraically, chooses layouts, splits the graph across devices — Kernel fusion, E-graphs (no page yet), Tensor layout transforms (no page yet), Operator graph partitioning (no page yet)
• Memory planner — assigns every intermediate tensor an offset in one arena — Tensor memory planning (no page yet)
• Loop and tile layer — turns one fused operator into a tiled, vectorised, thread-mapped loop nest — Loop optimizations (no page yet), Tiled matrix multiplication, SIMD programming, CUDA and GPU languages
• Schedule search — tries candidate schedules, measures them, learns a cost model — Kernel autotuning (no page yet), Roofline model
• Code generator — SSA to instructions to registers, emitting machine code or PTX — Instruction selection, Register allocation, SSA form

 model code ─► capture / trace ─► graph IR (SSA, shapes known)
                                      │  graph passes: shape inference, simplify, e-graph rewrite,
                                      │  fusion, layout, partition, memory plan
                                      ▼
                         fused kernels (tensor expressions)
                                      │  schedule = tile sizes, loop order, vector width, thread map
                    ┌─────────────────┴──────────────┐
                    │ autotuner: propose ► build ► run on device ► cost model update
                    └─────────────────┬──────────────┘
                                      ▼
                       loop-level IR (SSA) ─► instruction selection ─► register allocation
                                      ▼
                       machine code / PTX ─► cache keyed by (graph, shapes, dtype, target)
                                      ▼
                       runtime: launch kernels over the planned arena

How the parts fit together

Capture and JIT driver

The compiler first needs a graph. Frameworks either trace Python execution (PyTorch's Dynamo, JAX's jit) or import a serialised graph. The driver is a JIT compiler in the strict sense: the first call with a given set of input shapes and dtypes compiles, later calls reuse the cached binary. The contract with the rest of the system is a cache key and guards: the key names the graph, shapes, dtypes and target; guards are cheap checks that the cached code is still valid for this call. In PyTorch 2 these guards come from TorchDynamo, and a failed guard triggers recompilation (PyTorch 2, ASPLOS 2024, see Sources).

Graph IR and analyses

The graph is kept in SSA form: every tensor value is defined once, so a rewrite can reason about a value without asking who else wrote to it. That makes the classic analyses cheap. Dataflow analysis gives liveness (when a tensor is last read, needed by the memory planner), use counts (can this producer be fused into its single consumer, or must its output be materialised for several?) and constant propagation. Shape inference (no page yet) propagates shapes forward through operator rules (a matmul of [m,k] by [k,n] gives [m,n], a reshape checks element counts). Almost every later decision depends on it: tile sizes, buffer sizes, whether two operators can fuse, whether a broadcast is needed. The seam to watch is dynamic shapes. With static shapes every size is a constant and everything can be planned at compile time; with symbolic sizes the compiler carries expressions (batch, seq_len), must prove equalities between them, and plans memory in terms of formulas, or compiles one specialisation per size and relies on the cache.

Graph optimizer

Four jobs share this stage; the first is usually worth the most.

• Fusion. Elementwise operations (add, scale, activation) do little arithmetic per byte moved, so run alone they are limited by memory bandwidth, not by compute. Kernel fusion merges a producer with its consumers so the intermediate stays in registers or on-chip memory and never goes to DRAM. The rule of thumb for where fusion pays comes from the Roofline model: an operator with low arithmetic intensity (FLOPs per byte) is bandwidth-bound, and cutting its memory traffic raises its speed directly. The same idea, carried to attention, is FlashAttention (Dao et al., 2022): one kernel, tiled, with no materialised score matrix.
• Algebraic rewriting. A fixed pass order of local rewrite rules is fragile: applying rule A first may hide rule B, the phase-ordering problem. E-graphs (no page yet) store all equivalent forms at once (equality saturation) and pick the cheapest afterwards, which removes the ordering decision. Research systems use it for tensor graphs (TENSAT, Yang et al., 2021, built on the egg library); production compilers mostly still use ordered passes because saturation can blow up in size.
• Layout. Hardware prefers specific orders of dimensions: channels-last for some GPU and NPU convolutions, blocked layouts (such as NCHWc) for CPU SIMD. Tensor layout transforms (no page yet) choose a layout per operator and insert the minimum number of transposes between operators that disagree; the cost model must count a transpose as a real memory pass or the pass ping-pongs.
• Partitioning. When a graph cannot all run on one backend (an unsupported operator, or a CPU plus accelerator target), Operator graph partitioning (no page yet) cuts it into subgraphs, one per backend, and the cut edges become copies. It is a graph-cut problem with the transfer cost on the edges: cut where tensors are small.

The contract between this stage and the next is a graph of fused kernels, each described as a tensor expression (what to compute, no loops yet) with its input and output tensors and layouts fixed.

Memory planner

Intermediate tensors are the dominant memory cost of inference. Tensor memory planning (no page yet) takes the liveness intervals from the dataflow analysis and the sizes from shape inference, and assigns each tensor an offset in one pre-allocated arena so that tensors alive at the same time never overlap while tensors with disjoint lifetimes reuse the same bytes. Finding the optimal packing is NP-hard in general; compilers use greedy heuristics (largest first, or in execution order) that are close enough. Planning needs a fixed execution order and static sizes, which is why it comes after fusion (fewer, larger tensors to plan) and why dynamic shapes make it harder. The runtime then does zero allocation on the hot path.

Loop and tile layer

A fused operator becomes a loop nest, and Loop optimizations (no page yet) (tiling, interchange, unrolling, vectorisation, parallelisation) decide how fast it runs. The most important shape is the blocked matrix multiply: Tiled matrix multiplication keeps a tile of each operand in fast memory and reuses it, which raises arithmetic intensity by the tile size.

The target decides what the inner loop looks like. On a CPU the inner tile maps to vector registers (SIMD programming); on a GPU the tile maps to a thread block, shared memory and, for matmul, tensor-core instructions (CUDA and GPU languages). Systems differ in who writes the tile: Halide and TVM take a compute definition plus a schedule; Triton lets the user write the program at the level of one tile and has the compiler handle the memory coalescing, shared-memory use and thread mapping inside the tile (Tillet et al., 2019). That is a deliberate shift of the split point toward the programmer.

Schedule search

Which tile size, loop order and unroll factor is fastest is not predictable by hand: it depends on cache sizes, register count and vector width. Kernel autotuning (no page yet) turns the choice into a search. The system generates candidate schedules, builds them, runs them on the device, and keeps the fastest; AutoTVM (Chen et al., 2018) added a learned cost model so that only the most promising candidates are measured, and Ansor (Zheng et al., 2020) generates the search space itself instead of requiring a template written by hand. Triton exposes the same idea at the user level with an autotune decorator over a list of configurations.

The Roofline model is the sanity check around this loop: it gives the upper bound a kernel can reach (the lower of the compute peak and bandwidth times arithmetic intensity), so the tuner knows when to stop and a reviewer knows whether a slow kernel is a tuning problem or a design problem. Results are stored in a tuning database keyed by operator, shape and hardware, so the cost is paid once per target, not per run.

Code generator

The loop IR goes through the same machinery as any compiler back end, which is why ML compilers build on LLVM or MLIR instead of rewriting it. Instruction selection maps IR operations to target instructions (fused multiply-add, vector loads, tensor-core operations, address modes). Register allocation assigns values to physical registers and spills the rest. For ML kernels this step has an unusual weight: on a GPU, registers per thread limit how many threads can be resident at once (occupancy), so a tile size that spills or uses too many registers is slow, and the tuner sees that only as a worse measured time. The final artifact is machine code or PTX (NVIDIA's virtual assembly, finished by the driver).

Data and control flow

One model, first call (JIT path).

1. The driver traces the model with example inputs and records a graph plus guards on shapes and dtypes.
2. Shape inference runs; fusion, layout and partitioning passes rewrite the graph into fused kernels. Reads the graph, writes a new graph; this is CPU time and is seconds to minutes for a large model (order of).
3. The memory planner assigns arena offsets from liveness and sizes.
4. For each fused kernel the compiler looks up the tuning database; on a miss it either takes a default schedule (fast compile) or starts a search (slow compile, fast kernel).
5. The loop IR is lowered to SSA, instructions are selected, registers allocated, machine code emitted and loaded.
6. The result is cached under the cache key. Later calls run guards (microseconds, order of) and launch.

One tuning search.

1. Take a fused kernel and its shape.
2. Generate N candidate schedules from a template or a generated space.
3. The cost model ranks them; the top few are compiled and measured on the real device with warmup and repeats.
4. Measured times update the cost model and the best candidate is kept; repeat until the budget is spent or the best is within a margin of the roofline bound.
5. Store (operator, shape, device) → schedule in the tuning database.

Time goes to measuring: every candidate must be compiled and run.

Scaling

Limits in the order they bite, with the lever for each:

1. Compile time. Large graphs with many distinct shapes make the first call slow. Levers: cache across processes, ahead-of-time build for known shapes, bucket shapes (pad to a few sizes), skip search for cold kernels.
2. Search cost. Candidates times kernels times device time. Levers: a learned cost model that prunes before measuring, a reusable tuning database, tuning only the hot kernels (a few matmul and attention shapes usually dominate runtime).
3. Kernel quality. After fusion and tuning, the remaining gap is usually the big matmuls. Lever: call a vendor library for those and let the compiler fuse the surrounding operators, or use tensor-core paths.
4. Memory. The arena size is set by the peak of live tensors. Levers: better planning order, recomputation instead of storing, smaller dtypes (Weight quantization).

Failure modes

Real implementations

Sources

• T. Chen et al., TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, OSDI 2018. paper
• T. Chen et al., Learning to Optimize Tensor Programs (AutoTVM), 2018. paper
• L. Zheng et al., Ansor: Generating High-Performance Tensor Programs for Deep Learning, 2020. paper
• C. Lattner et al., MLIR: A Compiler Infrastructure for the End of Moore's Law, 2020. paper
• P. Tillet, H.-T. Kung, D. Cox, Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, MAPL 2019 (ACM; no free link verified).
• J. Ragan-Kelley et al., Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines, PLDI 2013. paper
• Y. Yang et al., Equality Saturation for Tensor Graph Superoptimization (TENSAT), 2021. paper
• M. Willsey et al., egg: Fast and Extensible Equality Saturation, POPL 2021. paper
• S. Williams, A. Waterman, D. Patterson, Roofline: An Insightful Visual Performance Model for Multicore Architectures, CACM 2009; tech report EECS-2008-134.
• T. Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022. paper
• J. Ansel et al., PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation, ASPLOS 2024. paper
• Official docs: OpenXLA, TVM, Triton, MLIR.