mirror of
https://github.com/tensorflow/tensorflow.git
synced 2026-09-28 05:13:36 +08:00
Jian Cai
b5d8801b1e
Root Cause:
`MemorySpaceAssignment::Process()` rewrites the HLO graph in two sequential loops over all `Allocation` objects:
1. Loop 1 calls the virtual method `Allocation::Process()` (`CopyAllocation::Process()`, `PinnedAllocation::Process()`, and `ParentAllocation::Process()`).
2. Loop 2 calls the virtual method `Allocation::PostProcess()` (`ParentAllocation::PostProcess()`).
Several HLO transformations across these two loops failed to maintain or propagate `OriginalValue`:
1. While-loop widening in `ParentAllocation::Process()` (Loop 1) and `ParentAllocation::PostProcess()` (Loop 2):
When `ParentAllocation` passes a buffer from a parent computation into a `while` loop, it widens the `while` loop tuple from N elements to N+1 elements at `new_tuple_index`:
- `ParentAllocation::Process()` widens the `while` loop inputs by creating a new input tuple (`new_while_operand = TupleUtil::ReplaceTupleWith(...)`) and mutating the shapes of `calling_instruction_` (the `kWhile` instruction), `while_condition->parameter_instruction(0)`, and `while_body->parameter_instruction(0)` in-place. Because the in-place mutated instructions kept their old N-element `OriginalValue` trees while their shapes became (N+1)-tuples, `HloVerifier` failed with shape mismatches. Meanwhile, `new_while_operand` was created with `original_value() == nullptr`.
- To keep downstream users of the `while` loop valid, `ParentAllocation::Process()` creates `tuple_with_old_shape = TupleUtil::ExtractPrefix(calling_instruction_, new_tuple_index)` and replaces uses of `calling_instruction_` with `tuple_with_old_shape`. `tuple_with_old_shape` and its `GetTupleElement` operands were created with `original_value() == nullptr`, dropping value tracking for all users after the loop.
- `ParentAllocation::PostProcess()` runs in Loop 2 (after all Loop 1 allocations inside `while_body` have finished) to widen the `while_body` root tuple (`new_while_body_root = TupleUtil::ReplaceTupleWith(added_element, old_body_root, ...)`), which also left `new_while_body_root` with `original_value() == nullptr` and dropped the appended element's `OriginalValue` at `{new_tuple_index}`.
2. Tuple replacement and async copy creation in `CopyAllocation::Process()` and `Allocation::UpdateUses()` (Loop 1):
- `CopyAllocation::Process()` inserts `copy-start` and `copy-done` instructions to move buffers between default memory (HBM) and alternate memory (VMEM), but left `copy_done_` with `original_value() == nullptr` instead of copying `producing_instruction`'s `OriginalValue`.
- `CopyAllocation::Process()`, `PinnedAllocation::Process()`, and `ParentAllocation::Process()` all call the helper `Allocation::UpdateUses()` during Loop 1 to rewire consumer instructions to use `producing_instruction`. When a consumer is a `kTuple` instruction (such as `while_body->root_instruction()` or `while_init`), `Allocation::UpdateUses()` calls `TupleUtil::ReplaceTupleWith`, which constructs a brand-new `kTuple` instruction with `original_value() == nullptr`. Consequently, any Loop 1 allocation updating an element of `while_body->root_instruction()` wiped out its `OriginalValue` to `nullptr` before Loop 2 (`ParentAllocation::PostProcess()`) even ran.
What This CL Fixes:
- Add `TupleTree::num_children(ShapeIndexView index = {})` in `tuple_tree.h` to query the number of children of a tuple node directly from the index table.
- Introduce `AppendToTupleOriginalValue(src_tuple, dest_tuple, appended_instr)` in `allocation.cc`, which uses `CopyOriginalValue` from `hlo_original_value_util.h` to map existing tuple elements from `src_tuple` to the widened `dest_tuple` and copies `appended_instr->original_value()->tree()` into the last element (`dest_tuple->shape().tuple_shapes().size() - 1`) via `CopyCompatibleSubtreeFrom`.
- In `ParentAllocation::Process()`, before widening `calling_instruction_` in-place, copy its N-element `OriginalValue` to `tuple_with_old_shape` via `CopyOriginalValue` and propagate subtrees to its `GetTupleElement` operands. Then widen `new_while_operand` with `AppendToTupleOriginalValue` and update `calling_instruction_`, `while_condition` parameter 0, and `while_body` parameter 0 via `AppendToWhileLoopOriginalValue`.
- In `AppendToWhileLoopOriginalValue` (`while_util.cc`), only update `while_body->root_instruction()` if its shape has already been widened to match `while_instr->shape()`, supporting two-phase while-loop widening where `while_body` root widening is deferred to `ParentAllocation::PostProcess()`.
- In `Allocation::AddGetTupleElements()`, propagate the `OriginalValue` subtree from `defining_position().instruction` at `defining_position().index` to the newly created `GetTupleElement` instruction.
- In `ParentAllocation::PostProcess()`, call `AppendToTupleOriginalValue(old_body_root, new_while_body_root, added_element)` so the widened `while_body` root tuple preserves both the body root's `OriginalValue` and the pass-through element's `OriginalValue`.
- In `CopyAllocation::Process()`, propagate `producing_instruction`'s `OriginalValue` to `copy_done_` (and any intermediate bitcast).
- In `Allocation::UpdateUses()`, when `TupleUtil::ReplaceTupleWith` replaces a `kTuple` instruction, clone `tuple_inst`'s `OriginalValue` onto `replacement_instruction` and update `{use.operand_index}` from `producing_instruction->original_value()`.
PiperOrigin-RevId: 987774950
XLA
XLA (Accelerated Linear Algebra) is an open-source machine learning (ML) compiler for GPUs, CPUs, and ML accelerators.
The XLA compiler takes models from popular ML frameworks such as PyTorch, TensorFlow, and JAX, and optimizes them for high-performance execution across different hardware platforms including GPUs, CPUs, and ML accelerators.
openxla.org is the project's website.
Get started
If you want to use XLA to compile your ML project, refer to the corresponding documentation for your ML framework:
If you're not contributing code to the XLA compiler, you don't need to clone and build this repo. Everything here is intended for XLA contributors who want to develop the compiler and XLA integrators who want to debug or add support for ML frontends and hardware backends.
Contribute
If you'd like to contribute to XLA, review How to Contribute and then see the developer guide.
Contacts
- For questions, contact the maintainers - maintainers at openxla.org
Resources
Code of Conduct
While under TensorFlow governance, all community spaces for SIG OpenXLA are subject to the TensorFlow Code of Conduct.