Keep delegate-consumed mutable buffers above the delegate (#21507)

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

Author: Conarnar

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):

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