From 7a8b8686e0eb22e400390465450dde8e7fd269b2 Mon Sep 17 00:00:00 2001 From: Conan Truong Date: Wed, 19 Aug 2026 20:36:30 -0700 Subject: [PATCH] Arena-aware sharing of mutable buffers via shared_buffer_fqns (#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 --- exir/passes/memory_planning_pass.py | 243 +++++++++++++++++- exir/tests/test_memory_planning.py | 384 ++++++++++++++++++++++++++++ 2 files changed, 617 insertions(+), 10 deletions(-) diff --git a/exir/passes/memory_planning_pass.py b/exir/passes/memory_planning_pass.py index 99a5f3dd8ec..1535def817b 100644 --- a/exir/passes/memory_planning_pass.py +++ b/exir/passes/memory_planning_pass.py @@ -4,6 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from __future__ import annotations + import itertools import logging import warnings @@ -30,6 +32,7 @@ from torch import fx from torch.export.exported_program import ExportGraphSignature from torch.fx import Node +from torch.utils import _pytree as pytree # copied from https://stackoverflow.com/questions/75582932/python-how-can-i-print-the-function-name-of-a-partial-function @@ -136,6 +139,119 @@ def _check_default_mem_ids(gm: torch.fx.GraphModule): ) +def _align_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +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"] may be a nested pytree (SpecPropPass), so flatten to leaves + for spec in pytree.tree_leaves(node.meta.get("spec")): + if not isinstance(spec, TensorSpec) or id(spec) in seen: + continue + seen.add(id(spec)) + specs.append(spec) + return specs + + +def _resolve_inplace_root(spec: TensorSpec) -> TensorSpec: + """Follow ``spec.inplace_base`` to the terminal spec it aliases. + + In-place ops produce a distinct result spec that aliases an input's spec + (greedy gives it the same mem_id/mem_offset). The chain can be several links + long; a ``seen`` set guards against a pathological cycle. + """ + seen: set[int] = set() + cur = spec + while getattr(cur, "inplace_base", None) is not None and id(cur) not in seen: + seen.add(id(cur)) + cur = cur.inplace_base + return cur + + +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: + # A declared shared buffer may be read-only in one method while mutated in + # another; collect it whenever it is a buffer so it front-packs to the same + # slot in every method. + is_buffer, fqn = _is_buffer(node, graph_signature) + if is_buffer and fqn in declared: + assert fqn is not None + specs_by_fqn[fqn] = _get_spec_from_node(node) + return specs_by_fqn + + +def _relocate_around_front_region( + graph_module: torch.fx.GraphModule, + declared_front: dict[int, int], + reserved: dict[int, int], + vacated: dict[int, list[tuple[int, int]]], +) -> None: + """Move each non-declared spec up by its arena's reserved front region, less + the declared columns vacated below it (so the arena does not grow). A spec that + aliases a declared buffer in place follows that buffer to the front -- both + ``mem_id`` and ``mem_offset`` -- so the aliased write lands on the shared + buffer rather than a shifted copy, even if the alias were given a different + arena. + """ + for spec in _iter_unique_specs(graph_module): + if id(spec) in declared_front: + continue + root = _resolve_inplace_root(spec) + root_front = declared_front.get(id(root)) + if root_front is not None: + spec.mem_id = root.mem_id + spec.mem_offset = root_front + continue + mem_id = spec.mem_id + mem_offset = spec.mem_offset + if mem_id is None or mem_offset is None: + continue + below = sum(size for off, size in vacated.get(mem_id, []) if off < mem_offset) + spec.mem_offset = reserved.get(mem_id, 0) + mem_offset - below + + +def _grow_arena_sizes( + graph_module: torch.fx.GraphModule, + reserved: dict[int, int], + vacated: dict[int, list[tuple[int, int]]], +) -> None: + """Update each front-packed arena's size. The scalar ``front - reclaimed`` term + preserves any non-spec size the algorithm folded into ``non_const_buffer_sizes`` + (submodule / XNNPACK padding); the ``max`` against the real high-water mark of + the final placements keeps the arena from being recorded smaller than actual + usage if an algorithm ever leaves an unaligned slot. + """ + sizes = list(graph_module.meta.get("non_const_buffer_sizes", [])) + high_water: dict[int, int] = {} + for spec in _iter_unique_specs(graph_module): + if spec.mem_id is None or spec.mem_offset is None: + continue + end = spec.mem_offset + spec.allocated_memory + if end > high_water.get(spec.mem_id, 0): + high_water[spec.mem_id] = end + for mem_id, front in reserved.items(): + while len(sizes) <= mem_id: + sizes.append(0) + reclaimed = sum(size for _, size in vacated.get(mem_id, [])) + sizes[mem_id] = max( + sizes[mem_id] + front - reclaimed, high_water.get(mem_id, 0) + ) + graph_module.meta["non_const_buffer_sizes"] = sizes + + @dataclass class _MemoryPlanningState: mutable_buffers: Dict[str, Set[TensorSpec]] = field(default_factory=dict) @@ -153,12 +269,20 @@ def __init__( alloc_mutable_buffers: bool = True, share_mutable_buffers: bool = False, alignment: int = ALIGNMENT, + shared_buffer_fqns: frozenset[str] | None = None, ) -> None: r""" alloc_graph_input/alloc_graph_output will have 4 different combinations to control if the memory planning algorithm need allocate memory for the graph input/output. The default behavior is the algorithm will allocate memory for both graph input and output. + + shared_buffer_fqns opts into the arena-aware sharing path. When set (and + share_mutable_buffers is True), the named mutable buffers keep the real + arena each is planned onto and are front-packed to an identical offset in + every method, instead of being forced onto the legacy dedicated mem_id 2 + arena. This lifts the single-default-arena requirement so programs with a + device/accelerator arena can still share a mutable buffer. """ if memory_planning_algo is None: memory_planning_algo = MemoryPlanningAlgorithmSuite() @@ -166,14 +290,26 @@ def __init__( raise ValueError( "share_mutable_buffers is only meaningful when alloc_mutable_buffers is True" ) + if shared_buffer_fqns is not None and not share_mutable_buffers: + raise ValueError("shared_buffer_fqns requires share_mutable_buffers=True") + if shared_buffer_fqns is not None and not shared_buffer_fqns: + raise ValueError( + "shared_buffer_fqns is empty; pass None for legacy sharing of all " + "mutable buffers, or a non-empty set to front-pack named buffers" + ) self.memory_planning_algo: Callable[..., List[int]] = memory_planning_algo self.allow_lifetime_and_storage_overlap = allow_lifetime_and_storage_overlap self.alloc_graph_input = alloc_graph_input self.alloc_graph_output = alloc_graph_output self.alloc_mutable_buffers = alloc_mutable_buffers self.share_mutable_buffers = share_mutable_buffers + self.shared_buffer_fqns: frozenset[str] | None = shared_buffer_fqns self.alignment = alignment self.state = _MemoryPlanningState() + # Resulting (mem_id, mem_offset, allocated_memory) of each declared + # shared buffer from the first method it appears in. Later methods must + # agree; a mismatch would alias two different tensors in one arena. + self._shared_placement: dict[str, tuple[int | None, int | None, int]] = {} # Set by EdgeProgramManager.to_executorch() from the top-level # ExecutorchBackendConfig. When True, apply_algo partitions specs by # device so non-CPU buffers get their own memory arenas. @@ -251,6 +387,14 @@ def run( # passes/stages is quite natural and avoid yet another 'context' data structure # to do the job. + # Shared mutable buffers are excluded from the main algo (and placed on + # the dedicated mem_id 2 arena later in run_multimethod) ONLY on the legacy + # path. The arena-aware path (shared_buffer_fqns set) keeps them in the + # algo so each lands on its real arena, then front-packs them below. + plan_mutable_buffers_in_algo = self.alloc_mutable_buffers and ( + not self.share_mutable_buffers or self.shared_buffer_fqns is not None + ) + _ = apply_algo( self.memory_planning_algo, graph_module, @@ -258,16 +402,17 @@ def run( graph_signature, self.alloc_graph_input, self.alloc_graph_output, - # If mutable buffers are shared, then do not allocate them in the - # main memory planning algo; they are allocated in run_multimethod. - self.alloc_mutable_buffers and not self.share_mutable_buffers, + plan_mutable_buffers_in_algo, self.enable_non_cpu_memory_planning, ) if self.share_mutable_buffers and graph_signature is not None: - self.state.graph_modules.append(graph_module) - _check_default_mem_ids(graph_module) - _insert_mutable_buffer_specs(self.state, graph_module, graph_signature) + if self.shared_buffer_fqns is None: + self.state.graph_modules.append(graph_module) + _check_default_mem_ids(graph_module) + _insert_mutable_buffer_specs(self.state, graph_module, graph_signature) + else: + self._front_pack_shared_buffers(graph_module, graph_signature) # TODO: make the verifier do the work recursively to handle # control flow @@ -275,10 +420,7 @@ def run( graph_module, self.alloc_graph_input, self.alloc_graph_output, - # If mutable buffers are shared, they are allocated after the - # main memory planning algo in run_multimethod, and should be - # skipped in the Verifier. - self.alloc_mutable_buffers and not self.share_mutable_buffers, + plan_mutable_buffers_in_algo, graph_signature, ) @@ -300,6 +442,87 @@ def run( verifier.verify_storage_reuse() return PassResult(graph_module, True) + def _front_pack_shared_buffers( + self, + graph_module: torch.fx.GraphModule, + graph_signature: ExportGraphSignature, + ) -> None: + """Pin the declared shared buffers to a deterministic front region. + + Each declared buffer keeps its own arena (``mem_id``) and is packed, in + sorted-FQN order, into a region reserved at the front of that arena. + Every other spec in the arena is relocated up by that region, less the + size of the declared columns that sat below it. A shared buffer has an + infinite lifetime, so the algorithm gave it an exclusive slot that no + other spec reuses; reclaiming that vacated slot in the same pass leaves no + dead space, so the arena does not grow. A buffer has the same ``mem_id`` + and size in every method, so its resulting front offset is identical + across methods by construction. + """ + declared = self.shared_buffer_fqns + assert declared is not None + specs_by_fqn = _collect_declared_shared_specs( + graph_module, graph_signature, declared + ) + missing = declared - specs_by_fqn.keys() + if missing: + raise ValueError( + "shared_buffer_fqns declares buffer(s) not present as buffers in " + f"this method: {sorted(missing)}" + ) + + reserved: dict[int, int] = {} + placement: dict[str, int] = {} + # id() of each declared buffer's placeholder spec -> its front offset, so + # specs that alias a declared buffer in place land on the same offset. + declared_front: dict[int, int] = {} + # Per arena, the (algo offset, size) each declared buffer was assigned. + # Those slots are vacated by the move to the front and reclaimed so the + # arena grows by nothing rather than by the whole front region. + vacated: dict[int, list[tuple[int, int]]] = {} + for fqn in sorted(specs_by_fqn): + spec = specs_by_fqn[fqn] + mem_id = spec.mem_id + if mem_id is None: + raise ValueError( + f"Declared shared buffer '{fqn}' was not assigned a memory arena" + ) + if spec.mem_offset is None: + raise ValueError( + f"Declared shared buffer '{fqn}' was not assigned an offset" + ) + vacated.setdefault(mem_id, []).append( + (spec.mem_offset, spec.allocated_memory) + ) + offset = reserved.get(mem_id, 0) + placement[fqn] = offset + declared_front[id(spec)] = offset + reserved[mem_id] = _align_up(offset + spec.allocated_memory, self.alignment) + + # Relocate the other specs around the reserved front region (declared + # buffers are shifted first so nothing lands back on the region), then + # place the declared buffers at the front and grow each arena to fit. + _relocate_around_front_region(graph_module, declared_front, reserved, vacated) + for fqn, offset in placement.items(): + specs_by_fqn[fqn].mem_offset = offset + _grow_arena_sizes(graph_module, reserved, vacated) + + self._validate_shared_placement(specs_by_fqn) + + def _validate_shared_placement(self, specs_by_fqn: dict[str, TensorSpec]) -> None: + for fqn, spec in specs_by_fqn.items(): + placement = (spec.mem_id, spec.mem_offset, spec.allocated_memory) + prior = self._shared_placement.get(fqn) + if prior is None: + self._shared_placement[fqn] = placement + elif prior != placement: + raise ValueError( + f"Shared buffer '{fqn}' has inconsistent placement across " + f"methods: {prior} (first method) != {placement} (this " + "method); a declared shared buffer must have identical size " + "and resulting placement in every method" + ) + def run_multimethod(self): """Resolve any memory planning done across entry points, called after run is called on all entry points.""" if self.share_mutable_buffers: diff --git a/exir/tests/test_memory_planning.py b/exir/tests/test_memory_planning.py index 31f3b1844c2..be1eca59971 100644 --- a/exir/tests/test_memory_planning.py +++ b/exir/tests/test_memory_planning.py @@ -1765,3 +1765,387 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: if node.op == "call_function" ) self.assertFalse(has_inplace) + + +class SharedStateArenaModel(nn.Module): + """A mutable buffer plus an unrelated activation that a custom pass pins to + a non-default arena, so the shared buffer coexists with a device-style arena. + """ + + def __init__(self) -> None: + super().__init__() + self.register_buffer("state", torch.zeros(4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.state.add_(x) + y = x * x + return self.state + y + + def reset(self, z: torch.Tensor) -> None: + self.state.copy_(z) + + +class _ForceArenaThreePass(MemoryPlanningPass): + """Pins every ``mul.out`` output onto arena 3, mimicking a program that has + an unrelated tensor on a non-default (e.g. device) arena. + """ + + def run( + self, + graph_module: torch.fx.GraphModule, + graph_signature: Optional[ExportGraphSignature] = None, + ) -> PassResult: + for node in graph_module.graph.nodes: + if node.op == "call_function" and node.target == torch.ops.aten.mul.out: + for spec in get_node_tensor_specs(node): + spec.mem_id = 3 + return super().run(graph_module, graph_signature) + + +class TestSharedBufferMemoryPlanning(unittest.TestCase): + """Arena-aware sharing of mutable buffers via ``shared_buffer_fqns``.""" + + def _prepare_state_model( + self, size: int + ) -> Tuple[GraphModule, ExportGraphSignature]: + class StateModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("state", torch.zeros(size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.state.add_(x) + return self.state * 2 + + model = StateModel().eval() + edge = to_edge(export(model, (torch.ones(size),), strict=True)) + gm = edge.exported_program().graph_module + gs = edge.exported_program().graph_signature + gm = PassManager(passes=[SpecPropPass(), ToOutVarPass()])(gm).graph_module + return gm, gs + + def _declared_specs( + self, + gm: GraphModule, + gs: ExportGraphSignature, + fqns: frozenset[str], + ) -> dict[str, TensorSpec]: + mutated = set(gs.buffers_to_mutate.values()) + out: dict[str, TensorSpec] = {} + for node in gm.graph.nodes: + if node.op == "placeholder" and isinstance(node.target, str): + fqn = gs.inputs_to_buffers.get(node.target) + if fqn in fqns and fqn in mutated: + out[fqn] = get_node_tensor_specs(node)[0] + return out + + def _build_shared_program( + self, mem_pass: MemoryPlanningPass + ) -> Any: # pyre-ignore[3] + model = SharedStateArenaModel().eval() + forward_ep = export(model, (torch.ones(4),)) + with patch_forward(model, model.reset): + reset_ep = export(model, (torch.zeros(4),)) + edge = to_edge({"forward": forward_ep, "reset": reset_ep}) + return edge.to_executorch( + ExecutorchBackendConfig( + memory_planning_pass=mem_pass, + emit_mutable_buffer_names=True, + ) + ) + + def test_multi_arena_shared_buffer(self) -> None: + """Shared buffer gets identical (mem_id, offset) across methods, and an + unrelated non-default-arena tensor does not trip the old blanket guard. + """ + et = self._build_shared_program( + _ForceArenaThreePass( + share_mutable_buffers=True, + shared_buffer_fqns=frozenset({"state"}), + ) + ) + + placements = [] + for plan in et.executorch_program.execution_plan: + state_vals = [ + v + for v in plan.values + if hasattr(v.val, "extra_tensor_info") + and v.val.extra_tensor_info is not None + and v.val.extra_tensor_info.fully_qualified_name == "state" + ] + self.assertEqual(len(state_vals), 1) + ai = state_vals[0].val.allocation_info + placements.append((ai.memory_id, ai.memory_offset_low)) + + self.assertEqual(placements[0], placements[1]) + self.assertEqual(placements[0][0], 1) + self.assertEqual(placements[0][1], 0) + + # The unrelated arena-3 tensor is present (would have raised before). + forward_plan = et.executorch_program.execution_plan[0] + self.assertGreaterEqual(len(forward_plan.non_const_buffer_sizes), 4) + self.assertGreater(forward_plan.non_const_buffer_sizes[3], 0) + + def test_legacy_path_rejects_non_default_arena(self) -> None: + """Without shared_buffer_fqns, the legacy guard rejects a non-default arena.""" + with self.assertRaises(ValueError): + self._build_shared_program(_ForceArenaThreePass(share_mutable_buffers=True)) + + def test_orphan_shared_buffer_fqns_raises(self) -> None: + with self.assertRaises(ValueError): + MemoryPlanningPass(shared_buffer_fqns=frozenset({"state"})) + + def test_size_mismatch_raises(self) -> None: + mem_pass = MemoryPlanningPass( + share_mutable_buffers=True, + shared_buffer_fqns=frozenset({"state"}), + ) + gm1, gs1 = self._prepare_state_model(4) + mem_pass.run(gm1, gs1) + + gm2, gs2 = self._prepare_state_model(8) + with self.assertRaises(ValueError) as cm: + mem_pass.run(gm2, gs2) + self.assertIn("state", str(cm.exception)) + + def test_declared_but_absent_raises(self) -> None: + mem_pass = MemoryPlanningPass( + share_mutable_buffers=True, + shared_buffer_fqns=frozenset({"state", "nonexistent"}), + ) + gm, gs = self._prepare_state_model(4) + with self.assertRaises(ValueError) as cm: + mem_pass.run(gm, gs) + self.assertIn("nonexistent", str(cm.exception)) + + def test_declared_buffer_on_non_default_arena(self) -> None: + """The declared buffer itself lives on a non-default arena; it is + front-packed within that arena and agrees across methods. + """ + mem_pass = MemoryPlanningPass( + share_mutable_buffers=True, + shared_buffer_fqns=frozenset({"state"}), + ) + placements = [] + for _ in range(2): + gm, gs = self._prepare_state_model(4) + state_spec = self._declared_specs(gm, gs, frozenset({"state"}))["state"] + state_spec.mem_id = 2 # pin the buffer itself onto a non-default arena + mem_pass.run(gm, gs) + placements.append((state_spec.mem_id, state_spec.mem_offset)) + # The buffer is the only tensor on arena 2, so the arena must be + # exactly its size -- relocating it to the front leaves no hole. + self.assertEqual( + gm.meta["non_const_buffer_sizes"][2], state_spec.allocated_memory + ) + self.assertEqual(placements[0], (2, 0)) + self.assertEqual(placements[0], placements[1]) + + def test_shared_arena_no_hole(self) -> None: + """A buffer sharing an arena with activations is relocated to the front + without growing the arena: the compacting shift reclaims its old slot. + + Compared against plain planning of the same graph, so no size is + hardcoded; a residual (interior) hole shows up as extra bytes on arena 1. + """ + gm_base, gs_base = self._prepare_state_model(4) + MemoryPlanningPass().run(gm_base, gs_base) + baseline = gm_base.meta["non_const_buffer_sizes"][1] + + gm, gs = self._prepare_state_model(4) + MemoryPlanningPass( + share_mutable_buffers=True, + shared_buffer_fqns=frozenset({"state"}), + ).run(gm, gs) + self.assertEqual(gm.meta["non_const_buffer_sizes"][1], baseline) + + def test_multi_buffer_sorted_front_pack(self) -> None: + """Two declared buffers on one arena pack in sorted-FQN order at + accumulating offsets. + """ + + class TwoBufferModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("cache_v", torch.zeros(8)) + self.register_buffer("cache_k", torch.zeros(4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.cache_k.add_(x[:4]) + self.cache_v.add_(x) + return self.cache_k.sum() + self.cache_v.sum() + + model = TwoBufferModel().eval() + edge = to_edge(export(model, (torch.ones(8),), strict=True)) + gm = edge.exported_program().graph_module + gs = edge.exported_program().graph_signature + gm = PassManager(passes=[SpecPropPass(), ToOutVarPass()])(gm).graph_module + + declared = frozenset({"cache_k", "cache_v"}) + MemoryPlanningPass( + share_mutable_buffers=True, + shared_buffer_fqns=declared, + ).run(gm, gs) + + specs = self._declared_specs(gm, gs, declared) + k, v = specs["cache_k"], specs["cache_v"] + self.assertEqual(k.mem_id, v.mem_id) + # "cache_k" < "cache_v", so cache_k is packed first at offset 0 and + # cache_v accumulates directly after it. + self.assertEqual(k.mem_offset, 0) + self.assertEqual(v.mem_offset, k.allocated_memory) + + def test_multi_arena_declared_buffers(self) -> None: + """Declared buffers on different arenas are each front-packed at offset 0 + of their OWN arena; per-arena reserved/vacated bookkeeping keeps them + independent (one arena's front region does not shift the other's). + """ + + class TwoBufferModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("cache_k", torch.zeros(4)) + self.register_buffer("cache_v", torch.zeros(8)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.cache_k.add_(x[:4]) + self.cache_v.add_(x) + return self.cache_k.sum() + self.cache_v.sum() + + model = TwoBufferModel().eval() + edge = to_edge(export(model, (torch.ones(8),), strict=True)) + gm = edge.exported_program().graph_module + gs = edge.exported_program().graph_signature + gm = PassManager(passes=[SpecPropPass(), ToOutVarPass()])(gm).graph_module + + declared = frozenset({"cache_k", "cache_v"}) + specs = self._declared_specs(gm, gs, declared) + specs["cache_v"].mem_id = 2 # pin cache_v onto its own arena + + MemoryPlanningPass( + share_mutable_buffers=True, + shared_buffer_fqns=declared, + ).run(gm, gs) + + k, v = specs["cache_k"], specs["cache_v"] + # Each buffer front-packs at offset 0 of its own arena: cache_v on arena 2 + # does not push cache_k off offset 0 on arena 1, nor vice versa. + self.assertEqual((k.mem_id, k.mem_offset), (1, 0)) + self.assertEqual((v.mem_id, v.mem_offset), (2, 0)) + # cache_v is alone on arena 2, so it is sized to exactly the buffer. + self.assertEqual(gm.meta["non_const_buffer_sizes"][2], v.allocated_memory) + + def test_inplace_alias_follows_shared_buffer(self) -> None: + """A spec that aliases a declared buffer in place (inplace_base chains to + it) stays pinned to the buffer's front offset instead of being shifted. + + This mirrors the annotation-only reinplace path where a distinct alloc + spec aliases a write-only mutated buffer. Without pinning, the aliased + write silently diverges from the buffer. + """ + + class WriteOnlyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("state", torch.zeros(4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.state.add_(x) + return x * x + + model = WriteOnlyModel().eval() + edge = to_edge(export(model, (torch.ones(4),), strict=True)) + gm = edge.exported_program().graph_module + gs = edge.exported_program().graph_signature + gm = PassManager(passes=[SpecPropPass(), ToOutVarPass()])(gm).graph_module + + state_spec = self._declared_specs(gm, gs, frozenset({"state"}))["state"] + mul_spec = None + for node in gm.graph.nodes: + if node.op == "call_function" and node.target == torch.ops.aten.mul.out: + mul_spec = node.meta["spec"] + self.assertIsNotNone(mul_spec) + self.assertIsNot(mul_spec, state_spec) + # Alias the mul result onto the buffer, as _set_alloc_node_spec would for + # an annotation-only reinplaced write. + mul_spec.inplace_base = state_spec + + MemoryPlanningPass( + share_mutable_buffers=True, + shared_buffer_fqns=frozenset({"state"}), + ).run(gm, gs) + + self.assertEqual(state_spec.mem_offset, 0) + self.assertEqual(mul_spec.mem_id, state_spec.mem_id) + self.assertEqual(mul_spec.mem_offset, state_spec.mem_offset) + + def test_empty_shared_buffer_fqns_raises(self) -> None: + """An empty (but non-None) shared_buffer_fqns with sharing on is a + silently-broken config (no legacy and no arena-aware sharing), so reject it. + """ + with self.assertRaises(ValueError): + MemoryPlanningPass( + share_mutable_buffers=True, + shared_buffer_fqns=frozenset(), + ) + + def test_iter_unique_specs_flattens_nested_pytree(self) -> None: + """node.meta['spec'] can be an arbitrary pytree; a nested spec must still + be enumerated (and thus relocated), not only top-level list/tuple entries. + """ + from executorch.exir.passes.memory_planning_pass import _iter_unique_specs + + gm, _ = self._prepare_state_model(4) + nested_spec = TensorSpec.from_tensor(torch.zeros(4)) + placeholder = next(n for n in gm.graph.nodes if n.op == "placeholder") + placeholder.meta["spec"] = {"outer": [nested_spec]} + + found = {id(s) for s in _iter_unique_specs(gm)} + self.assertIn(id(nested_spec), found) + + def test_read_only_in_one_method_shared_buffer(self) -> None: + """A declared buffer mutated in one method but read-only in another is + collected (via _is_buffer) and front-packed to the same slot in both, + instead of being rejected as 'not present' in the read-only method. + """ + + class TwoMethodModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("state", torch.zeros(4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.state.add_(x) # mutates + return self.state * 2 + + def peek(self, x: torch.Tensor) -> torch.Tensor: + return self.state + x # read-only + + model = TwoMethodModel().eval() + forward_ep = export(model, (torch.ones(4),)) + with patch_forward(model, model.peek): + peek_ep = export(model, (torch.ones(4),)) + + mem_pass = MemoryPlanningPass( + share_mutable_buffers=True, + shared_buffer_fqns=frozenset({"state"}), + ) + + offsets = [] + for ep in (forward_ep, peek_ep): + edge = to_edge(ep) + gm = edge.exported_program().graph_module + gs = edge.exported_program().graph_signature + gm = PassManager(passes=[SpecPropPass(), ToOutVarPass()])(gm).graph_module + mem_pass.run(gm, gs) # must not raise, including the read-only method + spec = None + for node in gm.graph.nodes: + if node.op == "placeholder" and isinstance(node.target, str): + if gs.inputs_to_buffers.get(node.target) == "state": + spec = get_node_tensor_specs(node)[0] + self.assertIsNotNone(spec) + offsets.append(spec.mem_offset) + + self.assertEqual(offsets[0], 0) + self.assertEqual(offsets[0], offsets[1])