Add create_mutable_buffer util #13813

digantdesai · 2025-08-29T15:26:11Z

Rationale:

* We want to support stateful custom ops inserted after to_edge in a non delegated cases like op-lib-backends i.e. cortex-m::cmsis_nn::linear

Details:

* Allows a pass to add a mutable buffer.
* Essentially redoes what export, to_edge does for a buffer registered in an nn.module.
* Also plays nice with to_executorch passes like conversion to in-place, and write-back for mutated buffers.
* Verifies above with tests, adds a HelperPass for someone looking to leverage this util as an example too.

To Test:

$ python -m unittest backends.transforms.test.test_create_mutable_buffer

pytorch-bot · 2025-08-29T15:26:14Z

🔗 Helpful Links

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

📄 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 c6bb918 with merge base 6c12956 ():
💚 Looks good so far! There are no failures yet. 💚

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

github-actions · 2025-08-29T15:26:46Z

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.

mergennachin

Add backends/transforms in https://github.com/pytorch/executorch/blob/main/pytest.ini

backends/transforms/utils.py

kimishpatel · 2025-09-03T14:59:20Z

backends/transforms/utils.py

+
+    if not user_input_indices_or_consts:
+        # If no user inputs or constants exist, insert at the end of all existing inputs
+        node_index = len(exp_program.graph_signature.input_specs)


unused in this if block

it is used outside the block

backends/transforms/utils.py

backends/transforms/test/test_create_mutable_buffer.py

kimishpatel · 2025-09-03T15:10:43Z

backends/transforms/test/test_create_mutable_buffer.py

+        )
+        return executorch_program_manager
+
+    def _test_eager(self, model, example_inputs, num_lifted_args=1):


dont follow what this is doing for eager?

eager might be a bad name, but the idea is to compare the ExecuTorch graph from two different origins for a mutable buffer (1) written in torch.nn.Module by the end user as in eager mode, and (2) inserted via this util + pass, and both ET graphs should be indistinguishable in every way.

kimishpatel · 2025-09-03T15:11:54Z

backends/transforms/test/test_create_mutable_buffer.py

+
+        self.compare(et_1, et_2, example_inputs)
+
+    def test_basic_with_param(self):


are you testing when you dont have any inputs in the graph? there is a path for that in your pass.

Not sure I follow. x is the normal input, and I am adding a param (beyond the test_basic module) to test if the signature update is correct.

kimishpatel

looks pretty good. left some comments

backends/transforms/test/test_create_mutable_buffer.py

backends/transforms/utils.py

JacobSzwejbka · 2025-09-03T16:38:16Z

backends/transforms/utils.py

