diff --git a/exir/backend/test/test_partitioner.py b/exir/backend/test/test_partitioner.py index b963cf5c931..b10cdcd436b 100644 --- a/exir/backend/test/test_partitioner.py +++ b/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): diff --git a/exir/backend/utils.py b/exir/backend/utils.py index bd76a8d4e10..fc8d0ceeb70 100644 --- a/exir/backend/utils.py +++ b/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)