Repository navigation
Commit b584d82
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
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
185 | 185 | | |
186 | 186 | | |
187 | 187 | | |
| 188 | + | |
| 189 | + | |
| 190 | + | |
| 191 | + | |
| 192 | + | |
| 193 | + | |
| 194 | + | |
| 195 | + | |
| 196 | + | |
| 197 | + | |
| 198 | + | |
| 199 | + | |
| 200 | + | |
188 | 201 | | |
189 | 202 | | |
190 | 203 | | |
| |||
0 commit comments