Fix mutated-buffer tagging after indirect mutation detection change

exir/backend/test/test_partitioner.py

@@ -29,7 +29,11 @@
     AddAttributePartitionerDemo,
     AllNodesPartitionerDemo,
 )
-from executorch.exir.backend.utils import get_delegates, tag_constant_data
+from executorch.exir.backend.utils import (
+    get_delegates,
+    tag_constant_data,
+    tag_mutated_buffer,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.tests.models import MLP
 from executorch.extension.pybindings.portable_lib import (  # @manual=//executorch/extension/pybindings:portable_lib
@@ -652,6 +656,35 @@ def forward(self, x):
         )
         self.assertNotIn("delegation_tag", state_node.meta)
 
+    def test_tag_mutated_buffer_detects_indirect_buffer_mutation(self) -> None:
+        class MutableStateModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("state", torch.zeros(1))
+
+            def forward(self, x):
+                result = x + self.state
+                self.state.copy_(x)
+                return result
+
+        edge = exir.to_edge(
+            torch.export.export(MutableStateModule(), (torch.ones(1),), strict=True)
+        )
+        exported_program = edge.exported_program()
+        state_node = next(
+            node
+            for node in exported_program.graph.nodes
+            if exported_program.graph_signature.inputs_to_buffers.get(node.name)
+            == "state"
+        )
+        delegate_tag = "test_partition"
+        for user in state_node.users:
+            user.meta["delegation_tag"] = delegate_tag
+
+        tag_mutated_buffer(exported_program)
+
+        self.assertEqual(state_node.meta.get("delegation_tag"), delegate_tag)
+
     def test_buffer_mutation1(self):
         class TestModule(torch.nn.Module):
             def __init__(self):

exir/backend/utils.py

@@ -406,18 +406,13 @@ def tag_mutated_buffer(edge_program: ExportedProgram) -> None:
     # Cache signature lookups to avoid rebuilding dicts on every access.
     sig = edge_program.graph_signature
     buffers_map = sig.inputs_to_buffers
-    buffers_to_mutate = sig.buffers_to_mutate
+    mutated_buffer_targets = set(sig.buffers_to_mutate.values())
 
     for node in edge_program.graph.nodes:
-        # Determine whether this node is a mutated buffer
-        is_mutated_buffer_node = False
-        if node.op == "placeholder" and node.name in buffers_map:
-            for node_user in node.users:
-                if node_user.name in buffers_to_mutate:
-                    is_mutated_buffer_node = True
-                    break
-        # This node is mutated buffer, tag it
-        if is_mutated_buffer_node:
+        if (
+            node.op == "placeholder"
+            and buffers_map.get(node.name) in mutated_buffer_targets
+        ):
             user_tags = set()
             for user in node.users:
                 user_tag = user.meta.get("delegation_tag", None)