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()