From 26aed9f0c075debf359c6273bd9986cbf20df4b8 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 12 Aug 2026 16:25:48 -0700 Subject: [PATCH] Add a capacity-aware memory planner over a target memory map ExecuTorch plans tensors into numbered arenas and the runtime supports several of them, but nothing tells the planner how large an arena is. greedy honors whatever mem_id a custom pool pass assigned, puts everything else in arena 1, and over-subscription only surfaces as a link error or a runtime abort. That makes it awkward to target a part whose fast memory is a small tightly-coupled region alongside a larger SRAM. This adds TargetMemoryMap, an ordered list of banks each with a byte capacity, and banked_greedy, which plans against it: buffers fill the fastest bank that can hold them and spill to the next, and a buffer that fits nowhere fails the export rather than the device. With a single bank the plan is bit-identical to greedy, asserted per-spec at alignments 16, 32 and 256. The map is deliberately a subset of the target's pools. Declaring a bank hands that arena to the planner; an undeclared mem_id behaves exactly as it does under greedy, so a region that must hold only what a pass deliberately pinned there -- a DMA-visible pool, accelerator-private scratch -- is protected by leaving it out. Custom pool passes therefore compose: the mem_id they assign is honored as a pin, pins are placed before any unpinned buffer so nothing can take the space they need, and unpinned buffers still share storage with them. Reviewers may find these the least obvious parts. Placement is one pass in size-descending order, which is what gives each bank the ordering pick_shared_obj requires without any repacking; the sort keys on the incoming size, as greedy's does, because realigning first manufactures ties whose order then differs. The unpinned phase cannot call pick_shared_obj directly, because pins and unpinned buffers are each sorted but their concatenation is not, so _reusable_object applies the same two reuse rules read-only and filters undersized objects where that function asserts. A bank's alignment must be a multiple of the graph's, which is what keeps per-bank realignment order-preserving. Banking costs total bytes, and that is inherent rather than incidental: a shared object never spans banks, so once a buffer spills the sum of the arenas can exceed what greedy needed for one. On a 30-buffer graph with a 6 KiB fast bank that is 12288 against greedy's 10752. Pinning many buffers into a declared bank costs a little more again, since pins take first claim on its capacity. The trade buys residency in fast memory, not a smaller total. Nothing consumes this yet. The runtime still backs every arena from one pool, so per-bank regions in a runner are a follow-up, and an example belongs with that change rather than ahead of it. share_mutable_buffers is also deferred: it reserves arena 2 and needs a change to core memory planning. Authored with assistance from Claude Code. --- exir/BUCK | 13 + exir/banked_memory_planning.py | 477 ++++++++++++ exir/tests/targets.bzl | 21 + exir/tests/test_banked_memory_planning.py | 863 ++++++++++++++++++++++ 4 files changed, 1374 insertions(+) create mode 100644 exir/banked_memory_planning.py create mode 100644 exir/tests/test_banked_memory_planning.py diff --git a/exir/BUCK b/exir/BUCK index d70900c02ae..4646e68de28 100644 --- a/exir/BUCK +++ b/exir/BUCK @@ -185,6 +185,19 @@ fbcode_target(_kind = runtime.python_library, ], ) +fbcode_target(_kind = runtime.python_library, + name = "banked_memory_planning", + srcs = [ + "banked_memory_planning.py", + ], + deps = [ + ":memory_planning", + ":tensor", + "//caffe2:torch", + "//executorch/exir/passes:memory_planning_pass", + ], +) + fbcode_target(_kind = runtime.python_library, name = "common", srcs = [ diff --git a/exir/banked_memory_planning.py b/exir/banked_memory_planning.py new file mode 100644 index 00000000000..87e46956814 --- /dev/null +++ b/exir/banked_memory_planning.py @@ -0,0 +1,477 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +r"""Capacity-aware memory planning across a target's memory hierarchy. + +ExecuTorch plans tensors into numbered arenas (``mem_id``) and the runtime supports +several, but nothing tells the planner how large an arena is: ``greedy`` honors +whatever ``mem_id`` a custom pool pass assigned, puts everything else in arena 1, +and over-subscription surfaces as a link error or a runtime abort. + +:class:`TargetMemoryMap` supplies the missing sizes. Buffers fill the fastest bank +that can hold them and spill to the next; one that fits nowhere fails the export +rather than the device. With a single bank the plan is bit-identical to ``greedy``. + +Custom pool passes still work ahead of this planner, as described in +``docs/source/compiler-memory-planning.md``; the ``mem_id`` they assign is honored +as a pin. Per-buffer constraints richer than capacity -- DMA reachability, +cacheability -- are deliberately absent until a consumer can enforce them. +""" + +import bisect +import logging +from dataclasses import dataclass +from typing import Dict, Iterable, List, Optional, Sequence, Set, Tuple + +import torch + +from executorch.exir.memory_planning import ( + _compute_total_sizes, + _does_not_overlap, + _find_max_overlapping_allocations_offset, + _resolve_inplace_specs, + AllocationSpec, + get_node_tensor_specs, + MemoryAlgoResult, + MemoryPlanningAlgorithmSuite, + pick_shared_obj, + SharedObject, + SpecAllocResult, +) +from executorch.exir.passes.memory_planning_pass import MemoryPlanningPass +from executorch.exir.schema import DeviceType +from executorch.exir.tensor import ALIGNMENT, TensorSpec +from torch.export.exported_program import ExportGraphSignature + + +@dataclass(frozen=True) +class Bank: + """A memory region the planner may fill. + + ``size`` is the budget for *this model*, usually less than the physical region: + the stack, the method allocator and anything else resident there come out of it + first. ``mem_id`` is the runtime arena index, declared rather than inferred from + list position because the map is sparse. ``alignment`` tightens the graph's + alignment for this bank and must be a multiple of it; ``None`` inherits it. + """ + + name: str + size: int + mem_id: int + alignment: Optional[int] = None + + +class TargetMemoryMap: + """The pools the planner may fill, in preference order (fastest first). + + Declaring a bank hands that arena to the planner: it spills unpinned buffers + into it and enforces its capacity. An undeclared ``mem_id`` behaves as it does + under ``greedy`` -- a pass may pin buffers to it and they are laid out normally + -- but the planner never places anything there on its own. That is how a + DMA-visible pool or accelerator-private scratch is protected: leave it out. + + List order is preference; ``Bank.mem_id`` is identity. They are separate + because the map is sparse. + """ + + def __init__(self, banks: Sequence[Bank]) -> None: + if not banks: + raise ValueError("A target memory map needs at least one bank") + names = [bank.name for bank in banks] + duplicates = {name for name in names if names.count(name) > 1} + if duplicates: + raise ValueError(f"Duplicate bank names in target memory map: {duplicates}") + ids = [bank.mem_id for bank in banks] + clashes = {i for i in ids if ids.count(i) > 1} + if clashes: + raise ValueError(f"Duplicate mem_ids in target memory map: {clashes}") + for bank in banks: + if bank.size <= 0: + raise ValueError(f"Bank {bank.name} has non-positive size {bank.size}") + if bank.alignment is not None and bank.alignment < 1: + raise ValueError( + f"Bank {bank.name} has non-positive alignment {bank.alignment}" + ) + if bank.mem_id < 1: + raise ValueError( + f"Bank {bank.name} has mem_id {bank.mem_id}; arena 0 is reserved " + "for constants, so planned banks start at 1" + ) + self.banks: List[Bank] = list(banks) + self._by_id: Dict[int, Bank] = {bank.mem_id: bank for bank in banks} + + def mem_ids(self) -> List[int]: + """Declared arena ids, in preference order.""" + return [bank.mem_id for bank in self.banks] + + def manages(self, mem_id: int) -> bool: + return mem_id in self._by_id + + def bank(self, mem_id: int) -> Bank: + if mem_id not in self._by_id: + raise ValueError(f"mem_id {mem_id} is not in the target memory map") + return self._by_id[mem_id] + + def mem_id(self, name: str) -> int: + for bank in self.banks: + if bank.name == name: + return bank.mem_id + raise ValueError(f"No bank named {name!r} in the target memory map") + + def describe(self) -> str: + return ", ".join( + f"{bank.name}(mem_id={bank.mem_id}, size={bank.size})" + for bank in self.banks + ) + + +class BankPlacementError(Exception): + """The target memory map cannot hold the graph's planned buffers.""" + + +class BankedGreedy: + """Plan specs across a :class:`TargetMemoryMap`, honoring bank capacity. + + One pass, largest buffer first, offering each to the fastest bank that can take + it: a bank accepts if the buffer fits an existing shared object for free, or if + it has headroom for a new one. Running out of banks fails the export. Every bank + therefore receives a size-descending subsequence, which is the ordering + ``pick_shared_obj`` requires, so a single-bank map reproduces ``greedy`` exactly. + + Pins -- specs whose ``mem_id`` a custom pool pass already set -- go down first, + so nothing can take the space they need, and a pin that does not fit is an error + rather than a silent relocation. A pin to an undeclared ``mem_id`` is laid out + but not budgeted. + + Unpinned buffers then share storage with pins freely: every pin is already + placed, so admitting one cannot block a pin still to come. That phase uses + :func:`_reusable_object` rather than ``pick_shared_obj``, because pins and + unpinned buffers are each size-descending but their concatenation is not, and + that function asserts on the ordering it would then see. + + Buffers are sorted by their incoming size, as ``greedy`` sorts, and realigned to + whichever bank they are offered to. A bank's alignment must be a multiple of the + graph's, which is what keeps that realignment order-preserving. + """ + + def __init__( + self, + target: TargetMemoryMap, + allow_overlapping_allocations: bool = True, + ) -> None: + self.target = target + self.allow_overlapping_allocations = allow_overlapping_allocations + # Named so MemoryPlanningAlgorithmSuite can report which algo it picked. + self.__name__ = "banked_greedy" + + def __call__( + self, + alignment: int, + specs: Iterable[TensorSpec], + graph_module: torch.fx.GraphModule, + graph_signature: Optional[ExportGraphSignature] = None, + extra_padding: int = 0, + ) -> MemoryAlgoResult: + result = MemoryAlgoResult({}, []) + self._check_bank_alignments(alignment) + all_specs = list(specs) + + # In-place outputs share their base's storage, so they inherit its bank + # rather than being placed or counted against capacity. + planned = [spec for spec in all_specs if spec.inplace_base is None] + deferred_inplace = [spec for spec in all_specs if spec.inplace_base is not None] + for spec in deferred_inplace: + spec.realign(alignment) + + self._reject_non_cpu(planned) + + objects: Dict[int, List[SharedObject]] = {} + used: Dict[int, int] = {} + for mem_id in self.target.mem_ids(): + objects[mem_id] = [] + used[mem_id] = _submodule_reserved(graph_module, mem_id) + spec2obj: Dict[TensorSpec, SharedObject] = {} + + # Pins first: nothing else may take the space they need, and they have no + # fallback bank. extra_padding is charged only to banks that hold something, + # as greedy pads only the arenas it touched. + pinned = [spec for spec in planned if spec.mem_id is not None] + for spec in _packing_order(pinned): + mem_id = spec.mem_id + bank_objects = objects.setdefault(mem_id, []) + spec.realign(self._bank_alignment(mem_id, alignment)) + before = len(bank_objects) + sobj = pick_shared_obj( + bank_objects, spec, self.allow_overlapping_allocations + ) + spec2obj[spec] = sobj + result.spec_dict[spec] = SpecAllocResult(mem_id, 0, 0) + if len(bank_objects) == before: + continue # rode along inside an object already there + if not self.target.manages(mem_id): + # Not our pool: lay it out, but claim no budget over it. + continue + bank = self.target.bank(mem_id) + cost = spec.allocated_memory + ( + extra_padding if len(bank_objects) == 1 else 0 + ) + if used[mem_id] + cost > bank.size: + raise BankPlacementError( + f"A {spec.allocated_memory}-byte buffer (shape " + f"{list(spec.shape)}) is pinned to bank {bank.name} " + f"(mem_id={mem_id}), which has {bank.size - used[mem_id]} of " + f"{bank.size} bytes left. A pinned buffer has no fallback bank." + ) + used[mem_id] += cost + + for spec in _packing_order([s for s in planned if s.mem_id is None]): + mem_id, sobj = self._place_free( + spec, objects, used, alignment, extra_padding + ) + result.spec_dict[spec] = SpecAllocResult(mem_id, 0, 0) + spec2obj[spec] = sobj + + for spec in deferred_inplace: + result.spec_dict[spec] = SpecAllocResult(0, 0, 0) + _resolve_inplace_specs(deferred_inplace, spec2obj, result) + + result.bufsizes = _compute_total_sizes( + objects, graph_module, 0, result, len(spec2obj) + ) + for mem_id, bank_objects in objects.items(): + if bank_objects and mem_id < len(result.bufsizes): + result.bufsizes[mem_id] += extra_padding + self._check_capacity(result.bufsizes) + logging.debug(f"banked greedy returns bufsizes: {result.bufsizes}") + return result + + def _check_bank_alignments(self, alignment: int) -> None: + for bank in self.target.banks: + if bank.alignment is not None and bank.alignment % alignment: + raise BankPlacementError( + f"Bank {bank.name} declares alignment {bank.alignment}, which is " + f"not a multiple of the graph's alignment {alignment}. A bank may " + f"tighten alignment, but only to a multiple: otherwise buffers " + f"land at offsets the export did not promise, and the packing " + f"order stops being size-descending." + ) + + def _bank_alignment(self, mem_id: int, default: int) -> int: + if not self.target.manages(mem_id): + return default + return self.target.bank(mem_id).alignment or default + + def _reject_non_cpu(self, specs: Sequence[TensorSpec]) -> None: + """``apply_algo`` plans per device, which would apply one map to each.""" + offenders = {spec.device for spec in specs if spec.device != DeviceType.CPU} + if offenders: + names = ", ".join(sorted(d.name for d in offenders)) + raise BankPlacementError( + f"Banked memory planning supports CPU specs only, but this graph has " + f"specs on {names}. A target memory map describes one address space, " + f"and apply_algo plans each device separately, so one map cannot " + f"describe them all." + ) + + def _place_free( + self, + spec: TensorSpec, + objects: Dict[int, List[SharedObject]], + used: Dict[int, int], + alignment: int, + extra_padding: int, + ) -> Tuple[int, SharedObject]: + """Put an unpinned spec in the fastest declared bank that can take it.""" + for mem_id in self.target.mem_ids(): + spec.realign(self._bank_alignment(mem_id, alignment)) + reuse = _reusable_object( + objects[mem_id], spec, self.allow_overlapping_allocations + ) + if reuse is not None: + # Fits inside an existing object: the bank does not grow. + sobj, offset = reuse + sobj.first_used_index = min(sobj.first_used_index, spec.lifetime[0]) + sobj.last_used_index = max(sobj.last_used_index, spec.lifetime[1]) + sobj.allocations.append(AllocationSpec(offset, spec)) + return mem_id, sobj + cost = spec.allocated_memory + (extra_padding if not objects[mem_id] else 0) + if used[mem_id] + cost <= self.target.bank(mem_id).size: + used[mem_id] += cost + sobj = SharedObject( + len(objects[mem_id]), + -1, + spec.allocated_memory, + spec.lifetime[0], + spec.lifetime[1], + ) + sobj.allocations.append(AllocationSpec(0, spec)) + objects[mem_id].append(sobj) + return mem_id, sobj + + remaining = ", ".join( + f"{self.target.bank(m).name} {self.target.bank(m).size - used[m]} of " + f"{self.target.bank(m).size} free" + for m in self.target.mem_ids() + ) + raise BankPlacementError( + f"No declared bank can hold a {spec.allocated_memory}-byte buffer " + f"(shape {list(spec.shape)}): {remaining}. Give a bank more room, or " + f"declare another one. Undeclared pools are never used as spill " + f"targets. Target memory map: {self.target.describe()}" + ) + + def _check_capacity(self, bufsizes: Sequence[int]) -> None: + """Backstop. Unreachable when the incremental accounting is right.""" + for mem_id in self.target.mem_ids(): + bank = self.target.bank(mem_id) + size = bufsizes[mem_id] if mem_id < len(bufsizes) else 0 + if size > bank.size: + raise BankPlacementError( + f"Bank {bank.name} (mem_id={mem_id}) needs {size} bytes but is " + f"only {bank.size} bytes. Target memory map: " + f"{self.target.describe()}" + ) + + +def _packing_order(specs: Sequence[TensorSpec]) -> List[TensorSpec]: + """Size-descending, tie-broken exactly as ``exir.memory_planning.greedy``.""" + sorted_specs: List[TensorSpec] = [] + for spec in specs: + bisect.insort(sorted_specs, spec, key=lambda x: x.allocated_memory) + sorted_specs.reverse() + return sorted_specs + + +def _reusable_object( + objects: Sequence[SharedObject], + spec: TensorSpec, + allow_overlapping_allocations: bool, +) -> Optional[Tuple[SharedObject, int]]: + """The object and offset that hold spec without growing the bank, if any. + + Mirrors ``pick_shared_obj``'s two reuse paths, read-only, and *filters* on + ``sobj.size >= spec.allocated_memory`` where that function asserts on it. + Unpinned buffers are offered after the pins, so the combined sequence a bank + sees is not size-descending and an undersized object must be skipped rather + than tripping the assert. + """ + for sobj in objects: + if sobj.size >= spec.allocated_memory and _does_not_overlap(sobj, spec): + return sobj, 0 + if allow_overlapping_allocations: + for sobj in objects: + max_offset = _find_max_overlapping_allocations_offset(sobj, spec) + if max_offset > 0 and max_offset + spec.allocated_memory <= sobj.size: + return sobj, max_offset + return None + + +def _submodule_reserved(graph_module: torch.fx.GraphModule, mem_id: int) -> int: + """Bytes a control-flow submodule already reserved in this bank.""" + bufsizes = getattr(graph_module, "input_mem_buffer_sizes", None) + if not bufsizes or len(bufsizes) <= mem_id: + return 0 + return bufsizes[mem_id] + + +def banked_memory_planning_pass( + target: TargetMemoryMap, + alignment: int = ALIGNMENT, + allow_overlapping_allocations: bool = True, + **kwargs: object, +) -> MemoryPlanningPass: + """A ``MemoryPlanningPass`` that plans across ``target``'s banks. + + The planner is wrapped in ``MemoryPlanningAlgorithmSuite`` because that is + what writes the winning placement back onto each spec. Remaining keyword + arguments go to ``MemoryPlanningPass`` (``alloc_graph_input`` and friends). + + Pass it as the *only* algorithm, as this helper does. The suite picks whichever + algorithm reports the smallest ``sum(bufsizes)``, which values a byte of slow + memory exactly like a byte of fast memory and applies no capacity check of its + own; putting ``greedy`` alongside this planner would let an unvalidated + single-arena plan win and silently over-subscribe a bank. + + To combine this with a custom pool pass, subclass that pass as + ``docs/source/compiler-memory-planning.md`` describes and hand it this + algorithm; the tags it writes are honored as pins. + + Note that a ``mem_id`` of 1 is a no-op tag under ``greedy`` -- it names the + arena everything already lands in -- but here it is a pin to the first declared + bank, which gets first claim on that bank's capacity. A pass written against + ``greedy`` that tags buffers with 1 will constrain them to the fastest bank; drop + those tags, or declare a bank 1 that can hold them. + """ + if kwargs.get("share_mutable_buffers"): + # run_multimethod hardcodes shared state to arena 2 and + # _check_default_mem_ids requires every other buffer on arena 1, neither of + # which a multi-bank plan can satisfy. Supporting it means placing shared + # state one arena past the last planned one, which is a change to core + # memory planning and belongs in its own review. + raise ValueError( + "share_mutable_buffers is not yet supported with banked memory " + "planning: it reserves arena 2 for shared state and requires every " + "other buffer on arena 1." + ) + + planner = BankedGreedy( + target, allow_overlapping_allocations=allow_overlapping_allocations + ) + return MemoryPlanningPass( + memory_planning_algo=MemoryPlanningAlgorithmSuite(algo_list=[planner]), + alignment=alignment, + **kwargs, # pyre-ignore[6] + ) + + +def planned_mem_ids(graph_module: torch.fx.GraphModule) -> List[int]: + """The bank each planned buffer ended up in, read back off the graph.""" + seen: Set[int] = set() + mem_ids: List[int] = [] + for node in graph_module.graph.nodes: + for spec in get_node_tensor_specs(node): + if id(spec) in seen or spec.const or spec.mem_id is None: + continue + seen.add(id(spec)) + mem_ids.append(spec.mem_id) + return mem_ids + + +def format_placement_report( + target: TargetMemoryMap, + bufsizes: Sequence[int], + mem_ids: Iterable[int] = (), +) -> str: + """A per-bank summary of where a plan put things. + + Takes the arena sizes and the planned buffers' ``mem_id``s rather than reading + state off a planner, so it is correct for any one method of a multi-method + program. ``mem_ids`` comes from :func:`planned_mem_ids` when all you have is + the exported program. + """ + counts: Dict[int, int] = {} + for mem_id in mem_ids: + counts[mem_id] = counts.get(mem_id, 0) + 1 + + lines = [ + "Memory bank placement:", + f" {'bank':<12}{'mem_id':>7}{'used':>12}{'size':>12}{'util':>8}{'buffers':>9}", + ] + for mem_id in target.mem_ids(): + bank = target.bank(mem_id) + used = bufsizes[mem_id] if mem_id < len(bufsizes) else 0 + util = 100.0 * used / bank.size if bank.size else 0.0 + lines.append( + f" {bank.name:<12}{mem_id:>7}{used:>12}{bank.size:>12}" + f"{util:>7.1f}%{counts.get(mem_id, 0):>9}" + ) + total_used = sum( + bufsizes[mem_id] if mem_id < len(bufsizes) else 0 for mem_id in target.mem_ids() + ) + lines.append(f" total planned: {total_used} bytes over {len(target.banks)} banks") + return "\n".join(lines) diff --git a/exir/tests/targets.bzl b/exir/tests/targets.bzl index 1bfe6634b7f..4fce5dc0319 100644 --- a/exir/tests/targets.bzl +++ b/exir/tests/targets.bzl @@ -160,6 +160,27 @@ def define_common_targets(is_fbcode = False): ], ) + runtime.python_test( + name = "banked_memory_planning", + srcs = [ + "test_banked_memory_planning.py", + ], + preload_deps = [ + "//executorch/kernels/portable:custom_ops_generated_lib", + ], + deps = [ + "//caffe2:torch", + "//caffe2/functorch:functorch_src", + "//executorch/exir:banked_memory_planning", + "//executorch/exir:lib", + "//executorch/exir:memory_planning", + "//executorch/exir:pass_manager", + "//executorch/exir:tensor", + "//executorch/exir/passes:lib", + "//executorch/extension/pybindings:portable_lib", # @manual + ], + ) + runtime.python_test( name = "memory_planning", srcs = [ diff --git a/exir/tests/test_banked_memory_planning.py b/exir/tests/test_banked_memory_planning.py new file mode 100644 index 00000000000..dfe0e31f949 --- /dev/null +++ b/exir/tests/test_banked_memory_planning.py @@ -0,0 +1,863 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import unittest +from typing import Dict, List, Mapping, Sequence, Tuple + +import torch +from executorch.exir import ExecutorchBackendConfig, to_edge +from executorch.exir.banked_memory_planning import ( + Bank, + banked_memory_planning_pass, + BankedGreedy, + BankPlacementError, + format_placement_report, + TargetMemoryMap, +) +from executorch.exir.memory_planning import ( + collect_specs_from_nodes, + greedy, + materialize_buffer, + MemoryPlanningAlgorithmSuite, + pick_shared_obj, + update_all_tensors_lifetime, + Verifier, +) +from executorch.exir.pass_manager import PassManager +from executorch.exir.passes import MemoryPlanningPass, SpecPropPass, ToOutVarPass +from executorch.exir.schema import DeviceType +from executorch.exir.tensor import TensorSpec +from functorch.experimental.control_flow import map as torch_map +from torch.export import export +from torch.export.exported_program import ExportGraphSignature +from torch.fx import GraphModule + +try: + from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch_from_buffer, + ) + + _HAS_RUNTIME = True +except ImportError: + _HAS_RUNTIME = False + + +KiB = 1024 +MiB = 1024 * 1024 + + +def make_spec(nbytes: int, lifetime: Tuple[int, int]) -> TensorSpec: + """A uint8 spec of exactly nbytes, with an explicit lifetime.""" + spec = TensorSpec(dtype=torch.uint8, shape=[nbytes]) + spec.lifetime = [lifetime[0], lifetime[1]] + return spec + + +def empty_graph_module() -> GraphModule: + return GraphModule(torch.nn.Module(), torch.fx.Graph()) + + +class ToyModelForBankPlanning(torch.nn.Module): + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + o = a + for _ in range(10): + o = o * a + o = o + b + return o + + def get_random_inputs(self) -> Tuple[torch.Tensor, ...]: + return (torch.randn(10), torch.randn(10)) + + +def prepare_toy_model() -> Tuple[GraphModule, ExportGraphSignature]: + model = ToyModelForBankPlanning() + edge = to_edge(export(model, model.get_random_inputs(), strict=True)) + gm = edge.exported_program().graph_module + gs = edge.exported_program().graph_signature + gm = PassManager(passes=[SpecPropPass(), ToOutVarPass()])(gm).graph_module + # Give the alloc nodes their specs, as MemoryPlanningPass does before + # planning; without this the graph exposes only its boundary tensors. + MemoryPlanningPass()._set_alloc_node_spec(gm) + update_all_tensors_lifetime(gm, gs) + return gm, gs + + +def flat_map(size: int = 1 << 20) -> TargetMemoryMap: + return TargetMemoryMap([Bank(name="sram", size=size, mem_id=1)]) + + +def two_banks(fast: int, slow: int) -> TargetMemoryMap: + return TargetMemoryMap( + [Bank(name="fast", size=fast, mem_id=1), Bank(name="slow", size=slow, mem_id=2)] + ) + + +def plan( + planner: BankedGreedy, + specs: List[TensorSpec], + graph_module: GraphModule = None, + extra_padding: int = 0, +) -> Tuple[List[int], Dict[int, Tuple[int, int]]]: + result = planner( + 16, specs, graph_module or empty_graph_module(), None, extra_padding + ) + return result.bufsizes, { + id(spec): (alloc.mem_id, alloc.mem_offset) + for spec, alloc in result.spec_dict.items() + } + + +def packed_size(specs: Sequence[TensorSpec]) -> int: + """Pack specs the way the planner does and return the bytes used.""" + objects = [] + ordered = sorted(specs, key=lambda s: s.allocated_memory, reverse=True) + for spec in ordered: + pick_shared_obj(objects, spec, True) + return materialize_buffer(objects) + + +def assert_no_free_readmission( + test: unittest.TestCase, + target: TargetMemoryMap, + specs: Sequence[TensorSpec], + alloc: Mapping[int, Tuple[int, int]], +) -> None: + """No spec may sit in a slow bank when a faster one would have taken it free. + + This is the invariant that a rank-only eviction rule violates: evicting a + buffer that shares a shared object frees no bytes, so it gets stranded in + slow memory for nothing. + """ + by_bank: Dict[int, List[TensorSpec]] = {m: [] for m in target.mem_ids()} + for spec in specs: + by_bank[alloc[id(spec)][0]].append(spec) + + for mem_id in target.mem_ids(): + for spec in by_bank[mem_id]: + for faster in range(1, mem_id): + residents = by_bank[faster] + before = packed_size(residents) + after = packed_size(list(residents) + [spec]) + if after == before and after <= target.bank(faster).size: + test.fail( + f"spec of {spec.allocated_memory} B sits in " + f"{target.bank(mem_id).name} but would ride free in " + f"{target.bank(faster).name} " + f"(packed size unchanged at {before})" + ) + + +class TestTargetMemoryMap(unittest.TestCase): + def test_banks_are_addressed_by_their_declared_mem_id(self) -> None: + target = two_banks(1 * KiB, 8 * KiB) + self.assertEqual(target.mem_id("fast"), 1) + self.assertEqual(target.mem_id("slow"), 2) + self.assertEqual(target.bank(1).name, "fast") + self.assertEqual(target.mem_ids(), [1, 2]) + + def test_declaration_order_is_the_preference_order(self) -> None: + target = two_banks(64 * KiB, 64 * KiB) + specs = [make_spec(256, (0, 4))] + _, alloc = plan(BankedGreedy(target), specs) + self.assertEqual(alloc[id(specs[0])][0], 1) + + def test_duplicate_bank_names_rejected(self) -> None: + with self.assertRaises(ValueError): + TargetMemoryMap( + [ + Bank(name="sram", size=16, mem_id=1), + Bank(name="sram", size=16, mem_id=2), + ] + ) + + def test_empty_map_rejected(self) -> None: + with self.assertRaises(ValueError): + TargetMemoryMap([]) + + def test_non_positive_bank_size_rejected(self) -> None: + with self.assertRaises(ValueError): + TargetMemoryMap([Bank(name="sram", size=0, mem_id=1)]) + + +class TestSingleBankEquivalence(unittest.TestCase): + """A flat map must plan bit-identically to exir's greedy.""" + + def _fresh(self) -> List[TensorSpec]: + return [ + make_spec(1024, (0, 4)), + make_spec(512, (1, 6)), + make_spec(2048, (5, 9)), + make_spec(256, (7, 12)), + make_spec(512, (11, 15)), + ] + + def test_matches_greedy_on_synthetic_specs(self) -> None: + banked_specs = self._fresh() + banked_sizes, banked_alloc = plan(BankedGreedy(flat_map()), banked_specs) + + greedy_specs = self._fresh() + greedy_result = greedy(16, greedy_specs, empty_graph_module(), None, 0) + greedy_alloc = { + i: ( + greedy_result.spec_dict[spec].mem_id, + greedy_result.spec_dict[spec].mem_offset, + ) + for i, spec in enumerate(greedy_specs) + } + banked_by_index = { + i: banked_alloc[id(spec)] for i, spec in enumerate(banked_specs) + } + + self.assertEqual(banked_sizes, greedy_result.bufsizes) + self.assertEqual(banked_by_index, greedy_alloc) + + def test_matches_greedy_per_spec_on_a_real_graph(self) -> None: + """Compare every placement, not just the totals.""" + gm, gs = prepare_toy_model() + specs = list(collect_specs_from_nodes(gm.graph.nodes, gs, do_assertion=False)) + self.assertGreater(len(specs), 10) + + banked = BankedGreedy(flat_map())(16, specs, gm, gs, 0) + expected = greedy(16, specs, gm, gs, 0) + + self.assertEqual(banked.bufsizes, expected.bufsizes) + for spec in specs: + self.assertEqual( + ( + banked.spec_dict[spec].mem_id, + banked.spec_dict[spec].mem_offset, + banked.spec_dict[spec].mem_obj_id, + ), + ( + expected.spec_dict[spec].mem_id, + expected.spec_dict[spec].mem_offset, + expected.spec_dict[spec].mem_obj_id, + ), + ) + + def test_matches_greedy_at_alignment_32(self) -> None: + banked_specs = self._fresh() + greedy_specs = self._fresh() + banked = BankedGreedy(flat_map())( + 32, banked_specs, empty_graph_module(), None, 0 + ) + expected = greedy(32, greedy_specs, empty_graph_module(), None, 0) + self.assertEqual(banked.bufsizes, expected.bufsizes) + for got, want in zip(banked_specs, greedy_specs): + self.assertEqual( + banked.spec_dict[got].mem_offset, expected.spec_dict[want].mem_offset + ) + + def test_matches_greedy_at_a_coarse_alignment(self) -> None: + """Sort must key on the incoming size, as greedy does, not the realigned one. + + Realigning before sorting manufactures ties, whose insort order then differs + from greedy's. These sizes make that divergence visible; many do not. + """ + + def fresh() -> List[TensorSpec]: + sizes = [899, 534, 508, 696, 940, 910, 331] + return [make_spec(n, (i, i + 3)) for i, n in enumerate(sizes)] + + banked_specs, greedy_specs = fresh(), fresh() + banked = BankedGreedy(flat_map())( + 256, banked_specs, empty_graph_module(), None, 0 + ) + expected = greedy(256, greedy_specs, empty_graph_module(), None, 0) + self.assertEqual(banked.bufsizes, expected.bufsizes) + for got, want in zip(banked_specs, greedy_specs): + self.assertEqual( + (banked.spec_dict[got].mem_offset, banked.spec_dict[got].mem_obj_id), + ( + expected.spec_dict[want].mem_offset, + expected.spec_dict[want].mem_obj_id, + ), + ) + + def test_is_suite_compatible(self) -> None: + specs = [make_spec(1024, (0, 4)), make_spec(512, (1, 6))] + suite = MemoryPlanningAlgorithmSuite(algo_list=[BankedGreedy(flat_map())]) + bufsizes = suite(16, specs, empty_graph_module(), None, 0) + self.assertEqual(bufsizes[0], 0) + self.assertGreater(bufsizes[1], 0) + for spec in specs: + self.assertEqual(spec.mem_id, 1) + + +class TestPlacement(unittest.TestCase): + def test_reserves_buffer_index_zero_for_constants(self) -> None: + bufsizes, _ = plan( + BankedGreedy(two_banks(1 * KiB, 8 * KiB)), [make_spec(256, (0, 4))] + ) + self.assertEqual(len(bufsizes), 3) + self.assertEqual(bufsizes[0], 0) + + def test_places_everything_in_the_fastest_bank_when_it_fits(self) -> None: + specs = [make_spec(256, (0, 4)), make_spec(256, (1, 6))] + bufsizes, alloc = plan(BankedGreedy(two_banks(4 * KiB, 64 * KiB)), specs) + for spec in specs: + self.assertEqual(alloc[id(spec)][0], 1) + self.assertEqual(bufsizes[2], 0) + + def test_spills_to_the_next_bank_when_the_fast_bank_is_full(self) -> None: + specs = [make_spec(1024, (0, 10)) for _ in range(3)] + bufsizes, alloc = plan(BankedGreedy(two_banks(2 * KiB, 64 * KiB)), specs) + self.assertEqual(sorted(alloc[id(s)][0] for s in specs), [1, 1, 2]) + self.assertEqual(bufsizes[1], 2 * KiB) + self.assertEqual(bufsizes[2], 1 * KiB) + + def test_cascades_across_three_banks(self) -> None: + target = TargetMemoryMap( + [ + Bank(name="tcm", size=2 * KiB, mem_id=1), + Bank(name="sram", size=2 * KiB, mem_id=2), + Bank(name="dram", size=64 * KiB, mem_id=3), + ] + ) + specs = [make_spec(1024, (0, 10)) for _ in range(6)] + bufsizes, alloc = plan(BankedGreedy(target), specs) + self.assertEqual(sorted(alloc[id(s)][0] for s in specs), [1, 1, 2, 2, 3, 3]) + self.assertEqual(bufsizes[1], 2 * KiB) + self.assertEqual(bufsizes[2], 2 * KiB) + + def test_never_exceeds_any_bank_capacity(self) -> None: + target = two_banks(2 * KiB, 8 * KiB) + specs = [make_spec(512, (i, i + 20)) for i in range(16)] + bufsizes, _ = plan(BankedGreedy(target), specs) + for mem_id in target.mem_ids(): + self.assertLessEqual(bufsizes[mem_id], target.bank(mem_id).size) + + def test_preserves_lifetime_reuse_within_a_bank(self) -> None: + specs = [make_spec(1024, (0, 4)), make_spec(1024, (5, 9))] + bufsizes, alloc = plan(BankedGreedy(two_banks(2 * KiB, 8 * KiB)), specs) + self.assertEqual(alloc[id(specs[0])], alloc[id(specs[1])]) + self.assertEqual(bufsizes[1], 1024) + self.assertEqual(bufsizes[2], 0) + + def test_fails_when_no_bank_can_hold_a_spec(self) -> None: + with self.assertRaises(BankPlacementError) as caught: + plan( + BankedGreedy(two_banks(1 * KiB, 2 * KiB)), [make_spec(16 * KiB, (0, 4))] + ) + message = str(caught.exception) + self.assertIn("16384", message) + self.assertIn("fast", message) + self.assertIn("slow", message) + + def test_plan_is_deterministic(self) -> None: + def run() -> Tuple[List[int], List[int]]: + specs = [make_spec(256 * (i % 5 + 1), (i, i + 8)) for i in range(40)] + bufsizes, alloc = plan(BankedGreedy(two_banks(4 * KiB, 64 * KiB)), specs) + return bufsizes, [alloc[id(s)][0] for s in specs] + + self.assertEqual(run(), run()) + + +class TestFastBankIsFilled(unittest.TestCase): + """Buffers that would ride free in a faster bank must not land in a slow one.""" + + def test_does_not_strand_a_spec_that_would_ride_free_in_a_faster_bank(self) -> None: + target = two_banks(16 * KiB, 64 * MiB) + specs = [make_spec(4 * KiB + 16 * i, (i, i + 3)) for i in range(120)] + _, alloc = plan(BankedGreedy(target), specs) + assert_no_free_readmission(self, target, specs, alloc) + + def test_short_lifetimes_let_many_specs_share_the_fast_bank(self) -> None: + target = two_banks(8 * KiB, 64 * MiB) + specs = [make_spec(4 * KiB, (i, i + 1)) for i in range(200)] + _, alloc = plan(BankedGreedy(target), specs) + in_fast = sum(1 for s in specs if alloc[id(s)][0] == 1) + # Two 4 KiB objects in an 8 KiB bank, and lifetimes are disjoint in pairs, + # so most of the 200 specs should ride inside them. + self.assertGreater(in_fast, 100) + + def test_reports_the_bank_when_a_spec_fits_nowhere(self) -> None: + target = TargetMemoryMap([Bank(name="fast", size=2 * KiB, mem_id=1)]) + specs = [make_spec(1024, (0, 10)) for _ in range(3)] + with self.assertRaises(BankPlacementError) as caught: + plan(BankedGreedy(target), specs) + self.assertIn("fast", str(caught.exception)) + + def test_extra_padding_is_charged_only_to_banks_that_hold_something(self) -> None: + """greedy pads only the arenas it touched; an idle bank must stay at 0.""" + target = two_banks(4 * KiB, 64 * KiB) + bufsizes, _ = plan( + BankedGreedy(target), [make_spec(256, (0, 4))], extra_padding=64 + ) + self.assertEqual(bufsizes[1], 256 + 64) + self.assertEqual(bufsizes[2], 0) + + def test_padding_counts_against_bank_capacity(self) -> None: + """A buffer that fits a bank on its own can still be pushed out by padding.""" + target = two_banks(32, 64 * KiB) + bufsizes, _ = plan( + BankedGreedy(target), [make_spec(16, (0, 4))], extra_padding=64 + ) + # 16 B fits the 32 B bank; 16 B plus 64 B of padding does not. + self.assertEqual(bufsizes[1], 0) + self.assertEqual(bufsizes[2], 16 + 64) + + def test_respects_submodule_reserved_bytes(self) -> None: + target = two_banks(4 * KiB, 64 * KiB) + gm = empty_graph_module() + # Index 0 is the constants slot; bank 1 already has 3 KiB reserved by a + # control-flow submodule, leaving room for only one 1 KiB spec. + gm.input_mem_buffer_sizes = [0, 3 * KiB, 0] + specs = [make_spec(1024, (0, 10)) for _ in range(3)] + bufsizes, alloc = plan(BankedGreedy(target), specs, graph_module=gm) + self.assertLessEqual(bufsizes[1], 4 * KiB) + self.assertEqual(sorted(alloc[id(s)][0] for s in specs), [1, 2, 2]) + + +class TestPreAssignedMemIds(unittest.TestCase): + """A custom pool pass may run ahead of the planner and pin specs.""" + + def test_pinned_spec_stays_in_its_bank(self) -> None: + target = two_banks(4 * KiB, 64 * KiB) + pinned = make_spec(256, (0, 10)) + pinned.mem_id = 2 # a custom pool pass put this in the slow bank + free = make_spec(256, (0, 10)) + + _, alloc = plan(BankedGreedy(target), [pinned, free]) + self.assertEqual(alloc[id(pinned)][0], 2) + self.assertEqual(alloc[id(free)][0], 1) + + def test_free_specs_cascade_around_a_pin(self) -> None: + target = two_banks(2 * KiB, 64 * KiB) + pinned = make_spec(1024, (0, 10)) + pinned.mem_id = 1 + free = [make_spec(1024, (0, 10)) for _ in range(3)] + + _, alloc = plan(BankedGreedy(target), [pinned] + free) + self.assertEqual(alloc[id(pinned)][0], 1) + self.assertEqual(sorted(alloc[id(s)][0] for s in free), [1, 2, 2]) + + def test_capacity_is_reserved_so_a_pin_is_never_squeezed_out(self) -> None: + """The pin is smallest, so a first-come rule would fill the bank first.""" + target = two_banks(2 * KiB, 64 * KiB) + pinned = make_spec(256, (0, 10)) + pinned.mem_id = 1 + free = [make_spec(1024, (0, 10)) for _ in range(4)] + + _, alloc = plan(BankedGreedy(target), free + [pinned]) + self.assertEqual(alloc[id(pinned)][0], 1) + # 2 KiB less the 256 B reservation leaves room for exactly one 1 KiB spec. + self.assertEqual(sorted(alloc[id(s)][0] for s in free), [1, 2, 2, 2]) + + def test_pinned_and_unpinned_buffers_share_storage(self) -> None: + """A pin must not wall off its bytes from disjoint unpinned buffers.""" + target = two_banks(4 * KiB, 64 * KiB) + pinned = make_spec(1024, (0, 4)) + pinned.mem_id = 1 + free = make_spec(1024, (5, 9)) # disjoint lifetime: can share + + bufsizes, alloc = plan(BankedGreedy(target), [pinned, free]) + self.assertEqual(alloc[id(free)], alloc[id(pinned)]) + self.assertEqual(bufsizes[1], 1024) + self.assertEqual(bufsizes[2], 0) + + def test_unpinned_buffer_larger_than_a_pin_takes_its_own_object(self) -> None: + """Placing pins first costs bytes when a pin is the smaller buffer. + + A shared object's size is fixed by its first occupant, so the 64 B pin's + object cannot host the 512 B buffer. greedy, which sorts purely by size, + would place the 512 B buffer first and let the pin ride inside it for 512 + total. This is the price of giving pins first claim on capacity, and it is + an ordering effect -- not a restriction on sharing. + """ + target = two_banks(4 * KiB, 64 * KiB) + pinned = make_spec(64, (0, 4)) + pinned.mem_id = 1 + free = make_spec(512, (5, 9)) + + bufsizes, alloc = plan(BankedGreedy(target), [pinned, free]) + self.assertEqual(alloc[id(free)][0], 1) + self.assertEqual(bufsizes[1], 576) # 64 + 512, vs greedy's 512 + + def test_pinned_graph_matches_greedy_when_one_bank_is_declared(self) -> None: + """Sharing across the pin boundary is what makes this hold.""" + + def fresh(): + specs = [ + make_spec(1024, (0, 4)), + make_spec(512, (5, 9)), + make_spec(256, (10, 14)), + make_spec(2048, (0, 14)), + ] + for spec in specs[:2]: + spec.mem_id = 1 + return specs + + banked = BankedGreedy(flat_map())(16, fresh(), empty_graph_module(), None, 0) + expected = greedy(16, fresh(), empty_graph_module(), None, 0) + self.assertEqual(banked.bufsizes, expected.bufsizes) + + def test_a_pin_that_rides_free_is_not_charged_capacity(self) -> None: + """Two pins with disjoint lifetimes share one object, so one bank's worth fits.""" + target = two_banks(1024, 64 * KiB) + first, second = make_spec(1024, (0, 4)), make_spec(1024, (5, 9)) + first.mem_id = second.mem_id = 1 + + bufsizes, alloc = plan(BankedGreedy(target), [first, second]) + self.assertEqual(alloc[id(second)], alloc[id(first)]) + self.assertEqual(bufsizes[1], 1024) + + def test_pin_larger_than_its_bank_is_an_error(self) -> None: + target = two_banks(1 * KiB, 64 * KiB) + pinned = make_spec(4 * KiB, (0, 10)) + pinned.mem_id = 1 + with self.assertRaises(BankPlacementError) as caught: + plan(BankedGreedy(target), [pinned]) + message = str(caught.exception) + self.assertIn("4096", message) + self.assertIn("fast", message) + + def test_pins_that_collectively_overflow_are_an_error(self) -> None: + target = two_banks(2 * KiB, 64 * KiB) + pins = [] + for _ in range(3): + spec = make_spec(1024, (0, 10)) + spec.mem_id = 1 + pins.append(spec) + with self.assertRaises(BankPlacementError) as caught: + plan(BankedGreedy(target), pins) + message = str(caught.exception) + self.assertIn("pinned to bank fast", message) + self.assertIn("no fallback bank", message) + + def test_pin_outside_the_map_is_laid_out_but_not_budgeted(self) -> None: + """An undeclared pool belongs to the pass that wrote it.""" + target = two_banks(4 * KiB, 64 * KiB) + pinned = make_spec(256, (0, 10)) + pinned.mem_id = 7 + free = make_spec(256, (0, 10)) + + bufsizes, alloc = plan(BankedGreedy(target), [pinned, free]) + self.assertEqual(alloc[id(pinned)][0], 7) + self.assertEqual(alloc[id(free)][0], 1) + self.assertEqual(bufsizes[7], 256) + + def test_undeclared_pool_never_receives_unpinned_buffers(self) -> None: + """The reason to leave a semantically special region out of the map.""" + target = TargetMemoryMap([Bank(name="fast", size=2 * KiB, mem_id=1)]) + dma = make_spec(256, (0, 10)) + dma.mem_id = 4 # a DMA-visible pool the planner must not fill + free = [make_spec(1024, (0, 10)) for _ in range(2)] + + with self.assertRaises(BankPlacementError): + # The third buffer has nowhere to go: bank 1 is full and pool 4 is + # not a spill target. It must fail rather than land in the DMA pool. + plan(BankedGreedy(target), [dma] + free + [make_spec(1024, (0, 10))]) + + _, alloc = plan(BankedGreedy(target), [dma] + free) + self.assertEqual(alloc[id(dma)][0], 4) + for spec in free: + self.assertEqual(alloc[id(spec)][0], 1) + + def test_custom_pool_pass_composes_end_to_end(self) -> None: + """The documented tagging pattern, planned by the banked algorithm.""" + + class TagAddsToSlowBank(MemoryPlanningPass): + def run(self, graph_module, graph_signature=None): + for subgm in graph_module.modules(): + if not isinstance(subgm, GraphModule): + continue + for node in subgm.graph.nodes: + if node.op == "call_function" and "add" in str(node.target): + spec = node.meta.get("spec") + if isinstance(spec, TensorSpec): + spec.mem_id = 2 + return super().run(graph_module, graph_signature) + + class AddNet(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(x + x) + 1.0 + + target = two_banks(64 * KiB, 1024 * KiB) + memory_pass = TagAddsToSlowBank( + memory_planning_algo=MemoryPlanningAlgorithmSuite( + algo_list=[BankedGreedy(target)] + ) + ) + edge = to_edge(export(AddNet(), (torch.randn(1, 64),), strict=True)) + program = edge.to_executorch( + ExecutorchBackendConfig(memory_planning_pass=memory_pass) + ) + bufsizes = list( + program.executorch_program.execution_plan[0].non_const_buffer_sizes + ) + self.assertEqual(len(bufsizes), 3) + # Tagged adds landed in bank 2; everything else in bank 1. + self.assertGreater(bufsizes[2], 0) + self.assertGreater(bufsizes[1], 0) + + +class TestInPlaceSpecs(unittest.TestCase): + def test_inplace_spec_shares_its_base_placement(self) -> None: + base = make_spec(1024, (0, 4)) + aliased = make_spec(1024, (4, 8)) + aliased.inplace_base = base + + _, alloc = plan(BankedGreedy(two_banks(4 * KiB, 64 * KiB)), [base, aliased]) + self.assertEqual(alloc[id(aliased)], alloc[id(base)]) + + def test_inplace_spec_follows_its_base_into_the_slow_bank(self) -> None: + target = two_banks(2 * KiB, 64 * KiB) + big = [make_spec(1024, (0, 10)) for _ in range(2)] + base = make_spec(1024, (0, 10)) + aliased = make_spec(1024, (0, 10)) + aliased.inplace_base = base + + _, alloc = plan(BankedGreedy(target), big + [base, aliased]) + self.assertEqual(alloc[id(aliased)], alloc[id(base)]) + + def test_inplace_spec_follows_a_base_that_spilled(self) -> None: + """The base is forced to the slow bank; the alias must go with it.""" + target = two_banks(2 * KiB, 64 * KiB) + fillers = [make_spec(1024, (0, 20)) for _ in range(2)] + base = make_spec(512, (0, 20)) # smaller, so it is placed last and spills + aliased = make_spec(512, (0, 20)) + aliased.inplace_base = base + + _, alloc = plan(BankedGreedy(target), fillers + [base, aliased]) + self.assertEqual(alloc[id(base)][0], 2) + self.assertEqual(alloc[id(aliased)], alloc[id(base)]) + + def test_inplace_spec_does_not_double_count_against_capacity(self) -> None: + base = make_spec(1024, (0, 10)) + aliased = make_spec(1024, (0, 10)) + aliased.inplace_base = base + bufsizes, _ = plan(BankedGreedy(two_banks(1 * KiB, 64 * KiB)), [base, aliased]) + self.assertEqual(bufsizes[1], 1024) + self.assertEqual(bufsizes[2], 0) + + +class TestPlanningPassIntegration(unittest.TestCase): + class Net(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = torch.nn.Linear(64, 64) + self.fc2 = torch.nn.Linear(64, 64) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.fc2(torch.relu(self.fc1(x))) + + def _plan(self, target: TargetMemoryMap): + edge = to_edge(export(self.Net(), (torch.randn(1, 64),), strict=True)) + program = edge.to_executorch( + ExecutorchBackendConfig( + memory_planning_pass=banked_memory_planning_pass(target) + ) + ) + sizes = program.executorch_program.execution_plan[0].non_const_buffer_sizes + return program, list(sizes) + + def test_export_produces_one_arena_per_bank(self) -> None: + _, bufsizes = self._plan(two_banks(64 * KiB, 1024 * KiB)) + self.assertEqual(len(bufsizes), 3) + self.assertEqual(bufsizes[0], 0) + + def test_export_respects_capacity_and_has_no_aliasing(self) -> None: + target = two_banks(8 * KiB, 1024 * KiB) + program, bufsizes = self._plan(target) + for mem_id in target.mem_ids(): + self.assertLessEqual(bufsizes[mem_id], target.bank(mem_id).size) + gm = program.exported_program().graph_module + Verifier(gm, True, True, True, None).verify_storage_reuse() + + +class TestBankedProgramExecutes(unittest.TestCase): + """A banked plan must not just serialize -- it must run.""" + + class Net(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.ModuleList( + [torch.nn.Conv2d(8, 8, 3, padding=1) for _ in range(4)] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for conv in self.conv: + x = torch.relu(conv(x)) + return x + + @unittest.skipUnless(_HAS_RUNTIME, "portable_lib not built") + def test_spilled_plan_runs_and_matches_eager(self) -> None: + model, inputs = self.Net().eval(), (torch.randn(1, 8, 16, 16),) + # 8 KiB holds one 8192 B activation, so the plan must use both arenas. + target = two_banks(8 * KiB, 8 * MiB) + program = to_edge(export(model, inputs, strict=True)).to_executorch( + ExecutorchBackendConfig( + memory_planning_pass=banked_memory_planning_pass(target) + ) + ) + arenas = list( + program.executorch_program.execution_plan[0].non_const_buffer_sizes + ) + self.assertEqual(arenas, [0, 8 * KiB, 8 * KiB]) + + runtime = _load_for_executorch_from_buffer(program.buffer) + got = runtime.forward(list(inputs))[0] + torch.testing.assert_close(got, model(*inputs)) + + +class TestControlFlow(unittest.TestCase): + class MapNet(torch.nn.Module): + def forward(self, xs: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + def body(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x + y + + return torch_map(body, xs, y) + + def test_submodule_arenas_are_reserved_and_match_greedy(self) -> None: + """apply_algo re-enters the planner per submodule; the parent must reserve.""" + inputs = (torch.randn(4, 256), torch.randn(256)) + target = two_banks(64 * KiB, 1024 * KiB) + + def sizes(memory_pass): + edge = to_edge(export(self.MapNet(), inputs, strict=True)) + program = edge.to_executorch( + ExecutorchBackendConfig(memory_planning_pass=memory_pass) + ) + return list( + program.executorch_program.execution_plan[0].non_const_buffer_sizes + ) + + flat = sizes( + MemoryPlanningPass( + memory_planning_algo=MemoryPlanningAlgorithmSuite(algo_list=[greedy]) + ) + ) + banked = sizes(banked_memory_planning_pass(target)) + self.assertEqual(banked[1], flat[1]) + self.assertLessEqual(banked[1], 64 * KiB) + + +class TestPerBankAlignment(unittest.TestCase): + def test_bank_alignment_overrides_the_global_alignment(self) -> None: + strict = TargetMemoryMap( + [Bank(name="tcm", size=4 * KiB, mem_id=1, alignment=32)] + ) + bufsizes, _ = plan(BankedGreedy(strict), [make_spec(40, (0, 4))]) + self.assertEqual(bufsizes[1], 64) # 40 -> 64 at alignment 32 + + def test_bank_without_alignment_inherits_the_global_one(self) -> None: + bufsizes, _ = plan(BankedGreedy(flat_map()), [make_spec(40, (0, 4))]) + self.assertEqual(bufsizes[1], 48) # 40 -> 48 at alignment 16 + + def test_each_bank_applies_its_own_alignment(self) -> None: + target = TargetMemoryMap( + [ + Bank(name="tcm", size=64, mem_id=1, alignment=32), + Bank(name="sram", size=64 * KiB, mem_id=2, alignment=16), + ] + ) + first, second = make_spec(40, (0, 10)), make_spec(40, (0, 10)) + bufsizes, alloc = plan(BankedGreedy(target), [first, second]) + # tcm rounds 40 up to 64 and is then full; sram takes the other at 48. + self.assertEqual(sorted(alloc[id(s)][0] for s in (first, second)), [1, 2]) + self.assertEqual(bufsizes[1], 64) + self.assertEqual(bufsizes[2], 48) + + def test_bank_alignment_must_be_a_multiple_of_the_graph_alignment(self) -> None: + """Otherwise ties invert and pick_shared_obj trips a bare internal assert.""" + for bad in (4, 24): # looser, and stricter but not a multiple + target = TargetMemoryMap( + [Bank(name="tcm", size=4 * KiB, mem_id=1, alignment=bad)] + ) + with self.assertRaises(BankPlacementError) as caught: + plan(BankedGreedy(target), [make_spec(32, (0, 4))]) + self.assertIn("multiple", str(caught.exception)) + + def test_pins_respect_per_bank_alignment(self) -> None: + target = TargetMemoryMap( + [Bank(name="tcm", size=4 * KiB, mem_id=1, alignment=32)] + ) + pinned = make_spec(40, (0, 4)) + pinned.mem_id = 1 + bufsizes, _ = plan(BankedGreedy(target), [pinned]) + self.assertEqual(bufsizes[1], 64) + + def test_non_positive_bank_alignment_rejected(self) -> None: + with self.assertRaises(ValueError): + TargetMemoryMap([Bank(name="x", size=64, mem_id=1, alignment=0)]) + + +class TestOverlappingReuse(unittest.TestCase): + """_reusable_object re-implements pick_shared_obj's second reuse rule.""" + + def test_offset_reuse_matches_greedy_per_spec(self) -> None: + """The third spec is admitted above an allocation whose lifetime overlaps.""" + + def fresh() -> List[TensorSpec]: + return [ + make_spec(256, (6, 7)), + make_spec(128, (0, 4)), + make_spec(128, (1, 1)), + ] + + banked_specs, greedy_specs = fresh(), fresh() + banked = BankedGreedy(flat_map())( + 16, banked_specs, empty_graph_module(), None, 0 + ) + expected = greedy(16, greedy_specs, empty_graph_module(), None, 0) + + self.assertEqual(banked.bufsizes, expected.bufsizes) + for got, want in zip(banked_specs, greedy_specs): + self.assertEqual( + (banked.spec_dict[got].mem_id, banked.spec_dict[got].mem_offset), + (expected.spec_dict[want].mem_id, expected.spec_dict[want].mem_offset), + ) + # Reuse must land at a non-zero offset, or this only exercises path one. + self.assertGreater(max(a.mem_offset for a in banked.spec_dict.values()), 0) + + +class TestUnsupportedConfigurations(unittest.TestCase): + def test_non_cpu_specs_are_rejected(self) -> None: + """A memory map describes one address space, not one per device.""" + spec = make_spec(256, (0, 4)) + spec.device = DeviceType.CUDA + with self.assertRaises(BankPlacementError) as caught: + plan(BankedGreedy(two_banks(4 * KiB, 64 * KiB)), [spec]) + self.assertIn("CUDA", str(caught.exception)) + + def test_allow_overlapping_allocations_reaches_the_planner(self) -> None: + """Vulkan-style configs disable overlapping; the factory must forward it.""" + memory_pass = banked_memory_planning_pass( + flat_map(), allow_overlapping_allocations=False + ) + planner = memory_pass.memory_planning_algo.algo_list[0] + self.assertFalse(planner.allow_overlapping_allocations) + + def test_share_mutable_buffers_is_rejected_for_now(self) -> None: + """Supporting it needs a core change to where shared state is placed.""" + with self.assertRaises(ValueError) as caught: + banked_memory_planning_pass(flat_map(), share_mutable_buffers=True) + self.assertIn("share_mutable_buffers", str(caught.exception)) + + +class TestPlacementReport(unittest.TestCase): + def test_report_lists_bytes_and_occupancy_per_bank(self) -> None: + target = two_banks(2 * KiB, 64 * KiB) + specs = [make_spec(1024, (0, 10)) for _ in range(3)] + planner = BankedGreedy(target) + result = planner(16, specs, empty_graph_module(), None, 0) + report = format_placement_report( + target, + result.bufsizes, + [alloc.mem_id for alloc in result.spec_dict.values()], + ) + + self.assertIn("fast", report) + self.assertIn("slow", report) + self.assertIn("2048", report) + self.assertIn("100.0%", report) + + +if __name__ == "__main__": + unittest.main()