What does --xla_dump_to do?

It names a directory, and XLA writes debugging artifacts for every module it compiles into that directory. You normally set it through the XLA_FLAGS environment variable, which works the same way for JAX, TensorFlow and PyTorch/XLA: XLA_FLAGS="--xla_dump_to=/tmp/dump" ./your_program. Setting it alone is enough, because the dump options are built by a constructor that fills in what you did not say. If you named no output format, HLO-as-text is turned on for you. If you named a format but no directory, the dump goes to stdout instead. If you named neither a directory nor a format, nothing is dumped at all. What you get by default is the module as text at two points in the compile, plus the buffer assignment reports at the second point, plus a file listing every debug option whose value is not the default.

What do the files in an XLA dump directory mean?

Every filename is built by the same function, and it has the shape [prefix.]module_NNNN[.module_name].suffix. The four-digit number is the module unique id, not a counter of dumped files, so gaps in the numbering are normal. The module name is whatever the framework called the computation, which is why JAX dumps say jit_step or pmap_something. The suffix is the interesting part: before_optimizations and after_optimizations are the two named bookends of a compile; a suffix like 0007.simplification.after_algsimp.before_reshape-mover is a per-pass dump, where the leading number is a per-module step counter and the two pass names say which pass just finished and which one is about to run. The prefix slot holds a timestamp, but only if you asked for one with --xla_dump_include_timestamp. Trailing suffixes name the format: .txt, .hlo.pb, .dot, .html, .riegeli.

How do I dump the HLO between compiler passes?

Add --xla_dump_hlo_pass_re with a regular expression matching the pass names you care about. Per-pass dumping is off unless you ask for it: the flag is the only thing that turns it on, and without it the dump holds the two bookends only. --xla_dump_hlo_pass_re=spmd|propagation gives you the SPMD partitioner and sharding propagation; --xla_dump_hlo_pass_re=.* gives you every pass and a very large directory. There is a rule worth knowing about that last one. When the regular expression is exactly .*, the pipeline runner only writes a file for a pass that reported it changed the module, so a pass that ran and did nothing leaves no trace. Any other regular expression dumps every matching pass whether it changed anything or not. A companion flag, --xla_dump_hlo_pipeline_re, restricts the dump to named sub-pipelines instead of individual passes.

How do I run a dumped HLO file without JAX or PyTorch?

Two tools take a dumped file directly. run_hlo_module compiles and executes the module on a platform you name, and by default also runs it on the reference interpreter and compares the two results, which is how you turn "the model gives wrong numbers" into a self-contained reproduction: run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo. It operates on pre-optimization HLO, so the before_optimizations file is the one to hand it, and it accepts several input formats including HLO text, HloProto in binary or text form, and StableHLO. hlo-opt is the other one. It does not execute anything; it runs the compiler and prints a stage. --stage=hlo gives you HLO after all optimizations, and other stages depend on the platform: buffer-assignment, hlo-backend, llvm, ptx, html. It also runs individual passes by name with --passes=dce,algsimp, which is the fastest way to find out which pass made a change you saw in a dump.

Why is my XLA dump directory enormous, and how do I make it smaller?

Usually because --xla_dump_hlo_pass_re=.* is on, and a real GPU pipeline runs on the order of two hundred passes for every module compiled. Multiply that by every module a program compiles and by every output format you enabled, and the file count grows fast. Four flags cut it down. --xla_dump_hlo_module_re limits the dump to modules whose name matches a pattern, which is the right first move when you know which computation you are chasing. --xla_dump_hlo_pipeline_re limits per-pass dumping to named sub-pipelines. --xla_dump_max_hlo_modules puts a ceiling on how many distinct modules get dumped in one process, and note that when the ceiling is hit the compiler simply logs an error and stops writing, so a truncated dump can look like a missing one. --xla_dump_hlo_to_subfolder puts each module in its own subdirectory named after the module and the host and process id, which is what you want under multi-host runs where every process writes into the same place.

Reading an XLA Dump: What the Compiler Writes, and How to Read It

Someone asks you to attach the HLO dump. You set one flag, point a directory at it, and get back a pile of files with names like module_0000.jit_step.0007.simplification.after_algsimp.before_reshape-mover.txt. This reads that directory end to end at a pinned commit of openxla/xla: the four inference rules one flag triggers before anything is written, the filename grammar and what every segment of it means, the two named bookends of a compile and the four buffer-assignment reports that ride with the second one, the per-pass files and the rule that quietly drops any pass that changed nothing, and the flags that make the dump larger, smaller, or unreadable. Then the two tools that take a dumped file and do something with it without a framework anywhere in sight: hlo-opt, which re-runs named passes and prints any stage of the compile, and run_hlo_module, which executes the module and checks it against the interpreter. The signature exhibit assembles the actual directory listing from the real inference rules, one flag at a time.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of Reading an XLA Dump: What the Compiler Writes, and How to Read It — built to be understood, not skimmed.

