The pass pipeline, from a real dump

A pass in XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → is a class with two members that matter. name() returns the string the dump will print, and Run() does the work and reports whether it changed anything. The return type is what makes the rest of the pipeline's behaviour legible: absl::StatusOr<bool>, where the comment in the header says the bool means whether it modified the module.

That bool is the reason a dump has repeats in it. A pipeline can be wrapped so that it reruns its member passes until every one of them returns false, and the chapter above notices this when algsimp shows up several times in a row. The mechanism behind the observation is a loop reading return values, and there is a companion entry point for it, RunOnChangedComputations, which lets a pass in a fixpoint loop revisit only what the previous iteration touched instead of the whole module.

Most passes derive from HloModulePass rather than implementing the interface directly. The distinction is about scope: a module pass gets the whole HloModule and may rewrite across computation boundaries, which is what layout assignment and buffer assignment need. Knowing which base a pass uses tells you, before reading a line of its body, whether it could possibly have caused a change you found two computations away.

class HloPassInterface {
 public:
  virtual absl::string_view name() const = 0;

  // Run the pass on the given HLO module ...
  // Returns whether it modified the module.
  absl::StatusOr<bool> Run(
      HloModule* module,
      const absl::flat_hash_set<absl::string_view>& execution_threads = {});

  // Run the pass on computations changed from last iteration in given HLO
  // module.
  // ...
};

class HloModulePass : public HloPassInterface {
  // ...
};

The dump flags are defined in xla/debug_options_flags.cc with help strings, and the help strings answer questions the tools page does not. --xla_dump_hlo_pass_re dumps HLO before and after optimization passes which match this regular expression, in addition to dumping at the very beginning and end of compilation. Two facts follow from that sentence: the endpoints are always dumped whether or not your regex matches anything, and the per-pass dumps come in before-and-after pairs rather than snapshots.

--xla_dump_to takes a directory and has two special values. The literal - means stdout, and sponge or test_undeclared_outputs_dir route into the directory named by the TEST_UNDECLARED_OUTPUTS_DIR environment variable. If you point it at a path that does not exist you get nothing and no complaint, which is the most common reason a first dump attempt appears to do nothing at all.

The disabling flags carry a warning in their own help text that the path had to learn the hard way. --xla_disable_hlo_passes takes a comma-separated list, and the help says these names must exactly match the passes' names, with no whitespace around commas. Nothing in that contract promises the compiler will tell you when a name matched nothing. The retraction on /xla/fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → is exactly that failure, published and then corrected.

A dump filename encodes four things, and once you can read all four you no longer need to open most of the files. The parts are a zero-padded module number, the module name, a zero-padded step number, and then the interesting half: which pipeline the step belongs to, which pass just finished, and which pass is about to run. HLO_passes_through_layout_assignment.after_cse.before_cse_barrier_expander is a complete sentence about where the compiler was.

Nesting is recoverable from that alone, without reading any source. In this repo's CPU capture, step 15 sits in pipeline HLO_passes_after_layout_assignment with before_after_layout_assignment, and step 16 sits in pipeline after_layout_assignment going from pipeline-start to pipeline-end. A name that appears as a pass in one line and as the pipeline in the next is a nested pipeline being entered. Do that for a whole dump and you have drawn the tree.

The other thing the filenames give you free is the fixpoint. Steps 3 through 8 of that capture are simplification running algsimp, then simplify-sorts, then algsimp again from pipeline-start, twice more. Three entries into the same pipeline from its start is three iterations of the loop the previous section described, and the loop stopped when the third pass produced no change.

A nested pipeline announces itself twice: once as a pass name, then as the pipeline on the next line.

