Inside the lazy tensor

The chapter above this lesson says the operations get recorded rather than run, which is the right sentence to carry around and the wrong place to stop. Open XLATensor::Data and the recording turns out to be a four-way state machine, one state per constructor. A tensor holds a device handle, meaning real bytes on the chip. Or it holds an IR value, meaning a pending expression nothing has computed. Or it holds an at::Tensor still sitting on the host, uploaded to nothing yet. Or it holds a view, a chain of offsets into some other tensor's storage.

Reading a value forces the first state. Everything else in this arc is about which lines of your program do that reading and what it costs, so it pays to be able to answer, for any tensor in a loop you are debugging, which of the four it is in right now.

The transition into the recording has one detail people trip on. GetIrValue() is where a materialized tensor gets lazified: if the tensor currently holds only a device handle, the method mints a DeviceData IR node wrapping that handle. It deliberately does not clear the handle afterward. That is why calling it twice returns the same node instead of two parameters pointing at one buffer.

torch_xla does not own a graph walker. It subclasses the LazyTensor core that ships inside PyTorch itself and overrides the parts that need an xla::Shape where upstream wanted a torch::lazy::Shape. What XlaNode adds is worth enumerating, because two of the three additions are hashes and the third is the sharding annotation that feeds one of them.

The interior constructor's four initializers tell you the rule directly. node_hash_ combines the op kind with a seed. dag_hash_ folds in every operand's hash, so it identifies the whole subgraph rooted at this node rather than the node alone. Look at what is missing from both: the shape. An interior node's hash does not know its own dimensions.

Leaves are the exception, and GetOpHash is where the exception lives. A DeviceData or Scalar node hashes its op kind together with shape.ToString(), so dimensions enter the graph through the leaves and nowhere else. There is a third hash for sharding, updated whenever an annotation lands, folded into hash() only when the node actually carries one, so two identical graphs with different shardings can never share a compiled executable.

XlaNode::XlaNode(torch::lazy::OpKind op, torch::lazy::OpList operands,
                 std::vector<torch::lazy::Shape>&& shapes, xla::Shape xla_shape,
                 size_t num_outputs, torch::lazy::hash_t hash_seed)
    : torch::lazy::Node(op, operands, std::move(shapes), num_outputs),
      xla_shape_(std::move(xla_shape)),
      node_hash_(torch::lazy::HashCombine(op.hash(), hash_seed)),
      dag_hash_(GetOperandHashes(operands, node_hash_)) {}

torch::lazy::hash_t XlaNode::GetOpHash(torch::lazy::OpKind op,
                                       const xla::Shape& shape,
                                       torch::lazy::hash_t hash_seed) {
  torch::lazy::hash_t h =
      torch::lazy::HashCombine(op.hash(), torch::lazy::Hash(shape.ToString()));
  return torch::lazy::HashCombine(h, hash_seed);
}

Put the two constructors together and a debugging rule falls out. When a graph recompiles after you changed something in the middle of a model, the middle is not where the hash changed. Interior hashes cover op kind and operand structure, so a shape change reaches the graph hash by travelling up from the leaf that owns the shape.

The same rule has a mirror image that surprises people the other way. Two different device buffers of the same shape hash identically, because DeviceData uses the leaf constructor and the leaf constructor never saw a buffer address. Buffer identity reaches the compile decision through a separate channel: the parameter sequence, the ordered list of parameter indices recorded as the lowering walks the graph, hashed alongside everything else.

Shape inference itself is memoized in a global cache keyed by the node hash, sized by XLA_IR_SHAPE_CACHE_SIZE at a default of 12288 entries. A repeated subgraph shape is not re-inferred, which is a small thing until you are staring at a trace wondering why the second identical block cost nothing.

One function assembles the number that decides compile against cache hit, and reading as far as its first MergeHash gives you most of the recompilation story. The very first ingredient is not the graph. It is config.force_ltc_data, a boolean about whether this sync is allowed to alias buffers, merged before any tensor is looked at, with a comment saying exactly why.

Next in are the compilation environment hash and two git revisions, the pytorch one and the torch_xla one, generated into the build. Upgrading either library invalidates every entry in your persistent compilation cache, and that is the intent rather than an accident: a bug fixed in torch_xla should not be masked by a stale executable on disk.

After that come the per-tensor ingredients: each synced tensor's IR value hash, the parameter sequence, the buffer donor indices when there are any, and the auto-sharding settings. Absent by design are the bytes in any buffer and the HLO text itself.

Two runs with different weights and identical shapes land on the same cache entry. That is the point of hashing structure rather than contents.

  // The force_ltc_data controls aliasing compilation, so effectively the same
  // graph with on/off force_ltc_data should not match, hash wise.
  coll.hash = torch::lazy::MHash(config.force_ltc_data);
  // Ensure that the compilation environment and git revisions are reflected
  // in the hash, so that different versions of the code can produce different
  // hashes for the same graph.
  XLA_ASSIGN_OR_THROW(runtime::ComputationClient * absl_nonnull const client,
                      runtime::GetComputationClient());
  MergeHash(
      {client->HashCompilationEnv(), torch::lazy::StringHash(TORCH_GITREV),
       torch::lazy::StringHash(XLA_GITREV)},
      &coll.hash);

You can now derive the recompile list instead of memorizing it. An input shape changes, so a leaf hash changes. A new op appears, or the same ops run in a different order, so the DAG hash changes. A stray tensor is alive at the barrier that was not alive last step, so a different set of IR value hashes merges. A sharding annotation moves. A library upgrade lands. Each of those is one ingredient of one number.

The last entry on the list is not about your program at all. The compiled result lands in an LRU sized by XLA_COMPILATION_CACHE_SIZE, default 2048 entries, or in a persistent on-disk cache when XLA_PERSISTENT_CACHE_PATH is set. Evict an entry and the next hit becomes a miss, which on the lazy path costs a compile and on the dynamo path costs something worse. The third lesson in this arc takes that apart.

This is the same cache-key discipline chapter 6 taught for dynamo guards, paid a second time at a different layer, against a key you now know the ingredients of. The two layers do not share a key and they do not share a cache, so a program can be perfectly stable under dynamo and still recompile every step down here.