Arena-aware sharing of mutable buffers via shared_buffer_fqns (#21958) - #21958
Open
Conarnar wants to merge 1 commit into
Open
Arena-aware sharing of mutable buffers via shared_buffer_fqns (#21958)#21958Conarnar wants to merge 1 commit into
Conarnar wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21958
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 1 Unrelated FailureAs of commit 7a8b868 with merge base 3caaaca ( NEW FAILURE - The following job has failed:
BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
Contributor
|
@Conarnar has exported this pull request. If you are a Meta employee, you can view the originating Diff in D116686347. |
This PR needs a
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends MemoryPlanningPass(share_mutable_buffers=True) with an opt-in shared_buffer_fqns mechanism to share selected mutable buffers across multiple entry points while preserving their algorithm-assigned arena (mem_id). Instead of forcing shared buffers into a dedicated arena and rejecting any non-default arenas, the new path front-packs declared buffers to deterministic offsets per arena and validates consistent placement across methods.
Changes:
- Add
shared_buffer_fqnstoMemoryPlanningPassand route sharing logic into an arena-aware front-packing path. - Implement front-packing/relocation helpers and cross-method placement validation for declared shared buffers.
- Add unit tests covering multi-arena programs, legacy guard behavior, and several front-pack/aliasing scenarios.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
exir/passes/memory_planning_pass.py |
Adds shared_buffer_fqns support, front-packing relocation, and cross-method placement validation for arena-aware shared buffers. |
exir/tests/test_memory_planning.py |
Adds test coverage for shared-buffer front packing across multiple arenas and related edge cases. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+179
to
+190
| def _collect_declared_shared_specs( | ||
| graph_module: torch.fx.GraphModule, | ||
| graph_signature: ExportGraphSignature, | ||
| declared: frozenset[str], | ||
| ) -> dict[str, TensorSpec]: | ||
| specs_by_fqn: dict[str, TensorSpec] = {} | ||
| for node in graph_module.graph.nodes: | ||
| is_mutable, fqn = _is_mutable_buffer(node, graph_signature) | ||
| if is_mutable and fqn in declared: | ||
| assert fqn is not None | ||
| specs_by_fqn[fqn] = _get_spec_from_node(node) | ||
| return specs_by_fqn |
Comment on lines
+420
to
+426
| mem_id = spec.mem_id | ||
| assert ( | ||
| mem_id is not None | ||
| ), f"Declared shared buffer '{fqn}' was not assigned a memory arena" | ||
| assert ( | ||
| spec.mem_offset is not None | ||
| ), f"Declared shared buffer '{fqn}' was not assigned an offset" |
meta-codesync
Bot
changed the title
Conarnar
added a commit
to Conarnar/executorch
that referenced
this pull request
Aug 19, 2026
…h#21958) Summary: `MemoryPlanningPass(share_mutable_buffers=True)` shares a mutable buffer across methods by moving every mutable buffer onto a dedicated `mem_id=2` arena and rejecting any tensor placed on a non-default `mem_id` (`_check_default_mem_ids`). That makes buffer sharing impossible for a program that already uses a device/accelerator arena: the shared buffer cannot stay on its real arena, and the unrelated device tensors trip the blanket guard. This adds an opt-in `shared_buffer_fqns` frozenset to `MemoryPlanningPass`. When it is set (alongside `share_mutable_buffers=True`), the named mutable buffers keep the real arena the planner assigns them and are front-packed to a deterministic offset in every method, instead of being forced onto the `mem_id=2` arena. Sorted-FQN ordering makes the resulting offset identical across methods by construction, and `_validate_shared_placement` raises if a buffer resolves to a different arena/offset/size in any method. Each declared buffer's algorithm-assigned slot is reclaimed as the other tensors shift (a compacting relocation), so no arena grows to make room for the front region. The legacy path (`shared_buffer_fqns` unset) is unchanged: it still excludes mutable buffers from the algorithm, places them on `mem_id=2` in `run_multimethod`, and enforces `_check_default_mem_ids`. Reading order: the `__init__`/`run` changes wire the opt-in and decide whether mutable buffers are planned in the main algorithm; `_front_pack_shared_buffers` does the per-method relocation; the rest are small helpers (`_align_up`, `_iter_unique_specs`, `_resolve_inplace_root`, `_collect_declared_shared_specs`, `_validate_shared_placement`). Differential Revision: D116686347
Conarnar
force-pushed
the
export-D116686347
branch
from
August 19, 2026 23:12
fac3792 to
2bf0806
Compare
Comment on lines
+145
to
+161
| def _iter_unique_specs(graph_module: torch.fx.GraphModule) -> list[TensorSpec]: | ||
| """Every TensorSpec reachable from graph node metas, deduplicated by identity. | ||
|
|
||
| A spec can be referenced by several nodes; shifting one twice would corrupt | ||
| the layout, so dedupe on id() rather than equality. | ||
| """ | ||
| seen: set[int] = set() | ||
| specs: list[TensorSpec] = [] | ||
| for node in graph_module.graph.nodes: | ||
| meta_spec = node.meta.get("spec") | ||
| candidates = meta_spec if isinstance(meta_spec, (list, tuple)) else [meta_spec] | ||
| for spec in candidates: | ||
| if not isinstance(spec, TensorSpec) or id(spec) in seen: | ||
| continue | ||
| seen.add(id(spec)) | ||
| specs.append(spec) | ||
| return specs |
…h#21958) Summary: `MemoryPlanningPass(share_mutable_buffers=True)` shares a mutable buffer across methods by moving every mutable buffer onto a dedicated `mem_id=2` arena and rejecting any tensor placed on a non-default `mem_id` (`_check_default_mem_ids`). That makes buffer sharing impossible for a program that already uses a device/accelerator arena: the shared buffer cannot stay on its real arena, and the unrelated device tensors trip the blanket guard. This adds an opt-in `shared_buffer_fqns` frozenset to `MemoryPlanningPass`. When it is set (alongside `share_mutable_buffers=True`), the named mutable buffers keep the real arena the planner assigns them and are front-packed to a deterministic offset in every method, instead of being forced onto the `mem_id=2` arena. Sorted-FQN ordering makes the resulting offset identical across methods by construction, and `_validate_shared_placement` raises if a buffer resolves to a different arena/offset/size in any method. Each declared buffer's algorithm-assigned slot is reclaimed as the other tensors shift (a compacting relocation), so no arena grows to make room for the front region. The legacy path (`shared_buffer_fqns` unset) is unchanged: it still excludes mutable buffers from the algorithm, places them on `mem_id=2` in `run_multimethod`, and enforces `_check_default_mem_ids`. Reading order: the `__init__`/`run` changes wire the opt-in and decide whether mutable buffers are planned in the main algorithm; `_front_pack_shared_buffers` does the per-method relocation; the rest are small helpers (`_align_up`, `_iter_unique_specs`, `_resolve_inplace_root`, `_collect_declared_shared_specs`, `_validate_shared_placement`). Differential Revision: D116686347
Conarnar
force-pushed
the
export-D116686347
branch
from
August 20, 2026 03:36
2bf0806 to
7a8b868
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
MemoryPlanningPass(share_mutable_buffers=True)shares a mutable buffer across methods by moving every mutable buffer onto a dedicatedmem_id=2arena and rejecting any tensor placed on a non-defaultmem_id(_check_default_mem_ids). That makes buffer sharing impossible for a program that already uses a device/accelerator arena: the shared buffer cannot stay on its real arena, and the unrelated device tensors trip the blanket guard.This adds an opt-in
shared_buffer_fqnsfrozenset toMemoryPlanningPass. When it is set (alongsideshare_mutable_buffers=True), the named mutable buffers keep the real arena the planner assigns them and are front-packed to a deterministic offset in every method, instead of being forced onto themem_id=2arena. Sorted-FQN ordering makes the resulting offset identical across methods by construction, and_validate_shared_placementraises if a buffer resolves to a different arena/offset/size in any method. Each declared buffer's algorithm-assigned slot is reclaimed as the other tensors shift (a compacting relocation), so no arena grows to make room for the front region.The legacy path (
shared_buffer_fqnsunset) is unchanged: it still excludes mutable buffers from the algorithm, places them onmem_id=2inrun_multimethod, and enforces_check_default_mem_ids.Reading order: the
__init__/runchanges wire the opt-in and decide whether mutable buffers are planned in the main algorithm;_front_pack_shared_buffersdoes the per-method relocation; the rest are small helpers (_align_up,_iter_unique_specs,_resolve_inplace_root,_collect_declared_shared_specs,_validate_shared_placement).Differential Revision: D116686347