Questions

What does --xla_dump_to do?: It names a directory, and XLA writes debugging artifacts for every module it compiles into that directory. You normally set it through the XLA_FLAGS environment variable, which works the same way for JAX, TensorFlow and PyTorch/XLA: XLA_FLAGS="--xla_dump_to=/tmp/dump" ./your_program. Setting it alone is enough, because the dump options are built by a constructor that fills in what you did not say. If you named no output format, HLO-as-text is turned on for you. If you named a format but no directory, the dump goes to stdout instead. If you named neither a directory nor a format, nothing is dumped at all. What you get by default is the module as text at two points in the compile, plus the buffer assignment reports at the second point, plus a file listing every debug option whose value is not the default.
What do the files in an XLA dump directory mean?: Every filename is built by the same function, and it has the shape [prefix.]module_NNNN[.module_name].suffix. The four-digit number is the module unique id, not a counter of dumped files, so gaps in the numbering are normal. The module name is whatever the framework called the computation, which is why JAX dumps say jit_step or pmap_something. The suffix is the interesting part: before_optimizations and after_optimizations are the two named bookends of a compile; a suffix like 0007.simplification.after_algsimp.before_reshape-mover is a per-pass dump, where the leading number is a per-module step counter and the two pass names say which pass just finished and which one is about to run. The prefix slot holds a timestamp, but only if you asked for one with --xla_dump_include_timestamp. Trailing suffixes name the format: .txt, .hlo.pb, .dot, .html, .riegeli.
How do I dump the HLO between compiler passes?: Add --xla_dump_hlo_pass_re with a regular expression matching the pass names you care about. Per-pass dumping is off unless you ask for it: the flag is the only thing that turns it on, and without it the dump holds the two bookends only. --xla_dump_hlo_pass_re=spmd|propagation gives you the SPMD partitioner and sharding propagation; --xla_dump_hlo_pass_re=.* gives you every pass and a very large directory. There is a rule worth knowing about that last one. When the regular expression is exactly .*, the pipeline runner only writes a file for a pass that reported it changed the module, so a pass that ran and did nothing leaves no trace. Any other regular expression dumps every matching pass whether it changed anything or not. A companion flag, --xla_dump_hlo_pipeline_re, restricts the dump to named sub-pipelines instead of individual passes.
How do I run a dumped HLO file without JAX or PyTorch?: Two tools take a dumped file directly. run_hlo_module compiles and executes the module on a platform you name, and by default also runs it on the reference interpreter and compares the two results, which is how you turn "the model gives wrong numbers" into a self-contained reproduction: run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo. It operates on pre-optimization HLO, so the before_optimizations file is the one to hand it, and it accepts several input formats including HLO text, HloProto in binary or text form, and StableHLO. hlo-opt is the other one. It does not execute anything; it runs the compiler and prints a stage. --stage=hlo gives you HLO after all optimizations, and other stages depend on the platform: buffer-assignment, hlo-backend, llvm, ptx, html. It also runs individual passes by name with --passes=dce,algsimp, which is the fastest way to find out which pass made a change you saw in a dump.
Why is my XLA dump directory enormous, and how do I make it smaller?: Usually because --xla_dump_hlo_pass_re=.* is on, and a real GPU pipeline runs on the order of two hundred passes for every module compiled. Multiply that by every module a program compiles and by every output format you enabled, and the file count grows fast. Four flags cut it down. --xla_dump_hlo_module_re limits the dump to modules whose name matches a pattern, which is the right first move when you know which computation you are chasing. --xla_dump_hlo_pipeline_re limits per-pass dumping to named sub-pipelines. --xla_dump_max_hlo_modules puts a ceiling on how many distinct modules get dumped in one process, and note that when the ceiling is hit the compiler simply logs an error and stops writing, so a truncated dump can look like a missing one. --xla_dump_hlo_to_subfolder puts each module in its own subdirectory named after the module and the host and process id, which is what you want under multi-host runs where every process writes into the same place.

Reading an XLA Dump: What the Compiler Writes, and How to Read It

Questions

Related explainers