Pipeline

XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s optimizer is not one monolithic algorithm. It is roughly two hundred separate passes, each one a class implementing HloPassInterface and run inside HloPassPipelines (xla/hlo/pass/hlo_pass_pipeline.h), one after another. Pipelines can contain other pipelines, so the sequence you would see if you unrolled the whole thing is deep, hardware-aware, and different on CPU than on GPU or TPU.

None of that has to stay invisible. XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → can be told to write the module out after every single pass runs, and the resulting files are not a black box: they are the compiler explaining itself, one decision at a time, in the exact order it made them.

Two flags do the whole job. --xla_dump_to points at a directory; --xla_dump_hlo_pass_re=.* tells XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → to match every pass name against that regex and dump after each match, which for .* means every pass, no exceptions. Set both through XLA_FLAGS before jax is imported, compile anything, and the directory fills with one file per pass step.

This ran on this machine against the same small attention program from earlier chapters, and it produced forty-two numbered files. Forty-two decisions, on a program with three matmuls and a softmax, is already a useful thing to internalize: production models with far more operations produce dumps that are correspondingly larger, but the same story-reading skill applies at any size.

import os
os.environ["XLA_FLAGS"] = "--xla_dump_to=/tmp/xla-dump --xla_dump_hlo_pass_re=.*"

import jax
import jax.numpy as jnp

def attend(q, k, v):
    s = q @ k.T / jnp.sqrt(jnp.float32(q.shape[-1]))
    return jax.nn.softmax(s) @ v

x = jnp.ones((64, 64))
jax.jit(attend).lower(x, x, x).compile()   # 42 numbered dump files appear

The numbered filenames are not arbitrary; read them in order and they tell the actual story of the compile. Early files cluster under a name like simplification, and several of them repeat: algsimp runs, something changes, algsimp runs again, because that stage is wrapped to iterate until two consecutive runs report no change at all. That repetition is a fixed point, not a bug in the dump.

Partway through, the filenames shift to a different family: HLO_passes_through_layout_assignment. That name change marks a real boundary, not just a naming convention. Before it, shapes carry no physical layout at all; after layout-assignment itself runs, every shape in the module has one, and passes on the far side of that line are working with information the ones before it never had.

Near the very end, copy-insertion appears, adding copies wherever the analyses run earlier in the pipeline found interference they could not resolve any other way. Its position in the sequence is not an accident either: it exists precisely because it needs a finished layout to know where interference actually lives.

0003 simplification.after_pipeline-start.before_algsimp
0004 simplification.after_algsimp.before_simplify-sorts
0005 simplification.after_tree_reduction_rewriter.before_zero_sized_hlo_elimination
0010 HLO_passes_through_layout_assignment.after_transpose-folding.before_cse
0011 HLO_passes_through_layout_assignment.after_cse.before_cse_barrier_expander
0012 HLO_passes_through_layout_assignment.after_flatten-call-graph.before_layout-assignment
0013 HLO_passes_through_layout_assignment.after_layout-assignment.before_sub-byte-size-setter
0017 HLO_passes_after_layout_assignment.after_fusion.before_simplification_after_layout_assignment
0019 copy-insertion.after_adding_copies_to_resolve_interference
0022 HLO_passes_after_layout_assignment.after_copy-insertion.before_dce

If this chapter has one fact worth carrying forward, it is that layout assignment is the irreversible line in the whole pipeline. Every pass before it can reason about a module using shape alone; every pass after it has to reason using shape and physical layout together, and nothing later ever goes back to being layout-free. Chapter 6 spends its entire time on what a layout actually is and what layout assignment does to earn that boundary.

The habit worth building here is simpler than any one fact: when a decision in an optimized module looks mysterious, the dump usually already explains it. The file where an instruction first appears, set against the file just before it where the instruction did not yet exist, is the compiler showing its own reasoning, and reading that pair of files is very often faster than guessing.

The compiler will show you every decision, if you ask for the dump.

The driver above ran on both a CPU and a Colab TPU v6e, and the dumps disagree about almost everything. CPU produced 42 files for this program; the TPU produced 88. That difference is not a deeper pipeline in any interesting sense, it is a different pipeline, built for different hardware, as the section on hardware-aware passes claimed and the file names now prove.

The names are where it gets specific. The TPU dump carries stages a CPU compile has no reason to own: hlo_device_type_async_wrapper, Before_X64_rewriter, X64_elimination, Phase_1_pre_layout_assignment_passes, and targets like add-random-host-offloading and tpu-embedding-thread-annotator. Read that list next to chapter 9 and the shape of the backend shows through: host offloading, embedding hardware, and a 64-bit rewrite the TPU wants handled before layout assignment ever runs.

One practical detail costs an afternoon if you meet it cold. TPU dump files insert a build id between the module name and the step number, as in module_0007.jit_attend.cl_948136882.0000.hlo_device_type_async_wrapper, while CPU files go straight from the module to the step. A pattern written against one backend silently matches nothing on the other, and a diff of nothing against something looks exactly like a real answer. Print the filenames before you parse them.