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