From Pending IR to a Device Buffer: xla_graph_executor.cpp, Line by Line

On an XLA device your operations do not run when you write them. They pile up as unexecuted nodes, and then one call turns the pile into a single number. This is the 1,609-line C++ file that does it, read end to end. The compilation cache and the two environment variables that size it. The arena that knows every live tensor, and why a random seed has to be a small graph of its own. The six merge sites, scattered across nine hundred lines, that assemble the graph hash: a config flag a print sets differently from a sync, the two git revisions baked in at build time, one hash per synced tensor, the parameter order, the donated buffers, and the sharding mode. The lookup that decides compile or replay, and the two counters that tell you which happened. Buffer donation and the one condition under which it is safe, argued in the file’s own twenty-line counterexample. Then lowering, sharding annotations, parameter wrapping above 3200 inputs, and the single call that costs the seconds. The signature exhibit folds the real hash term by term, in the order the code folds it, so you can watch where one changed ingredient makes every later digest diverge.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of From Pending IR to a Device Buffer: xla_graph_executor.cpp, Line by Line — built to be understood, not skimmed.

Questions

What exactly goes into the torch_xla graph hash?
Six merge sites, in this order. First the hash is seeded from the boolean config.force_ltc_data (xla_graph_executor.cpp:646), which is true on a real sync and false on a value read. Then three values at once: whatever the computation client reports as its compilation environment, plus TORCH_GITREV and XLA_GITREV, the two git revisions stamped into the binary at build time (:652). Then, for every tensor that still needs syncing, that tensor’s root IR hash, which already covers the operation, the shape, and the hashes of its operands (:696). Then the parameter sequence produced by the post-order traversal, so two graphs with the same operations but a different input order stay distinct (:1562). Then the sorted list of donated buffer indices, but only if that list is non-empty (:1568). And finally the auto-sharding flag together with the contents of XLA_AUTO_SPMD_MESH (:1572). Because a hash fold is one-way, a change at any of these merges changes every digest after it, and the finished value is the entire cache key.
Why does my torch_xla program recompile every step?
Because some ingredient of the graph hash is changing between steps, and the cache is keyed on nothing else. The usual culprit is shape: a batch that is not a fixed size, a sequence padded to its actual length rather than a bucket, or an index tensor whose length depends on the data. Each distinct shape produces different IR hashes and therefore a different key. Two less obvious ones live in the same list. Reading a value (a print, a .cpu(), an .item()) syncs with force_ltc_data set to false, and that flag is the first thing merged into the hash, so the graph you compiled to answer the print is not the graph a later sync compiles. And a graph whose parameter order varies run to run hashes differently even when the operations are identical. Read UncachedCompile and CachedCompile in the metrics report: in a warm loop UncachedCompile stops climbing after the first few steps.
Why does printing a tensor and then syncing compile the same graph twice?
Because they are not the same graph as far as the cache is concerned. Printing goes through GetTensors (xla_graph_executor.cpp:493), which sets config.force_ltc_data = false at line 498. That is correct for a read: you want the value, and you do not want the tensor’s pending IR silently replaced by a device buffer as a side effect of looking at it. But line 646 seeds the graph hash from that exact flag, so the compiled program is filed under a key carrying force_ltc_data=false. When you later call torch_xla.sync() over the same operations, force_ltc_data is true, the key is different, the lookup misses, and the identical program is compiled a second time. Dumping the HLO instead is free: DumpHloComputation at line 391 collects the IR values and renders them, with no config, no hash, no cache lookup and no call into the runtime.
What is XLA_COMPILATION_CACHE_SIZE and when should I change it?
It is the maximum number of compiled programs kept in memory, read once at line 86 of xla_graph_executor.cpp with a default of 2048, and it sizes a plain LRU cache that evicts the least recently used entry on overflow. For an ordinary training loop with a handful of distinct graphs it never matters. It starts mattering when a program legitimately has more distinct graphs than that: many bucketed sequence lengths, many branches, or a long evaluation harness that runs dozens of different models. The failure it prevents is specific and unpleasant on the torch.compile path, where ExecuteComputationWithBarrier looks a program up by a hash computed earlier and hard-fails if the entry was evicted, rather than recompiling. Two related variables sit beside it: XLA_PERSISTENT_CACHE_PATH turns the cache into a directory on disk that survives the process, and XLA_PERSISTENT_CACHE_READ_ONLY stops it being written to.
When does torch_xla donate an input buffer, and why not always?
Buffer donation tells the compiler an input buffer is dead after the program runs, so an output can be written on top of it instead of into a fresh allocation. torch_xla will only claim that at a step barrier: GetBufferDonors computes donate_ltc_data as sync_ltc_data && force_ltc_data (line 1335), and only then collects the input buffers whose tensors are also outputs of this very execution. The reason is written out as a counterexample in the source. Take a tensor A, a graph that adds 0.4 to it, and two prints. If the first print aliases A’s buffer, the second print reads a buffer that has already been advanced once and advances it again, returning a wrong number, and views make it worse because a value read at two different times has to reflect changes to its source. Two switches sit around this: XLA_ENABLE_PARAM_ALIASING (default true) turns the whole feature off, and auto-sharding disables it because aliasing compares unpartitioned shapes. The donated indices reach the runtime as annotations on the compiled program, not as execution options.

Related explainers