ML Compiler Architecture and Optimization

Run a small PyTorch function containing matrix multiplication, bias addition, and GELU without compilation. Nsight Systems can show three or more GPU launches and two temporary tensor transfers. Run the same function through torch.compile and the first invocation pauses while software captures and compiles the program. Later invocations may show one library matrix call with a fused epilogue or a smaller number of generated kernels. A new input shape can cause another pause, while a Python print or data-dependent branch can split the compiled region.

Chapter 12 explains every stage that produced those observations. Program capture records tensor operations and the assumptions that make the recorded program valid. A guard is a runtime condition, such as a tensor having a particular datatype or a dimension remaining within a supported range. A graph break is a point where capture stops and ordinary eager execution resumes. Compiler logs report guards, recompilations, and graph breaks; those artifacts explain many timeline gaps before any assembly inspection is necessary.

A captured tensor program passes through intermediate representations, usually abbreviated IRs. An IR is a compiler-owned program form printed in logs or dump files; the form is neither the original Python source nor final machine instructions. A high-level IR can still say 'matrix multiplication,' while a later GPU IR names loops, memory spaces, thread blocks, and transfers. Reading successive dumps reveals where an operation was fused, tiled, converted to a library call, or expanded into lower-level instructions.

The running example begins on a central processing unit (CPU), then launches generated work on a graphics processing unit (GPU). A high-level matrix multiplication (GEMM) becomes tiled GPU work that reads operands from high-bandwidth memory (HBM). Compiler and profiler reports count floating-point operations (FLOPs) as arithmetic work and bytes as data movement. The intermediate representations (IRs) explain how the source expression became those measured operations.

PyTorch compilation, Multi-Level Intermediate Representation (MLIR), and Accelerated Linear Algebra (XLA) use different names and pipelines, but the diagnostic method remains consistent. Inspect the first stage whose output differs from the expected program, then connect that difference to a launch count, memory transaction count, generated instruction, compile pause, or numerical result. The chapter builds this method with one concrete compiled expression rather than treating a compiler as an unexplained black box.