One backend-specific trap costs an afternoon if you meet it cold. TPU dumps 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 dumps go straight from module name to step. A filename pattern written against one backend matches nothing on the other, and a diff of nothing against something looks exactly like a real result. List the filenames before you parse them.

  n  pipeline                                after            -> before
  0  sharding-removal                        pipeline-start      sharding-remover
  1  SubbytePacker_pipeline                  pipeline-start      sub-byte-size-setter
  2  HLO_passes_through_layout_assignment    pipeline-start      gather_scatter_normalizer
  3  simplification                          pipeline-start      algsimp
  4  simplification                          algsimp             simplify-sorts
  5  simplification                          tree_reduction_rewriter  zero_sized_hlo_elimination
  6  simplification                          pipeline-start      algsimp
  8  simplification                          pipeline-start      algsimp
 12  HLO_passes_through_layout_assignment    flatten-call-graph  layout-assignment
 13  HLO_passes_through_layout_assignment    layout-assignment   sub-byte-size-setter
 15  HLO_passes_after_layout_assignment      pipeline-start      after_layout_assignment
 16  after_layout_assignment                 pipeline-start      pipeline-end
 17  HLO_passes_after_layout_assignment      fusion              simplification_after_layout_assignment
 22  HLO_passes_after_layout_assignment      copy-insertion      dce

  n  pipeline                                after                 before
  0  sharding-removal                        pipeline-start        sharding-remover
  1  SubbytePacker_pipeline                  pipeline-start        sub-byte-size-setter
  2  HLO_passes_through_layout_assignment    pipeline-start        gather_scatter_normalizer
  3  simplification                          pipeline-start        algsimp
  4  simplification                          algsimp               simplify-sorts
  5  simplification                          tree_reduction_rewriter  zero_sized_hlo_elimination
  6  simplification                          pipeline-start        algsimp
  7  simplification                          algsimp               simplify-sorts
  8  simplification                          pipeline-start        algsimp
  9  HLO_passes_through_layout_assignment    simplification        bitcast_dtypes_expander
 10  HLO_passes_through_layout_assignment    transpose-folding     cse
 11  HLO_passes_through_layout_assignment    cse                   cse_barrier_expander
 12  HLO_passes_through_layout_assignment    flatten-call-graph    layout-assignment
 13  HLO_passes_through_layout_assignment    layout-assignment     sub-byte-size-setter
 14  hlo_normalization                       pipeline-start        reshape-decomposer
 15  HLO_passes_after_layout_assignment      pipeline-start        after_layout_assignment
 16  after_layout_assignment                 pipeline-start        pipeline-end
 17  HLO_passes_after_layout_assignment      fusion                simplification_after_layout_assignment
 18  simplification_after_layout_assignment  pipeline-start        algsimp
 22  HLO_passes_after_layout_assignment      copy-insertion        dce

The chapter above ships an instrument, EX·15, that steps twenty of these files so you can scrub the sequence, watch the reruns appear, and find the watershed at layout assignment. Stepping is the first half of the skill. The second half is a loop you run against your own dump, and it is four steps, in this order, because each one prevents a mistake the next one would otherwise hide.

List the filenames first and read the pattern off them, rather than assuming the pattern from another backend. Then grep the pass vocabulary before you name a pass in a flag, because names are backend-specific and a name lifted from a CPU dump can be simply absent from a TPU pipeline. Then diff the pair around a pass and count instructions rather than lines, since a metadata-only reprint can produce a hundred changed lines and zero changed program. Only then, if you disabled something, time it.

That third step is where this repo has receipts. Its pass corpus records that both the CPU and TPU pipelines register a pass called exactly fusion, and that disabling fusion on the TPU leaves an attention module byte-identical, while tpu_fusion and tpu_multi_output_fusion change it. So a flag can be accepted, name a pass that genuinely exists, and still be a no-op for your program. Diffing the compiled module is the only step that catches that, and it costs ten seconds.

Read the same corpus the other direction and it doubles as a vocabulary map. Names like add-random-host-offloading, tpu-embedding-thread-annotator, and legalize-scheduling-annotations appear only on TPU; tree_reduction_rewriter, cse_barrier_expander, and sharding-remover only on CPU. Two names, fusion and bitcast_dtypes_expander, appear in both. That overlap is small enough to be worth internalizing: almost nothing you learn about one backend's pipeline transfers to another by name.

The last skill in this lesson is turning a name in a dump into a file you can open. Hardware-independent passes live under xla/hlo/transforms, which has subdirectories for simplifiers, expanders, and collectives, plus a README and its own BUILD file. Backend passes live under xla/backends/<backend>/transforms. If a grep through xla/hlo/transforms comes back empty for a name your dump printed, that is usually the answer rather than a mystery: the pass belongs to a backend.

The name-to-filename mapping is by hand, in both directions, because there is no consistent convention. The printed string is whatever the pass's name() returns, and the corpus above shows both styles side by side: layout-assignment, transpose-folding, flatten-call-graph, and copy-insertion print with hyphens, while tree_reduction_rewriter, zero_sized_hlo_elimination, and gather_scatter_normalizer print with underscores. Filenames are snake_case regardless. Convert by hand and search both spellings.

Browsing that directory once is worth more than it sounds, because the file list is a catalogue of what a compiler at this level actually worries about. Host offloading gets three passes of its own (host_offloader, host_offload_legalize, host_offloading_prepare). There is a bfloat16_propagation and an operand_upcaster for precision, a while_loop_trip_count_annotator for loops, and a defuser that undoes fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla →, which exists because debugging and testing need a way back out.