Skip to content

Commit b584d82

Browse files
authored
Add a capacity-aware memory planner over a target memory map (pytorch#21821)
### Summary 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, which is asserted per-spec at three alignments. 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 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. Not covered here: the runtime still backs every arena from one pool, so per-bank regions in a runner are a follow-up, as is share_mutable_buffers, which reserves arena 2. Authored with assistance from Claude Code.
1 parent b55d8fc commit b584d82

4 files changed

Lines changed: 1419 additions & 0 deletions

File tree

‎exir/BUCK‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,19 @@ fbcode_target(_kind = runtime.python_library,
185185
],
186186
)
187187

188+
fbcode_target(_kind = runtime.python_library,
189+
name = "banked_memory_planning",
190+
srcs = [
191+
"banked_memory_planning.py",
192+
],
193+
deps = [
194+
":memory_planning",
195+
":tensor",
196+
"//caffe2:torch",
197+
"//executorch/exir/passes:memory_planning_pass",
198+
],
199+
)
200+
188201
fbcode_target(_kind = runtime.python_library,
189202
name = "common",
190203
srcs = [

0 commit comments

Comments
 (0)