Memory planning: Support shared_allocation with offset

[Erik-Lundell]

Previously, there was a mechanism that let
passes annotate nodes with a meta field
"_share_alloc_with_arg_idx", to indicate to
memory planning algorithms that a node output
tensor is shared. Generalize this by adding
the meta field "_shared_alloc_offset", to allow
setting an offset from the base shared allocation.

This meta is picked up by the MemoryPlanningPass,
and transfered to the TensorSpec. To do this,
a new field "storage_base_offset" is added to the
TensorSpec. "inplace_base" is renamed to
"storage_base", while keeping a inplace_base as
an alias. When inplace_base
is used, the storage_base_offset is not allowed
to be non-zero.

This work exposed an issue where planning of
inplace_base/storage_base backing tensors didn't
respect the lifetimes of backed tensors. Fix
this by extending lifetimes accordingly in
update_all_tensors_lifetime. Handling of unexpected values of TensorSpec in _move_memory_meta_to_spec
were also changed to errors instead of silent
returns.

Finally, modify the greedy algorithm to respect
the new offset.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21840

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 0f4e271 with merge base 79e3eae ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Erik-Lundell]

@claude review this PR

[claude]

Claude finished @Erik-Lundell's task in 3m 25s —— View job

---

PR Review: Memory planning — shared_allocation with offset

Reviewed the full diff against origin/main. This is a clean, well-tested generalization. Below are my findings, ordered by importance.

---

🐛 Bug: duplicated variable in alignment error message

exir/memory_planning.py:1015-1018

internal_assert(
    (base_alloc_offset + storage_base_offset) % spec.alignment == 0,
    f"Storage-backed TensorSpec allocation must respect alignment, got offset {storage_base_offset} inside parent with offset {storage_base_offset} for alignment {spec.alignment}.",
)

The message interpolates storage_base_offset twice. The second one — "inside parent with offset" — is clearly meant to be base_alloc_offset. As written the diagnostic hides the actual parent offset, which is exactly the value you'd want when debugging an alignment failure. Fix this →

---

⚠️ Weaker verification for common-root specs

exir/memory_planning.py:205-217

The overlap check was relaxed from "direct inplace pair" to "share a common storage_root":

is_common_base_pair = Verifier.storage_root(lhs_spec) is Verifier.storage_root(rhs_spec)

This is correct for the intended cases, but note it now unconditionally permits any storage+lifetime overlap between two specs sharing a root — including two sibling children whose [offset, offset+size) ranges genuinely overlap when they shouldn't (e.g. a pass sets bad offsets). The old code was narrower. Since _resolve_storage_base_specs already asserts each child fits within its base (storage_base_offset + allocated_memory <= base_allocated_memory), a truly overlapping-siblings bug would slip past the verifier but not past resolution — so this is acceptable, just worth being aware that the verifier no longer catches sibling overlaps. Consider a follow-up that checks actual byte-range overlap among common-root specs rather than short-circuiting.

---

🔎 Minor / questions

1. Removed public-ish materialize_buffer (memory_planning.py, was module-level, no underscore). Inlining into _compute_total_sizes is fine in-tree (grepped — no other callers), but it had no leading underscore, so out-of-tree code could theoretically import it. Low risk; flagging in case it's considered semi-public.

2. _shared_alloc_offset without _share_alloc_with_arg_idx (memory_planning_pass.py:_move_memory_meta_to_spec): if a node carries _shared_alloc_offset but no _share_alloc_with_arg_idx, the function returns early and the offset is silently ignored. Given you deliberately converted other unexpected states to errors, an offset-without-base combination might also deserve a warning/error rather than a silent no-op.

3. Multi-output out-var nodes: _move_memory_meta_to_spec is only invoked on the len(out_arg_names) == 1 branch (memory_planning_pass.py:237). This matches the prior behavior (single-output only), so no regression — just confirming storage sharing on multi-output nodes is intentionally out of scope.

4. Stricter _move_memory_meta_to_spec now raises TypeError/IndexError where the old code silently skipped. This is a behavioral change for any existing pass that set _share_alloc_with_arg_idx loosely (e.g. out-of-range index). The PR body calls this out as intended — good, just make sure downstream Arm passes are aligned.

---

✅ Strengths

• The lifetime-extension fix (_extend_storage_base_lifetimes called from update_all_tensors_lifetime) is the right place to solve the base-lifetime bug, and walks the full chain with a cycle guard.
• Chain resolution in _resolve_storage_base_specs correctly accumulates offsets through multi-level chains (base → child → grandchild), verified by test_greedy_resolves_chained_storage_base.
• Cycle detection is consistent across storage_root, _extend_storage_base_lifetimes, and _resolve_storage_base_specs.
• inplace_base alias with the zero-offset invariant (asserted in both getter and setter) is a clean backward-compat shim, and _ViewSpec._self_fields was correctly updated to include the new fields.
• Test coverage is thorough: offset placement, chained aliasing, lifetime reservation vs. reuse, verifier acceptance, and the inplace_base offset-0 guard.

---

Overall: solid, well-scoped change. Only the duplicated-variable message is a concrete bug worth fixing before merge; the rest are minor considerations.
· branch generalize-memory-planning