Let delegate backends declare memory-planned scratch buffers

[rascani]

Lets a delegate backend declare scratch memory that the AOT memory planner allocates in the arena, instead of taking it from the runtime temp allocator at execute() time.

Why

A delegate that needs scratch today pays two costs. The size is invisible: an integrator has to run the model, watch the temp allocator's high-water mark, and hard-code the result, because nothing in the program says how much the delegate needs. And the temp allocator is a single pool, so a mixed CPU/NPU design has to back it with whatever memory the NPU's large scratch requires — a system that would put NPU scratch in DDR while keeping CPU kernel temporaries in fast TCM cannot, and the CPU kernels pay for it.

Making the requirement part of the memory plan addresses both: it becomes a property of the model rather than a number found by measurement, and because planned buffers carry a mem_id it can be pinned to one pool while the CPU's temporaries stay in another.

Reuse against unrelated planned tensors is a bonus rather than the point — the temp allocator is already reset every instruction, so it too gives max-not-sum across delegate calls.

Op-library backends already work this way: an out-variant kernel declares scratch through get_scratch_metas, gets a memory.alloc node, and the planner allocates it like any other intermediate. This gives delegates the same thing.

What a backend writes

return PreprocessResult(
    processed_bytes=compiled.blob,
    scratch_specs=[DelegateScratchSpec(nbytes=compiled.scratch_bytes)],
)

preprocess() runs per partition, after that partition is compiled, so the size is known. InsertDelegateScratchPass turns each spec into a memory.alloc node appended to that delegate call's arguments.

The runtime contract

Scratch arrives as the trailing arguments of the delegate call, after the inputs and the outputs, and there is no runtime or program-schema change. The runtime attaches no roles to delegate arguments at all — it hands over a flat list, and every backend already derives its input and output counts from its own blob (Ethos-U reads handles.inputs->count; NXP reads cfg->numInputs). Scratch is the same kind of positional slot, and putting it last means existing input and output indices are unchanged.

Recommended adoption pattern, correct in both directions of version skew:

if (args.size() > n_inputs + n_outputs) {
  // planned scratch is present
} else {
  // legacy .pte: fall back to get_temp_allocator()
}

Design notes

• The count comes from the graph, not the declaration. The pass records it on the delegate node; emit paths that skip the pass (LoweredBackendModule.program()) now raise instead of slicing a real input into the scratch slot.
• strip_delegate_scratch_pass runs first in the pipeline because to_executorch() writes its lowered graph back over the caller's edge program, so a second call would otherwise retrace a delegate call with arguments its signature does not have.
• mem_id should come from a compile spec, not be hardcoded. Which id means which pool is a property of the target; on the Arm runner the fast SRAM pool is mem_id=3 and is compiled out by default, and a non-default id is rejected by share_mutable_buffers and by enable_non_cpu_memory_planning.
• No alignment field, deliberately. Adding one required a second alignment attribute on TensorSpec (because realign() overwrites), which then broke the greedy planner's sort invariant, and it never aligned the buffer's start anyway — offsets are running sums. The planner's 16 bytes is what Ethos-U and NXP ask for; a backend needing more should check the pointer it receives.
• serde carries scratch_specs, since a round trip that dropped it emitted a program without the argument the backend's blob still expects.

Tests

exir/backend/test/test_delegate_scratch.py (12): argument placement, per-call sizing, disjoint-lifetime sharing, two buffers keeping declaration order and their own pools, scratch inside a torch.cond submodule, emitting without the pass erroring, survival through serde and deepcopy, to_executorch() twice with and without scratch, and declaration-time validation.

runtime/executor/test/backend_integration_test.cpp asserts what the backend actually receives: one extra argument beyond the inputs and outputs, a 1-D byte tensor of the declared size, a non-null pointer distinct from every input and output, and a planned arena large enough to hold it. test/models/export_delegated_program.py gained --scratch_bytes to produce the fixture. Note the CMake target for that test file is disabled upstream (TODO T191569140), so it runs under Buck rather than in the OSS CMake run.

Known gaps

• Inspectability is half delivered: the size is in the arena, but nothing labels it as scratch in a shipped .pte, so attributing arena growth to a delegate needs a tooling follow-up.
• devtools/inspector identifies delegate outputs by counting back from the end of the logged argument list, so an adopting backend breaks AOT-versus-runtime comparison until the Inspector skips the scratch tail. Scratch also shows up in etdump.
• Vulkan, CUDA, Metal, OpenVINO and MLX assert exact argument counts and must be updated before they adopt. They are unaffected by other backends' scratch, since the pass only touches delegates that declare it.
• No in-tree adopter yet. Ethos-U is the natural first one — arm_vela.py already computes the scratch size and discards it.

Authored with Claude Code.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 358cba6 with merge base 869b5ef ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

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.