+            )
+
+    # Find indices of user inputs and constant tensors
+    user_input_indices_or_consts = [


I feel like a loop would be a cleaner way of writing this through line 324.

Since you have validated the Spec.kind and the nodes earlier you could just loop over the input_specs and index into graph.nodes

refactored but still not as clean as I like.

JacobSzwejbka · 2025-09-03T16:41:35Z

@angelayi @avikchaudhuri @tugsbayasgalan

I think in general it would be a good value add from compiler if these sorts of primitive graph operations lived in core.

Things like node.is_mutable_buffer(), ep.add_buffer() ep.add_mutable_buffer() etc

JacobSzwejbka

.

mergennachin · 2025-09-05T13:45:46Z

@digantdesai see my previous comment #13813 (review)

digantdesai · 2025-09-08T16:35:35Z

@digantdesai see my previous comment #13813 (review)

separate PR?

JacobSzwejbka · 2025-09-08T16:48:44Z

backends/transforms/test/test_create_mutable_buffer.py

+            if modified:
+                graph_module = super().call(graph_module).graph_module
+                graph_module.graph.lint()
+                graph_module.graph.eliminate_dead_code()


this fn is technically unsafe to call if the graph is not functional.

JacobSzwejbka · 2025-09-08T16:49:45Z

backends/transforms/utils.py

+
+    example_node = nodes[0]
+    if isinstance(
+        example_node.meta["val"], (tuple, torch.fx.immutable_collections.immutable_list)


seems like a bit of a strange check. Is this robust?

Hmm just refactored in this PR.. but what else it can be when number of tensors > 1?

digantdesai

Thanks @JacobSzwejbka

backends/transforms/utils.py

digantdesai · 2025-09-08T17:07:10Z

backends/transforms/utils.py

+
+    example_node = nodes[0]
+    if isinstance(
+        example_node.meta["val"], (tuple, torch.fx.immutable_collections.immutable_list)


Hmm just refactored in this PR.. but what else it can be when number of tensors > 1?

Rationale: * We want to support stateful custom ops inserted after to_edge in a non delegated cases like op-lib-backends i.e. cortex-m::cmsis_nn::linear Details: * Allows a pass to add a mutable buffer. * Essentially redoes what export, to_edge does for a buffer registered in an nn.module. * Also plays nice with to_executorch passes like conversion to in-place, and write-back for mutated buffers. * Verifies above with tests, adds a HelperPass for someone looking to leverage this util as an example too. To Test: $ python -m unittest backends.transforms.test.test_create_mutable_buffer

digantdesai requested a review from JacobSzwejbka August 29, 2025 15:26

digantdesai requested a review from kimishpatel as a code owner August 29, 2025 15:26

meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 29, 2025

digantdesai force-pushed the mutable_buffers branch 2 times, most recently from dbbd892 to 311eba6 Compare September 2, 2025 21:25

mergennachin reviewed Sep 3, 2025

View reviewed changes

kimishpatel reviewed Sep 3, 2025

View reviewed changes

backends/transforms/utils.py Outdated Show resolved Hide resolved

kimishpatel reviewed Sep 3, 2025

View reviewed changes

backends/transforms/utils.py Outdated Show resolved Hide resolved

kimishpatel reviewed Sep 3, 2025

View reviewed changes

backends/transforms/test/test_create_mutable_buffer.py Outdated Show resolved Hide resolved

kimishpatel reviewed Sep 3, 2025

View reviewed changes

kimishpatel approved these changes Sep 3, 2025

View reviewed changes

JacobSzwejbka reviewed Sep 3, 2025

View reviewed changes

backends/transforms/test/test_create_mutable_buffer.py Outdated Show resolved Hide resolved

JacobSzwejbka reviewed Sep 3, 2025

View reviewed changes

backends/transforms/utils.py Outdated Show resolved Hide resolved

JacobSzwejbka reviewed Sep 3, 2025

View reviewed changes

backends/transforms/utils.py Show resolved Hide resolved

JacobSzwejbka reviewed Sep 3, 2025

View reviewed changes

JacobSzwejbka requested changes Sep 3, 2025

View reviewed changes

digantdesai force-pushed the mutable_buffers branch 2 times, most recently from 332edae to 918b70b Compare September 5, 2025 10:30

JacobSzwejbka reviewed Sep 8, 2025

View reviewed changes

JacobSzwejbka approved these changes Sep 8, 2025

View reviewed changes

digantdesai commented Sep 8, 2025

View reviewed changes

digantdesai force-pushed the mutable_buffers branch from 7f74232 to c6bb918 Compare September 8, 2025 17:15

digantdesai merged commit 9ce07da into main Sep 8, 2025
116 checks passed

digantdesai deleted the mutable_buffers branch September 8, 2025 19:19


		self.compare(et_1, et_2, example_inputs)

		def test_basic_with_param(self):

Add create_mutable_buffer util #13813

Add create_mutable_buffer util #13813

Uh oh!

Conversation

digantdesai commented Aug 29, 2025

Rationale:

Details:

To Test:

Uh oh!

pytorch-bot bot commented Aug 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ No Failures

Uh oh!

github-actions bot commented Aug 29, 2025

This PR needs a release notes: label

Uh oh!

mergennachin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

digantdesai Sep 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kimishpatel left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JacobSzwejbka commented Sep 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

JacobSzwejbka left a comment

Choose a reason for hiding this comment

Uh oh!

mergennachin commented Sep 5, 2025

Uh oh!

digantdesai commented Sep 8, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

digantdesai left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

pytorch-bot bot commented Aug 29, 2025 •

edited

Loading

This PR needs a `release notes:` label

digantdesai Sep 3, 2025 •

edited

Loading

JacobSzwejbka commented Sep 3, 2025 •

edited

Loading