XLA: a review and a critique

Here is a small function and the entire program XLA compiled it into:

def simp(x):
    y = (x + 0.0) * 1.0
    return jnp.transpose(jnp.transpose(y))

ENTRY %main (x: f32[3,3]) -> f32[3,3] {
  %x = f32[3,3]{1,0} parameter(0)
  ROOT %copy = f32[3,3]{1,0} copy(%x)
}

The add-zero, the multiply-one, and both transposes are gone. XLA proved the whole thing is the identity and emitted a single copy. This is not a toy result; it is the same algebraic-simplification machinery that lets XLA take a page of dense linear algebra and hand back a fused kernel that runs at the memory-bandwidth limit of a TPU. XLA is a genuine optimizing compiler, and a very good one.

It is also rigid in ways that will bite you, and the interesting thing, the thing worth a long post, is that its strengths and its limitations are the same design decision. XLA freezes your shapes, sees the whole graph at once, allocates every buffer statically, and fuses against a global cost model you do not control. That is exactly why it is fast, and exactly why it is inflexible.

Where XLA sits, and what it eats

XLA is an ahead-of-time compiler for array programs. "Ahead of time" is relative to execution, not to your process: it compiles the first time a given shape signature appears, caches the result, and reuses it. Its input is a graph of high-level tensor operations, and the modern portable front door to that graph is StableHLO, an MLIR dialect that JAX emits and that XLA ingests. The full path:

jaxpr → StableHLO → HLO → [ optimization passes ] → optimized HLO → backend
                                                              ├─ LLVM  (CPU)
                                                              ├─ LLVM/PTX + cuBLAS/cuDNN/Triton (GPU)
                                                              └─ XLA:TPU (closed backend)

The naming is a genuine mess worth untangling once. There is the classic HLO, a hand-written C++ IR that predates MLIR. There is the MLIR world layered on later: MHLO, and then StableHLO, which is built on MHLO and adds serialization plus compatibility guarantees, so it can serve as a stable portability contract between frameworks and compilers. JAX lowers to StableHLO; XLA converts StableHLO into its internal HLO and optimizes that.

The bargain, and where it stops paying off

The bargain is: freeze shapes, see the whole graph, allocate statically, fuse against a global cost model. In exchange you get an optimizing compiler that can prove x + 0.0 is x and emit a single fused kernel running at bandwidth.

Where it stops paying off:

• Dynamic shapes cost you. The ahead-of-time model wants shapes frozen. Every dynamic dimension either forces recompilation or a fallback path.
• You do not control the fusion. The cost model decides what fuses and what does not. When it is wrong, the escape hatch is custom-call or Pallas.
• The TPU backend is closed. The optimizations are mostly not open source; only the machine model is readable.

XLA freezes your shapes, sees the whole graph at once, allocates every buffer statically, and fuses against a cost model you do not control. That is exactly why it is fast, and exactly why it is inflexible.

Back to the blog