From 6829cff64cdebc08d6270dfe4a84dc213603ca0f Mon Sep 17 00:00:00 2001 From: Conan Truong Date: Fri, 31 Jul 2026 11:22:53 -0700 Subject: [PATCH] Keep delegate-consumed mutable buffers above the delegate (#21507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: `tag_constant_data` freezes a const/param/buffer into a delegate when all of its users are partitioned. It recognized a *mutated* buffer only when a direct user's node name is a key in `buffers_to_mutate`. That misses the case where a buffer both feeds a delegate and is mutated *inside* it: the mutation is produced by a `getitem` off the `call_delegate`, so the buffer's only direct user is the delegate node (whose name is not a `buffers_to_mutate` key), and the buffer is wrongly frozen as constant data instead of staying a method-owned mutable buffer. Additionally treat a placeholder as a mutated buffer when its FQN is a `buffers_to_mutate` target, so a buffer a delegate updates in place stays owned above the delegate. This is the case anticipated by `test_not_delegate_mutable_buffers` ("consider when the delegate can consume the mutable buffer") — e.g. a backend that embeds a pre-compiled engine which updates a KV cache in place. Authored with AI assistance (Claude Code). Differential Revision: D114288921 --- exir/backend/test/test_partitioner.py | 152 ++++++++++++++++++++++++++ exir/backend/utils.py | 47 +++++--- 2 files changed, 186 insertions(+), 13 deletions(-) diff --git a/exir/backend/test/test_partitioner.py b/exir/backend/test/test_partitioner.py index f2ce8e14988..38a85fc2fe3 100644 --- a/exir/backend/test/test_partitioner.py +++ b/exir/backend/test/test_partitioner.py @@ -620,6 +620,158 @@ def partition( ] self.assertEqual(len(copy_node), 1) + def test_delegate_consumes_mutable_buffer(self) -> None: + """The case anticipated by test_not_delegate_mutable_buffers: a delegate + both reads AND updates a mutable buffer, so the buffer's only user is the + already-fused call_delegate and its mutation is produced by a getitem off + that delegate (as when a backend embeds a pre-compiled engine before + partitioning). tag_constant_data must keep the buffer owned above the + delegate -- it recognizes it because its FQN is a buffers_to_mutate + target, not because a direct user is the mutation producer (the direct + user is the call_delegate, whose name is not a mutation key). + """ + + class DelegatedMutableModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("my_state", torch.zeros(1)) + + def forward(self, x): + y = x + self.my_state + self.my_state.add_(x) + return y + + edge = exir.to_edge( + torch.export.export( + DelegatedMutableModule(), (torch.zeros(1),), strict=True + ) + ) + self.assertGreater( + len(edge.exported_program().graph_signature.buffers_to_mutate), 0 + ) + + # Fuse both adds (buffer read + mutation producer) into ONE call_delegate, + # WITHOUT tagging constant data yet -- so afterwards my_state feeds the + # fused call_delegate and its mutation is a getitem off it. + class FuseAddsPartitioner(Partitioner): + def __init__(self): + super().__init__() + self.delegation_spec = DelegationSpec( + ExecutorBackend.__name__, + [CompileSpec(key, value) for key, value in self.spec.items()], + ) + + def partition( + self, edge_exported_program: ExportedProgram + ) -> PartitionResult: + partition_tags = {} + for node in edge_exported_program.graph.nodes: + if node.op == "call_function" and node.target in [ + exir_ops.edge.aten.add.Tensor + ]: + node.meta["delegation_tag"] = "tag0" + partition_tags["tag0"] = self.delegation_spec + return PartitionResult( + tagged_exported_program=edge_exported_program, + partition_tags=partition_tags, + ) + + lowered = edge.to_backend(FuseAddsPartitioner()) + lowered_ep = lowered.exported_program() + + # Sanity: my_state now feeds a call_delegate whose getitem is the mutation. + gs = lowered_ep.graph_signature + self.assertIn("my_state", set(gs.buffers_to_mutate.values())) + mutate_producers = [ + name for name, buf in gs.buffers_to_mutate.items() if buf == "my_state" + ] + self.assertEqual(len(mutate_producers), 1) + self.assertTrue(mutate_producers[0].startswith("getitem")) + + # A backend that embeds a pre-compiled multi-output engine tags the + # delegate node itself, so the mutable buffer feeds a *tagged* delegate + # whose getitem is the mutation. tag_constant_data's freeze loop would + # pull such a buffer into the delegate unless it recognizes it as a + # mutation target. + for node in lowered_ep.graph.nodes: + if node.op == "call_function" and "call_delegate" in str(node.target): + node.meta["delegation_tag"] = "tag0" + + tag_constant_data(lowered_ep) + for node in lowered_ep.graph.nodes: + if ( + node.op == "placeholder" + and gs.inputs_to_buffers.get(node.name) == "my_state" + ): + self.assertIsNone( + node.meta.get("delegation_tag"), + "mutable buffer consumed by a delegate must stay above it, " + "not be frozen into the delegate as constant data", + ) + + def test_delegate_consumes_buffer_keeps_only_mutated_above(self) -> None: + """Companion to test_delegate_consumes_mutable_buffer: a delegate that + consumes both a mutated and a non-mutated buffer must keep only the + *mutated* one above the delegate; the non-mutated buffer is still frozen + into the delegate as constant data. Guards that the mutation-target check + in tag_constant_data is surgical and leaves plain constants untouched. + """ + + class MixedModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("my_state", torch.zeros(2)) + self.register_buffer("frozen_w", torch.ones(2)) + + def forward(self, x): + y = x + self.my_state + self.frozen_w + self.my_state.add_(x) + return y + + edge = exir.to_edge( + torch.export.export(MixedModule(), (torch.zeros(2),), strict=True) + ) + + class FuseAddsPartitioner(Partitioner): + def __init__(self): + super().__init__() + self.delegation_spec = DelegationSpec( + ExecutorBackend.__name__, + [CompileSpec(key, value) for key, value in self.spec.items()], + ) + + def partition( + self, edge_exported_program: ExportedProgram + ) -> PartitionResult: + partition_tags = {} + for node in edge_exported_program.graph.nodes: + if node.op == "call_function" and node.target in [ + exir_ops.edge.aten.add.Tensor + ]: + node.meta["delegation_tag"] = "tag0" + partition_tags["tag0"] = self.delegation_spec + return PartitionResult( + tagged_exported_program=edge_exported_program, + partition_tags=partition_tags, + ) + + lowered_ep = edge.to_backend(FuseAddsPartitioner()).exported_program() + for node in lowered_ep.graph.nodes: + if node.op == "call_function" and "call_delegate" in str(node.target): + node.meta["delegation_tag"] = "tag0" + tag_constant_data(lowered_ep) + + gs = lowered_ep.graph_signature + tag_by_buffer = {} + for node in lowered_ep.graph.nodes: + if node.op == "placeholder" and node.name in gs.inputs_to_buffers: + tag_by_buffer[gs.inputs_to_buffers[node.name]] = node.meta.get( + "delegation_tag" + ) + # mutated buffer stays above the delegate; non-mutated buffer is frozen in. + self.assertIsNone(tag_by_buffer.get("my_state")) + self.assertEqual(tag_by_buffer.get("frozen_w"), "tag0") + 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 9bdba810138..ff1560d8d81 100644 --- a/exir/backend/utils.py +++ b/exir/backend/utils.py @@ -342,6 +342,37 @@ def format_delegated_graph(graph_module: torch.fx.GraphModule) -> str: return graph_format_str +def _find_mutated_buffers( + edge_program: ExportedProgram, + params_map: Dict[str, str], + buffers_map: Dict[str, str], + constants_map: Dict[str, str], + buffers_to_mutate: Dict[str, str], +) -> Set[torch.fx.Node]: + """Return the const/param/buffer placeholder nodes that are mutated (and so + must not be frozen into a delegate as constant data). + + A buffer is mutated when a direct user produces a graph-signature buffer + mutation, or when its FQN is a mutation target -- the latter covers a buffer + mutated inside a delegate, whose mutation is produced by a getitem off the + call_delegate rather than by a direct user of the buffer placeholder. + """ + mutated_buffer_targets = set(buffers_to_mutate.values()) + mutated_buffer: Set[torch.fx.Node] = set() + for node in edge_program.graph.nodes: + if node.op != "placeholder" or not ( + node.name in params_map + or node.name in buffers_map + or node.name in constants_map + ): + continue + if any(user.name in buffers_to_mutate for user in node.users) or ( + buffers_map.get(node.name) in mutated_buffer_targets + ): + mutated_buffer.add(node) + return mutated_buffer + + def tag_constant_data(edge_program: ExportedProgram) -> None: """ Util function for partitioners. This function tags the const/param/buffers nodes @@ -357,19 +388,9 @@ def tag_constant_data(edge_program: ExportedProgram) -> None: constants_map = sig.inputs_to_lifted_tensor_constants buffers_to_mutate = sig.buffers_to_mutate - mutated_buffer = set() - for node in edge_program.graph.nodes: - if node.op == "placeholder" and ( - node.name in params_map - or node.name in buffers_map - or node.name in constants_map - ): - for node_user in node.users: - if node_user.name in buffers_to_mutate: - logging.info( - "The buffer node is a mutated buffer node, which is not constant." - ) - mutated_buffer.add(node) + mutated_buffer = _find_mutated_buffers( + edge_program, params_map, buffers_map, constants_map, buffers_to_mutate + ) for node in edge_program.graph.nodes: # go through const/param/buffer nodes, if all users of const/param/buffer nodes are partitioned then partition