Fix export of zero-stride constants and inputs, and of constants that delegates share

exir/backend/canonical_partitioners/duplicate_constant_node_pass.py

@@ -40,6 +40,7 @@ def duplicate_constant_node(
     """
     A pass to duplicate the attributes/constants node (the candidate_node) in the graph. Mostly used for duplicating light-weight data.
     If the data is too large, try tag it with "no_copy" to prevent high memory usage and make it as part of the output.
+    A mutated buffer is never copied, because a copy would not see its updates: this raises instead.
 
     Args:
         exported_program: the exported program to be modified. If constants nodes are copied, they will be added as new
@@ -76,9 +77,23 @@ def duplicate_constant_node(
                 raise RuntimeError(
                     f"placeholder node for non-params, non-buffer, and non-tensor constants should not be tagged: {node} "
                 )
-            users = list(node.users.keys())
-            for ith in range(1, len(node.users)):
-                copy_constant_or_attribute_fqn = node.name + f"_copy_{ith - 1}"
+            if (
+                old_signature.inputs_to_buffers.get(node.name)
+                in old_signature.buffers_to_mutate.values()
+            ):
+                raise RuntimeError(
+                    f"mutated buffer ({node}) cannot be copied, because a copy "
+                    "would not see its updates. Tag it only when every node "
+                    "that reads or writes it has the same tag, as "
+                    "tag_mutated_buffer does, or leave it untagged."
+                )
+            # The graph output keeps the original, which the output specs
+            # name, so then every other user reads a copy.
+            users = [user for user in node.users if user.op != "output"]
+            if len(users) == len(node.users):
+                users = users[1:]
+            for ith, user in enumerate(users):
+                copy_constant_or_attribute_fqn = node.name + f"_copy_{ith}"
                 with exported_program.graph.inserting_before(
                     constant_or_attribute_node
                 ):
@@ -96,26 +111,12 @@ def duplicate_constant_node(
                     copied_constant_or_attribute_node.meta["val"] = (
                         constant_or_attribute_node.meta["val"]
                     )
-                    new_args = tuple(
-                        [
-                            (
-                                arg
-                                if arg != constant_or_attribute_node
-                                else copied_constant_or_attribute_node
-                            )
-                            for arg in users[ith].args
-                        ]
+                    # This also reaches a constant that an op such as cat or
+                    # index_put reads from a list.
+                    user.replace_input_with(
+                        constant_or_attribute_node,
+                        copied_constant_or_attribute_node,
                     )
-                    new_kwargs = {
-                        key: (
-                            value
-                            if value != constant_or_attribute_node
-                            else copied_constant_or_attribute_node
-                        )
-                        for key, value in users[ith].kwargs.items()
-                    }
-                    users[ith].args = new_args
-                    users[ith].kwargs = new_kwargs
                     # The copy keeps the original target, so it reads the
                     # original tensor and needs no data of its own.
                     new_input_specs.append(

exir/backend/test/BUCK

@@ -372,12 +372,15 @@ fbcode_target(_kind = runtime.python_test,
         "test_utils.py",
     ],
     deps = [
+        ":backend_with_compiler_demo",
         ":op_partitioner_demo",
         "//caffe2:torch",
         "//executorch/exir:lib",
         "//executorch/exir/backend:backend_api",
+        "//executorch/exir/backend:compile_spec_schema",
         "//executorch/exir/backend:partitioner",
         "//executorch/exir/backend:utils",
+        "//executorch/exir/backend/test/demos/rpc:executor_backend_preprocess",
         "//executorch/exir/dialects:lib",
         "//executorch/extension/pybindings:portable_lib",  # @manual
         "//executorch/runtime/executor/test:test_backend_compiler_lib",

exir/backend/test/test_partitioner.py

@@ -658,32 +658,42 @@ def forward(self, x):
 
     def test_tag_mutated_buffer_detects_indirect_buffer_mutation(self) -> None:
         class MutableStateModule(torch.nn.Module):
-            def __init__(self):
+            def __init__(self, compute_new_value: bool):
                 super().__init__()
+                self.compute_new_value = compute_new_value
                 self.register_buffer("state", torch.zeros(1))
 
             def forward(self, x):
                 result = x + self.state
-                self.state.copy_(x)
+                self.state.copy_(x + 1 if self.compute_new_value else 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"
-        )
+        # A new value that is a graph input is never in a partition, so the
+        # buffer then stays in the owning program.
         delegate_tag = "test_partition"
-        for user in state_node.users:
-            user.meta["delegation_tag"] = delegate_tag
+        for compute_new_value, expected_tag in ((True, delegate_tag), (False, None)):
+            with self.subTest(compute_new_value=compute_new_value):
+                edge = exir.to_edge(
+                    torch.export.export(
+                        MutableStateModule(compute_new_value),
+                        (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"
+                )
+                for node in exported_program.graph.nodes:
+                    if node.op == "call_function":
+                        node.meta["delegation_tag"] = delegate_tag
 
-        tag_mutated_buffer(exported_program)
+                tag_mutated_buffer(exported_program)
 
-        self.assertEqual(state_node.meta.get("delegation_tag"), delegate_tag)
+                self.assertEqual(state_node.meta.get("delegation_tag"), expected_tag)
 
     def test_buffer_mutation1(self):
         class TestModule(torch.nn.Module):

exir/backend/test/test_passes.py

@@ -86,3 +86,26 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         self.assertEqual(len(copied_nodes), 1)
         self.assertNotEqual(user.kwargs["other"], buffer_node)
+
+    def test_duplicate_constant_node_refuses_a_mutated_buffer(self) -> None:
+        class UpdateConst(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("const", torch.zeros(2))
+
+            def forward(self, x):
+                y = x * self.const
+                self.const.add_(1)
+                return y
+
+        program = exir.to_edge(
+            export(UpdateConst(), (torch.ones(2),), strict=True)
+        ).exported_program()
+        name = next(
+            node.name
+            for node in program.graph.nodes
+            if node.op == "placeholder" and is_buffer(program, node)
+        )
+
+        with self.assertRaisesRegex(RuntimeError, "would not see its updates"):
+            duplicate_constant_node(program, name)

exir/backend/test/test_utils.py

@@ -10,17 +10,34 @@
 from executorch import exir
 from executorch.exir import to_edge
 from executorch.exir.backend.backend_api import to_backend
-from executorch.exir.backend.partitioner import Partitioner, PartitionResult
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+from executorch.exir.backend.partitioner import (
+    DelegationSpec,
+    Partitioner,
+    PartitionResult,
+)
+from executorch.exir.backend.test.backend_with_compiler_demo import (
+    BackendWithCompilerDemo,
+)
+from executorch.exir.backend.test.demos.rpc.executor_backend_preprocess import (
+    ExecutorBackend,
+)
 from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo
 from executorch.exir.backend.utils import (
     format_delegated_graph,
     get_delegated_payload,
     get_delegates,
     get_non_lowered_nodes,
     is_identical_graph,
+    tag_constant_data,
+    tag_mutated_buffer,
 )
 
 from executorch.exir.dialects._ops import bind_pattern_to_op, ops as exir_ops
+from executorch.extension.pybindings.portable_lib import (  # @manual
+    _get_registered_backend_names,
+    _load_for_executorch_from_buffer,
+)
 from torch.export import export, ExportedProgram
 from torch.fx import symbolic_trace
 from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
@@ -390,3 +407,155 @@ def forward(self, a, x, b):
             payload_names,
             f"Delegate names mismatch: {delegate_names} vs {payload_names}",
         )
+
+    def test_to_backend_keeps_a_mutated_buffer_out_of_a_partial_delegate(  # noqa: C901
+        self,
+    ):
+        class UpdateState(torch.nn.Module):
+            def __init__(self, read: bool, returned: bool, copy_input: bool):
+                super().__init__()
+                self.read = read
+                self.returned = returned
+                self.copy_input = copy_input
+                self.register_buffer("state", torch.zeros(2))
+
+            def forward(self, x):
+                if self.read:
+                    out = torch.sin(self.state) + (x + self.state)
+                else:
+                    out = torch.sin(x)
+                self.state.copy_(x if self.copy_input else x + x)
+                if self.returned:
+                    return out, self.state
+                # In read mode, the last add reads the new value.
+                return out + self.state if self.read else out
+
+        class TagNodes(Partitioner):
+            def __init__(self, tags):
+                super().__init__()
+                self.tags = tags
+
+            def partition(self, program: ExportedProgram) -> PartitionResult:
+                for node in program.graph.nodes:
+                    if node.name in self.tags:
+                        node.meta["delegation_tag"] = self.tags[node.name]
+                tag_mutated_buffer(program)
+                spec = DelegationSpec(
+                    BackendWithCompilerDemo.__name__,
+                    [CompileSpec("max_value", bytes([4]))],
+                )
+                return PartitionResult(
+                    tagged_exported_program=program,
+                    partition_tags={tag: spec for tag in self.tags.values()},
+                )
+
+        # The delegate owns the buffer only when it reads and writes all of
+        # it. Otherwise the owning program keeps the one buffer: a copy in the
+        # delegate would never see the updates.
+        sin, add, writer = "aten_sin_default", "aten_add_tensor", "aten_add_tensor_2"
+        # Also the add that makes "out", and the add that reads the new value.
+        whole = {sin, add, "aten_add_tensor_1", writer, "aten_add_tensor_3"}
+        has_demo_backend = (
+            BackendWithCompilerDemo.__name__ in _get_registered_backend_names()
+        )
+        for read, returned, copy_input, tags, owned_by_program in (
+            (True, False, False, dict.fromkeys(whole, "tag0"), False),
+            (True, False, False, {sin: "tag0", add: "tag0", writer: "tag0"}, True),
+            (True, False, False, {sin: "tag0", add: "tag0"}, True),
+            (True, False, False, {sin: "tag0"}, True),
+            (True, False, False, {sin: "tag0", add: "tag1"}, True),
+            # Lowering cannot give a delegate a buffer whose new value is returned.
+            (True, True, False, {sin: "tag0", add: "tag0", writer: "tag0"}, True),
+            # Without readers, "add" is the writer. Some partitioners tag such a
+            # buffer themselves.
+            (False, False, False, {add: "tag0"}, True),
+            (False, False, False, {add: "tag0", "b_state": "tag0"}, False),
+            (False, True, False, {add: "tag0", "b_state": "tag0"}, True),
+            # The new value is the graph input, which no partition holds.
+            (True, False, True, dict.fromkeys(whole, "tag0"), True),
+        ):
+            with self.subTest(
+                read=read, returned=returned, copy_input=copy_input, tags=tags
+            ):
+                x = torch.full((2,), 0.25)
+                lowered = to_edge(
+                    export(UpdateState(read, returned, copy_input), (x,), strict=True)
+                ).to_backend(TagNodes(tags))
+                self.assertEqual(
+                    [
+                        spec.arg.name
+                        for spec in lowered.exported_program().graph_signature.input_specs
+                        if spec.target == "state"
+                    ],
+                    ["b_state"] if owned_by_program else [],
+                )
+                if copy_input:
+                    # Every op is delegated, so only the new value keeps the
+                    # buffer out.
+                    self.assertEqual(
+                        get_non_lowered_nodes(lowered.exported_program().graph), []
+                    )
+                if tags == {sin: "tag0"} and has_demo_backend:
+                    eager = UpdateState(read, returned, copy_input)
+                    module = _load_for_executorch_from_buffer(
+                        lowered.to_executorch().buffer
+                    )
+                    # The demo backend approximates sin with two Taylor terms.
+                    for _ in range(2):
+                        torch.testing.assert_close(
+                            module.forward((x,))[0], eager(x), atol=1e-3, rtol=0
+                        )
+        if not has_demo_backend:
+            self.skipTest("the runtime has no demo backend to run the lowered program")
+
+    def test_to_backend_copies_a_constant_read_by_ops_in_two_places(self):
+        class CatConst(torch.nn.Module):
+            def __init__(self, returned: bool):
+                super().__init__()
+                self.returned = returned
+                self.register_buffer("const", torch.ones(2))
+
+            def forward(self, x):
+                out = torch.cat([x, self.const]), torch.cat([self.const, x])
+                return (*out, self.const) if self.returned else out
+
+        class TagEachCat(Partitioner):
+            def partition(self, program: ExportedProgram) -> PartitionResult:
+                tags = {}
+                for node in program.graph.nodes:
+                    if node.target == exir_ops.edge.aten.cat.default:
+                        tag = f"tag{len(tags)}"
+                        node.meta["delegation_tag"] = tag
+                        tags[tag] = DelegationSpec(ExecutorBackend.__name__, [])
+                tag_constant_data(program)
+                return PartitionResult(
+                    tagged_exported_program=program, partition_tags=tags
+                )
+
+        # cat reads the constant from a list. Each delegate gets its own copy,
+        # and the graph output keeps the original, which the output specs name.
+        for returned, constants in (
+            (False, [["b_const"], ["b_const_copy_0"]]),
+            (True, [["b_const_copy_0"], ["b_const_copy_1"]]),
+        ):
+            with self.subTest(returned=returned):
+                lowered = to_edge(
+                    export(CatConst(returned), (torch.ones(2),), strict=True)
+                ).to_backend(TagEachCat())
+                program = lowered.exported_program()
+                self.assertEqual(
+                    [
+                        [
+                            spec.arg.name
+                            for spec in module.original_module.graph_signature.input_specs
+                            if spec.target == "const"
+                        ]
+                        for _, module in program.graph_module.named_children()
+                    ],
+                    constants,
+                )
+                if returned:
+                    self.assertEqual(
+                        program.graph.output_node().args[0][-1].name, "b_const"
+                    )
+                lowered.to_executorch()