A tour of XLA, where MLIR lives

If you train or serve models, XLA is probably compiling them, whether or not you have ever named it. It is the compiler behind JAX, behind TensorFlow, and, through PyTorch/XLA, behind much of PyTorch on TPUs. It takes the tensor program your framework produces and turns it into fast code for CPUs, GPUs, and TPUs.

This is a tour of how it is built: the IR it works in, the optimizer at its center, how it generates code for three very different kinds of hardware, how it spreads a program across many chips, and where you can step outside it. The aim is to be able to read what XLA is actually doing to your model, and to see the shape of its design, where it is strong and where its assumptions start to bind.

The IR: HLO

XLA's intermediate representation is HLO, "High Level Operations", an SSA graph of operations over statically-shaped tensors. The operation set is deliberately small, on the order of 150 opcodes, covering the primitives a tensor program reduces to: dot (general contraction), convolution, reduce, dynamic-slice, the collectives (all-reduce, all-gather), and an escape valve, custom-call. A fragment of optimized HLO looks like this, with the f32[4,16]{1,0} shape-and-layout notation and the fusion op:

%fused_computation.1 (param_0: f32[4,16], param_1: f32[16]) -> f32[4,16] {
  %b   = f32[4,16]{1,0} broadcast(%param_1), dimensions={1}
  %add = f32[4,16]{1,0} add(%param_0, %b)
  ROOT %max = f32[4,16]{1,0} maximum(%add, %constant)
}
ENTRY %main {
  %f = f32[4,16]{1,0} fusion(%x, %W), kind=kCustom, calls=%fused_computation
}

HLO predates MLIR but these days frameworks (like JAX) do not hand XLA this C++ HLO directly. They emit StableHLO, a portable, versioned interchange that XLA imports and translates into HLO. StableHLO is the well designed contract between a framework and the compiler, the thing JAX's lower() produces.

The optimizer at the center

XLA's optimizer is a sequence of passes, each one a local rewrite of the graph. The passes cluster into families: simplify (algebraic simplifications, common subexpression elimination), layout (assign a memory layout to every tensor), fuse (merge producer-consumer chains into fusion nodes), allocate (assign buffer lifetimes and reuse), and schedule (order the operations for the target).

The fusion pass is the one that matters most for inference. A fusion node merges a producer-consumer chain into a single kernel, so the intermediate tensors never touch global memory. This is how a relu(x @ W + b) becomes one kernel instead of three.

Generating code for three very different hardwares

XLA's backends diverge sharply:

• CPU: LLVM emits an object file. The code is threaded and vectorized for the host.
• GPU: LLVM/PTX for custom kernels, cuBLAS/cuDNN for the heavy matmuls and convolutions, and Triton for the fused kernels. The GPU backend is where XLA and the inference engine world meet.
• TPU: Mosaic emits code for libtpu. The TPU backend is entirely different because the hardware is entirely different: no warps, a different memory model.

Spreading a program across many chips

XLA's SPMD partitioning handles the multi-chip case. Given shardings on some values, it infers a consistent sharding for every other value and inserts the collectives needed to make the sharded program equal the original. This is the machinery behind JAX's sharding, and it is the same machinery that serves a model across multiple GPUs.

Where you can step outside it

The escape valve is custom-call: an op that lets you hand XLA a function pointer to your own kernel. This is how frameworks drop in hand-written kernels for operations XLA does not fuse well. It is also the door to Pallas, which lets you write the kernel yourself and lower it through Triton on GPU or Mosaic on TPU.

XLA fuses at the granularity of HLO ops using a cost model you do not control. When that granularity is wrong, you stop describing what to compute and start describing how.

Back to the blog