Add option in memory planning to put shared state on same location across entry points #14230

JacobSzwejbka · 2025-09-11T21:27:35Z

Summary: API that lets you place the same state tensor on the same id and offset across entry points. Lets you have get and set state more natively in the runtime if the underlying arenas are the same.

Differential Revision: D82250153

pytorch-bot · 2025-09-11T21:27:38Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 1 Unrelated Failure

As of commit 12ecb5a with merge base e31cef6 ():

NEW FAILURES - The following jobs have failed:

pull / android / run-emulator (gh)
The process '/usr/bin/sh' failed with exit code 255
pull / test-eval_llama-wikitext-linux / linux-job (gh)
RuntimeError: Command docker exec -t 7318ae011b508d1326e9ac2e7edd79aef25927375ef83b298c983f05aeeb7c94 /exec failed with exit code 1

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

pull / test-models-linux-basic (mv3, xnnpack-quantization-delegation, cmake, linux.arm64.2xlarge, execut... / 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.

facebook-github-bot · 2025-09-11T21:27:45Z

This pull request was exported from Phabricator. Differential Revision: D82250153

github-actions · 2025-09-11T21:28:25Z

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.

GregoryComer

I don't have full context on the memory planning logic, but this looks reasonable to me. This will be a very nice feature for encoder-decoder models. As an aside, do you think it will be feasible to default share_mutable_buffers to true in the future?

larryliu0820 · 2025-09-12T19:20:10Z

exir/passes/memory_planning_pass.py

+    return specs[0]
+
+def _insert_mutable_buffer_specs(state: "_MemoryPlanningState", gm: torch.fx.GraphModule, gs: ExportGraphSignature):
+    for node in gm.graph.nodes:


Would be good to have some docstring here

larryliu0820 · 2025-09-12T19:22:01Z

exir/passes/memory_planning_pass.py

        return PassResult(graph_module, True)
+
+    def run_multimethod(self):
+        "Resolve any memory planning done across entry points"


This should be a docstring right?

Suggested change

"Resolve any memory planning done across entry points"

"""Resolve any memory planning done across entry points"""

larryliu0820 · 2025-09-12T19:27:28Z

exir/passes/memory_planning_pass.py

+            assert(fqn)
+            spec = _get_spec_from_node(node)
+            if getattr(spec, 'mem_id', None) is not None or getattr(spec, 'mem_offset', None) is not None:
+                raise ValueError("Cannot share mutable buffers if they already have a mem_id or mem_offset assigned")


Should this be an exception or just a warning?

I think Id rather start with an exception and relax it later if we have a use case that needs it. If a mem_id or mem_offset is already present that means someone is running a custom memory plan before this. This already probably doesn't compose well in that scenario because we just place the buffers on mem_id=2 and assert everything else is on 1 or not planned.

…ross entry points (pytorch#14230) Summary: API that lets you place the same state tensor on the same id and offset across entry points. Lets you have get and set state more natively in the runtime if the underlying arenas are the same. Reviewed By: GregoryComer Differential Revision: D82250153

facebook-github-bot · 2025-09-12T19:56:53Z