XLA and StableHLO overview

Two artefacts

StableHLO is an operation set with a specification and a compatibility promise: a small, versioned, serialisable graph language that frontends emit and compilers consume. XLA is a compiler that consumes it, with its own IR, a pass pipeline, and backends. The pair is worth reading because the boundary between them is the interesting decision: what belongs in a portable operation set, and what belongs to the compiler that runs it.

The op set is deliberately small

dot_general, reduce, broadcast_in_dim, convert, while, and a handful of elementwise operations. Shapes and dtypes are part of the contract; layouts are not. Read the specification for an operation you know well, such as dot_general, and notice how much is left to the compiler: batch dimensions, contracting dimensions, precision. That residual is where the optimisation space lives.

Layout assignment and fusion

XLA’s pipeline is the useful part. Algebraic simplification, then layout assignment, which picks physical layouts and inserts copies where a consumer wants something different, then fusion, which decides how much of the graph becomes one kernel, then backend lowering. Read the architecture page with a simple graph in mind and predict where the transposes will appear before the page tells you.

What the docs do not give

A cost model you can lift, or a promise that a StableHLO graph performs the same on every backend. Portability of semantics is not portability of speed.

What to read

The architecture overview, then the specification’s sections on operations and dynamic shapes. Skip the pages for backends you do not target.

Afterwards you should be able to take a short PyTorch sequence, write the corresponding StableHLO operations by hand, and say where layout assignment would insert a transpose.

Things that break

Treating StableHLO as a performance interface. Two graphs with identical semantics but different decomposition compile to very different kernels, because fusion and layout work on the structure you handed over. The tell is a rewrite that looks neutral, such as replacing a fused pattern with two operations, and costs a copy.