Add a capacity-aware memory planner over a target memory map

[rascani]

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.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21821

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 26aed9f with merge base a56af1c ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.