This talk by Eugene Zhulenev from Google provides a comprehensive overview of the XLA (Accelerated Linear Algebra) compiler's architecture, key design principles, and internal workings. It begins by defining XLA as a compiler designed to accelerate machine learning workloads by taking a high-level graph representation and generating optimized machine code for various hardware backends like TPUs, GPUs, and CPUs.
The presentation highlights the core design decisions that have shaped XLA: a reduced instruction set (HLO) to manage complexity, a reliance on statically shaped programs for predictable performance and memory planning, and a split architecture with a shared frontend for generic optimizations and specialized backends for device-specific code generation. The talk details fundamental XLA concepts such as the array data structure, memory layouts, tiling for TPUs, and the use of tokens for managing control dependencies. It then walks through the compiler pipeline, covering crucial frontend passes like operator fusion, layout assignment, and buffer assignment, which are key to XLA's performance. Finally, it contrasts the execution models of the different backends: the host-driven, kernel-by-kernel execution on GPUs and CPUs versus the on-device, single-binary execution on TPUs. The relationship between the high-level PjRt API for frameworks and XLA's internal runtime is also clarified.
JAX/OpenXLA DevLab 2025 - XLA Architecture
0:28 What is XLA?: XLA (Accelerated Linear Algebra) is a domain-specific compiler for linear algebra that takes a high-level graph representation and generates optimized machine code. Originally designed to program TPUs for TensorFlow, its architecture has proven effective for GPUs, CPUs, and other accelerators, and it is now a key backend for frameworks like JAX and PyTorch.
2:05 Core Design Philosophy: XLA was built on several key decisions:
Reduced Instruction Set: The HLO (High-Level Operations) instruction set is intentionally small (~100 operations) to make the compiler manageable, in contrast to the thousands of custom kernels found in frameworks.
Statically Shaped Programs: Programs must have known tensor shapes at compile time. This allows for precise memory planning and performance predictability, which is critical for hardware like TPUs. This principle is now evolving to support dynamic shapes.
Static Operation Scheduling: XLA creates a static schedule for operations at compile time, enabling optimizations for minimum memory (TPU/GPU) or maximum concurrency (CPU).
Frontend/Backend Split: A shared frontend performs hardware-agnostic optimizations, while multiple backends (TPU, GPU, CPU) handle device-specific code generation.
StableHLO: The modern, MLIR-based dialect that serves as the official input language for XLA. Frameworks target XLA by emitting StableHLO modules via the PjRt API.
HLO: XLA's internal, historical Intermediate Representation (IR). XLA converts incoming StableHLO into its internal HLO format for optimization.
Array and Shape: The fundamental data structure in XLA is the array, which is defined by its data type, dimensions, and layout. The layout specifies how dimensions are physically stored in memory (e.g., row-major or column-major).
Tiling (15:18): Primarily a TPU concept where the memory layout is further divided into tiles that match the hardware's register sizes, simplifying code generation. On CPU/GPU, this is typically handled implicitly during code generation.
Tokens (21:50): A special type used to convert control dependencies into data dependencies. This forces a specific execution order for side-effecting operations like collectives, which is crucial for preventing deadlocks in distributed training.
Fusions: The most important optimization. XLA groups multiple element-wise operations into a single "fusion" kernel. This minimizes memory bandwidth by ensuring intermediate values are kept in registers instead of being written to and read back from main memory.
Layout Assignment (36:59): The compiler analyzes operations (like matrix multiplications) and assigns the optimal memory layout to all intermediate arrays to ensure efficient, linear memory access on the target hardware.
Buffer Assignment (40:18): After scheduling, XLA performs a liveness analysis to determine the peak memory usage and assigns all intermediate values to slices of a single, pre-allocated temporary "scratch" buffer.
CPU/GPU: The host CPU drives the execution. The XLA runtime launches a sequence of individual kernels (corresponding to fusions). Concurrency is managed by CUDA streams on GPU or a thread pool on CPU.
TPU (46:36): The entire optimized HLO graph is compiled into a single, large binary. This binary is loaded onto the TPU and executes from start to finish without host intervention, only interrupting the host upon completion.
PjRt (Pluggable JAX Runtime): The high-level, stable C API that frameworks use to interact with XLA. It abstracts away the complexities of device management, compilation, and execution.
XLA Runtime (Thunks): A low-level implementation detail of the CPU/GPU backends. It is a simple interpreter that executes a sequence of commands called "thunks" (e.g., launch kernel, perform collective), which orchestrate the computation on the device.