Repository: pytorch/executorch
Pull request #10071: allow not memory planning mutable buffers
Author: login: JacobSzwejbka
Labels: name: CLA Signed
description: This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.
name: fb-exported
name: module: exir
description: Issues related to Export IR and the code under exir/
name: release notes: api
description: Changes to public facing apis (any interfaces, pybinded runtime methods, etc.)
Created: 2025-04-10T17:46:36Z | Merged: 2025-04-11T17:47:57Z

Summary:
Config option to not memory plan mutable buffers.

This when paired with a future runtime PR will allow users to retrieve buffers by name in the runtime and then set their dataptr.

Differential Revision: D72749868




cc @angelayi


Diff:
diff --git a/exir/emit/_emitter.py b/exir/emit/_emitter.py
index 9cc8a1e809c..fe18e49a623 100644
--- a/exir/emit/_emitter.py
+++ b/exir/emit/_emitter.py
@@ -1640,13 +1640,26 @@ def placeholder(  # noqa: C901
                 else:
                     spec.extra_tensor_info.fully_qualified_name = fqn
                     spec.extra_tensor_info.location = TensorDataLocation.EXTERNAL
-            if self.emitter_state.emit_mutable_buffer_names and is_mutable_buffer:
-                if spec.extra_tensor_info is None:
-                    spec.extra_tensor_info = ExtraTensorInfo(
-                        fully_qualified_name=fqn, location=TensorDataLocation.SEGMENT
+
+            if is_mutable_buffer:
+                # Emit names if we are supposed to.
+                if self.emitter_state.emit_mutable_buffer_names:
+                    if spec.extra_tensor_info is None:
+                        spec.extra_tensor_info = ExtraTensorInfo(
+                            fully_qualified_name=fqn,
+                            location=TensorDataLocation.SEGMENT,
+                        )
+                    else:
+                        spec.extra_tensor_info.fully_qualified_name = fqn
+                # if We aren't emitting the name then it needs to be memory planned.
+                elif spec.mem_id is None or spec.mem_offset is None:
+                    raise InternalError(
+                        self._emit_node_specific_error(
+                            self.node,
+                            # [2:] to remove the b_ prefix buffers get
+                            f'Mutable buffer "{target[2:]}" must have a memory id and offset if we are emitting it without a name. Please either memory plan your mutable buffers or call to_executorch with config=ExecutorchBackendConfig(emit_mutable_buffer_names=True)',
+                        )
                     )
-                else:
-                    spec.extra_tensor_info.fully_qualified_name = fqn
 
             # From the fqn find the corresponding tensor
             real_tensor = None
diff --git a/exir/emit/test/test_emit.py b/exir/emit/test/test_emit.py
index a2df53e8cf3..186c5a402ab 100644
--- a/exir/emit/test/test_emit.py
+++ b/exir/emit/test/test_emit.py
@@ -1838,8 +1838,40 @@ def forward(self, x):
         ep = to_edge(ep)
         # Lower the graph to executorch.
         ep = ep.to_executorch(
-            config=ExecutorchBackendConfig(emit_mutable_buffer_names=True)
+            config=ExecutorchBackendConfig(
+                emit_mutable_buffer_names=True,
+                memory_planning_pass=MemoryPlanningPass(alloc_mutable_buffers=False),
+            )
         )
         for val in ep.executorch_program.execution_plan[0].values:
             if isinstance(val, Tensor) and val.extra_tensor_info:
                 self.assertEqual(val.extra_tensor_info.fully_qualified_name, "buffer")
+                self.assertEqual(val.allocation_info, None)
+
+    def test_emit_mutable_buffer_names_fails(self) -> None:
+        class Net(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(2, 2)
+                self.register_buffer("buffer", torch.zeros(1, 2))
+
+            def forward(self, x):
+                self.buffer.add_(1)
+                return self.linear(x) + self.buffer
+
+        net = Net()
+
+        ep = export(net, (torch.randn(1, 2),), strict=True)
+        # Lower the graph to edge dialect.
+        ep = to_edge(ep)
+        # Lower the graph to executorch.
+        # Must emit mutable buffer names if we don't allocate mutable buffers
+        with self.assertRaises(InternalError):
+            ep.to_executorch(
+                config=ExecutorchBackendConfig(
+                    emit_mutable_buffer_names=False,
+                    memory_planning_pass=MemoryPlanningPass(
+                        alloc_mutable_buffers=False
+                    ),
+                )
+            )
diff --git a/exir/memory_planning.py b/exir/memory_planning.py
index 17640a9f7aa..83598940882 100644
--- a/exir/memory_planning.py
+++ b/exir/memory_planning.py
@@ -44,12 +44,14 @@ def __init__(
         graph_module: torch.fx.GraphModule,
         alloc_graph_input: bool,
         alloc_graph_output: bool,
+        alloc_mutable_buffers: bool,
         graph_signature: Optional[ExportGraphSignature] = None,
     ) -> None:
         self.graph_module = graph_module
         self.graph_signature = graph_signature
         self.alloc_graph_input = alloc_graph_input
         self.alloc_graph_output = alloc_graph_output
+        self.alloc_mutable_buffers = alloc_mutable_buffers
 
     @classmethod
     def mem_obj_id_match(
@@ -149,6 +151,7 @@ def verify_storage_reuse(
                 ignore_const=True,
                 ignore_graph_input=not self.alloc_graph_input,
                 ignore_graph_output=not self.alloc_graph_output,
+                ignore_mutable_buffers=not self.alloc_mutable_buffers,
                 do_assertion=False,
                 ignore_out_var_node=False,
                 dedup=True,
@@ -374,6 +377,7 @@ def collect_specs_from_nodes(  # noqa: C901
     graph_signature: Optional[ExportGraphSignature] = None,
     ignore_graph_input: bool = False,
     ignore_graph_output: bool = False,
+    ignore_mutable_buffers: bool = False,
     ignore_const: bool = True,
     ignore_out_var_node: bool = True,
     dedup: bool = True,
@@ -414,6 +418,9 @@ def collect_specs_from_nodes(  # noqa: C901
         if _is_inplace_node(node):
             continue
 
+        if _is_mutable_buffer(node, graph_signature) and ignore_mutable_buffers:
+            continue
+
         if do_assertion:
             internal_assert(
                 node.op in ("placeholder", "output")
@@ -469,6 +476,7 @@ def update_all_tensors_lifetime(
     Set the lifetime for all the tensors encountered in the Fx graph.
     """
     specs = set()
+
     for node_idx, node in enumerate(graph_module.graph.nodes):
         for spec in collect_specs_from_nodes(
             filter_nodes(itertools.chain([node], node.args, node.kwargs.values())),
@@ -1053,6 +1061,7 @@ def apply_algo(
     graph_signature: Optional[ExportGraphSignature] = None,
     alloc_graph_input: bool = True,
     alloc_graph_output: bool = True,
+    alloc_mutable_buffers: bool = True,
 ) -> List[int]:
     """
     Recursively apply algo to graph_module and its submodules for control flow.
@@ -1065,12 +1074,10 @@ def apply_algo(
        storage with tensors in the outer module.
     TODO: make these optimizations once we have some baseline working.
     """
-
     # Extract the nodes and their lifespans from the graph_module
     # Difficult to just filter the list of specs returned by this due to
     # how we flag trainable weights.
     _ = update_all_tensors_lifetime(graph_module, graph_signature)
-
     # Filter specs based on alloc_graph_input and alloc_graph_output
     specs = collect_specs_from_nodes(
         graph_module.graph.nodes,
@@ -1078,6 +1085,7 @@ def apply_algo(
         do_assertion=False,
         ignore_graph_input=not alloc_graph_input,
         ignore_graph_output=not alloc_graph_output,
+        ignore_mutable_buffers=not alloc_mutable_buffers,
     )
 
     # Get extra padding for XNNPACK if needed
diff --git a/exir/passes/memory_planning_pass.py b/exir/passes/memory_planning_pass.py
index f4b3ad8a8a7..9bd4ab20bf5 100644
--- a/exir/passes/memory_planning_pass.py
+++ b/exir/passes/memory_planning_pass.py
@@ -44,6 +44,7 @@ def __init__(
         allow_lifetime_and_storage_overlap: bool = False,
         alloc_graph_input: bool = True,
         alloc_graph_output: bool = True,
+        alloc_mutable_buffers: bool = True,
         alignment: int = ALIGNMENT,
     ) -> None:
         r"""
@@ -54,10 +55,11 @@ def __init__(
         """
         if memory_planning_algo is None:
             memory_planning_algo = MemoryPlanningAlgorithmSuite()
-        self.memory_planning_algo = memory_planning_algo
+        self.memory_planning_algo: Callable[..., List[int]] = memory_planning_algo
         self.allow_lifetime_and_storage_overlap = allow_lifetime_and_storage_overlap
         self.alloc_graph_input = alloc_graph_input
         self.alloc_graph_output = alloc_graph_output
+        self.alloc_mutable_buffers = alloc_mutable_buffers
         self.alignment = alignment
 
     def _set_alloc_node_spec(self, graph_module: torch.fx.GraphModule) -> None:
@@ -124,13 +126,15 @@ def run(
         # customized fields. Using the graph_module object to convey information across
         # passes/stages is quite natural and avoid yet another 'context' data structure
         # to do the job.
+
         _ = apply_algo(
-            self.memory_planning_algo,  # pyre-ignore[6]
+            self.memory_planning_algo,
             graph_module,
             self.alignment,
             graph_signature,
             self.alloc_graph_input,
             self.alloc_graph_output,
+            self.alloc_mutable_buffers,
         )
 
         # TODO: make the verifier do the work recursively to handle
@@ -139,6 +143,7 @@ def run(
             graph_module,
             self.alloc_graph_input,
             self.alloc_graph_output,
+            self.alloc_mutable_buffers,
             graph_signature,
         )
 
diff --git a/exir/tests/test_memory_planning.py b/exir/tests/test_memory_planning.py
index 52986aaa04c..b87ae2dfb58 100644
--- a/exir/tests/test_memory_planning.py
+++ b/exir/tests/test_memory_planning.py
@@ -241,6 +241,7 @@ def maketest(
     use_functionalization: bool = True,
     alloc_graph_input: bool = True,
     alloc_graph_output: bool = True,
+    alloc_mutable_buffer: bool = True,
     has_unused_graph_input: bool = False,
 ) -> Callable[..., None]:
     # parameterized.expand is not compatible with maketest. I'll just loop thru
@@ -282,10 +283,17 @@ def wrapper(self: "TestMemoryPlanning") -> None:
             )(graph_module).graph_module
 
             self.verify_reuse(
-                graph_module, expect_reuse, alloc_graph_input, alloc_graph_output
+                graph_module,
+                expect_reuse,
+                alloc_graph_input,
+                alloc_graph_output,
+                alloc_mutable_buffer,
             )
             self.verify_graph_input_output(
-                graph_module, alloc_graph_input, alloc_graph_output
+                graph_module,
+                alloc_graph_input,
+                alloc_graph_output,
+                alloc_mutable_buffer,
             )
 
             self.verify_overlap_placeholders(has_unused_graph_input, graph_module)
@@ -306,6 +314,7 @@ def verify_reuse(
         expect_reuse: bool,
         alloc_graph_input: bool,
         alloc_graph_output: bool,
+        alloc_mutable_buffer: bool,
     ) -> None:
         r"""
         Do sanity check and verify tensor storage reuse.
@@ -321,6 +330,7 @@ def verify_reuse(
             graph_module,
             alloc_graph_input=alloc_graph_input,
             alloc_graph_output=alloc_graph_output,
+            alloc_mutable_buffers=alloc_mutable_buffer,
         ).verify_storage_reuse()
 
         print(f"num_reuse_pairs is {num_reuse_pairs}")
@@ -334,9 +344,10 @@ def verify_graph_input_output(
         graph_module: torch.fx.GraphModule,
         alloc_graph_input: bool,
         alloc_graph_output: bool,
+        alloc_mutable_buffers: bool,
     ) -> None:
         Verifier(
-            graph_module, alloc_graph_input, alloc_graph_output
+            graph_module, alloc_graph_input, alloc_graph_output, alloc_mutable_buffers
         ).verify_graph_input_output()
 
     def verify_overlap_placeholders(
@@ -404,13 +415,16 @@ def verify_overlap_placeholders(
     )
 
     def test_graph_input_output(self) -> None:
-        for alloc_graph_input, alloc_graph_output in itertools.product(
-            [True, False], [True, False]
-        ):
+        for (
+            alloc_graph_input,
+            alloc_graph_output,
+            alloc_mutable_buffers,
+        ) in itertools.product([True, False], [True, False], [True, False]):
             case = maketest(
                 ModelWithDifferentTensorSizes,
                 alloc_graph_input=alloc_graph_input,
                 alloc_graph_output=alloc_graph_output,
+                alloc_mutable_buffer=alloc_mutable_buffers,
             )
             case(self)
 
@@ -535,6 +549,7 @@ def test_multiple_pools(
             graph_module,
             alloc_graph_input=True,
             alloc_graph_output=True,
+            alloc_mutable_buffers=True,
         )
         verifier.verify_storage_reuse()
         verifier.verify_graph_input_output()


Reviews:
user: login: JacobSzwejbka
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T20:14:04Z
user: login: tarun292
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T00:02:42Z
user: login: JacobSzwejbka
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T00:56:41Z
user: login: JacobSzwejbka
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T00:57:58Z

Review comments:
diff_hunk: @@ -1640,13 +1640,26 @@ def placeholder(  # noqa: C901
                 else:
                     spec.extra_tensor_info.fully_qualified_name = fqn
                     spec.extra_tensor_info.location = TensorDataLocation.EXTERNAL
-            if self.emitter_state.emit_mutable_buffer_names and is_mutable_buffer:
-                if spec.extra_tensor_info is None:
-                    spec.extra_tensor_info = ExtraTensorInfo(
-                        fully_qualified_name=fqn, location=TensorDataLocation.SEGMENT
+
+            if is_mutable_buffer:
path: exir/emit/_emitter.py
user: login: JacobSzwejbka
body: This is the actual core logic change, the rest of the changes are mostly piping a flag around
created_at: 2025-04-10T20:14:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1838,8 +1838,40 @@ def forward(self, x):
         ep = to_edge(ep)
         # Lower the graph to executorch.
         ep = ep.to_executorch(
-            config=ExecutorchBackendConfig(emit_mutable_buffer_names=True)
+            config=ExecutorchBackendConfig(
path: exir/emit/test/test_emit.py
user: login: tarun292
body: Maybe also add a __post_init__ in the `ExecutorchBackendConfig` dataclass that asserts if `emit_mutable_buffer_names` is `False` and `alloc_mutable_buffers` is also False.
created_at: 2025-04-11T00:01:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1640,13 +1640,26 @@ def placeholder(  # noqa: C901
                 else:
                     spec.extra_tensor_info.fully_qualified_name = fqn
                     spec.extra_tensor_info.location = TensorDataLocation.EXTERNAL
-            if self.emitter_state.emit_mutable_buffer_names and is_mutable_buffer:
-                if spec.extra_tensor_info is None:
-                    spec.extra_tensor_info = ExtraTensorInfo(
-                        fully_qualified_name=fqn, location=TensorDataLocation.SEGMENT
+
+            if is_mutable_buffer:
+                # Emit names if we are supposed to.
+                if self.emitter_state.emit_mutable_buffer_names:
+                    if spec.extra_tensor_info is None:
path: exir/emit/_emitter.py
user: login: tarun292
body: A little confused by this as i haven't kept track of the tensor info changes, what does it mean if `extra_tensor_info` is `None` and if it isn't we overwrite the fqn again? Maybe add a small comment here too.
created_at: 2025-04-11T00:02:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1838,8 +1838,40 @@ def forward(self, x):
         ep = to_edge(ep)
         # Lower the graph to executorch.
         ep = ep.to_executorch(
-            config=ExecutorchBackendConfig(emit_mutable_buffer_names=True)
+            config=ExecutorchBackendConfig(
path: exir/emit/test/test_emit.py
user: login: JacobSzwejbka
body: I considered this but since the MemoryPlanningPass is user configureable there is no guarantee memory_planning_pass.alloc_mutable_buffers exists. So I just check the end result in the emitter.
created_at: 2025-04-11T00:56:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1640,13 +1640,26 @@ def placeholder(  # noqa: C901
                 else:
                     spec.extra_tensor_info.fully_qualified_name = fqn
                     spec.extra_tensor_info.location = TensorDataLocation.EXTERNAL
-            if self.emitter_state.emit_mutable_buffer_names and is_mutable_buffer:
-                if spec.extra_tensor_info is None:
-                    spec.extra_tensor_info = ExtraTensorInfo(
-                        fully_qualified_name=fqn, location=TensorDataLocation.SEGMENT
+
+            if is_mutable_buffer:
+                # Emit names if we are supposed to.
+                if self.emitter_state.emit_mutable_buffer_names:
+                    if spec.extra_tensor_info is None:
path: exir/emit/_emitter.py
user: login: JacobSzwejbka
body: None just means no one has created one, if its not none they might not have populated the fqn so I populate it. I guess I could check what it is before overwriting it but the name is unique so it should always be safe to do this. 

extra_tensor_info is where we put optional info in the flatbuffer to not regress the size of every tensor by too much in embedded cases that dont need it.
created_at: 2025-04-11T00:57:58Z
author_association: CONTRIBUTOR

PR comments:
user: login: pytorch-bot[bot]
created_at: 2025-04-10T17:46:39Z
body: <!-- drci-comment-start -->

## :link: Helpful Links
### :test_tube: See artifacts and rendered test results at [hud.pytorch.org/pr/pytorch/executorch/10071](https://hud.pytorch.org/pr/pytorch/executorch/10071)
* :page_facing_up: Preview [Python docs built from this PR](https://docs-preview.pytorch.org/pytorch/executorch/10071/index.html)

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


## :white_check_mark: No Failures
As of commit d8fa917acf260cc0d6f97afc26829f7b51d4ef64 with merge base d42676f93fc928341730bec7efc265a09b05fb1a (<sub><sub><img alt="image" width=70 src="https://img.shields.io/date/1744313667?label=&color=FFFFFF&style=flat-square"></sub></sub>):
:green_heart: Looks good so far! There are no failures yet. :green_heart:


This comment was automatically generated by Dr. CI and updates every 15 minutes.
<!-- drci-comment-end -->
author_association: NONE
user: login: facebook-github-bot
created_at: 2025-04-10T17:46:50Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D72749868](https://www.internalfb.com/diff/D72749868)
author_association: CONTRIBUTOR
user: login: facebook-github-bot
created_at: 2025-04-10T17:54:36Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D72749868](https://www.internalfb.com/diff/D72749868)
author_association: CONTRIBUTOR
user: login: facebook-github-bot
created_at: 2025-04-10T17:57:20Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D72749868](https://www.internalfb.com/diff/D72749868)
author_association: CONTRIBUTOR
user: login: facebook-github-bot
created_at: 2025-04-10T18:18:37Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D72749868](https://www.internalfb.com/diff/D72749868)
author_association: CONTRIBUTOR
user: login: facebook-github-bot
created_at: 2025-04-10T20:12:54Z
body: This pull request was **exported** from Phabricator. Differential Revision: [D72749868](https://www.internalfb.com/diff/D72749868)
author_association: CONTRIBUTOR