Allow delegate to consume buffer mutations

Summary:
Fixing #4209

Edge Program:
```
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, b_b: "f32[3, 3]", x: "f32[3, 3]"):
             # File: /data/users/angelayi/fbsource/buck-out/v2/gen/fbcode/389acaeb40d57230/executorch/exir/backend/test/__test_partitioner__/test_partitioner#link-tree/executorch/exir/backend/test/test_partitioner.py:631 in forward, code: self.b.add_(x)
            aten_add_tensor: "f32[3, 3]" = executorch_exir_dialects_edge__ops_aten_add_Tensor(b_b, x);  b_b = None

             # File: /data/users/angelayi/fbsource/buck-out/v2/gen/fbcode/389acaeb40d57230/executorch/exir/backend/test/__test_partitioner__/test_partitioner#link-tree/executorch/exir/backend/test/test_partitioner.py:632 in forward, code: return x + self.b
            aten_add_tensor_1: "f32[3, 3]" = executorch_exir_dialects_edge__ops_aten_add_Tensor(x, aten_add_tensor);  x = None
            return (aten_add_tensor, aten_add_tensor_1)

Graph signature: ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.BUFFER: 3>, arg=TensorArgument(name='b_b'), target='b', persistent=True), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='x'), target=None, persistent=None)], output_specs=[OutputSpec(kind=<OutputKind.BUFFER_MUTATION: 3>, arg=TensorArgument(name='aten_add_tensor'), target='b'), OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='aten_add_tensor_1'), target=None)])
```

Partitioned / lowered Exported Program (buffer mutation gets removed):
```
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[3, 3]"):
            # No stacktrace found for following nodes
            lowered_module_0 = self.lowered_module_0
            executorch_call_delegate = torch.ops.higher_order.executorch_call_delegate(lowered_module_0, x);  lowered_module_0 = x = None

             # File: /data/users/angelayi/fbsource/buck-out/v2/gen/fbcode/389acaeb40d57230/executorch/exir/backend/test/__test_partitioner__/test_partitioner#link-tree/executorch/exir/backend/test/test_partitioner.py:632 in forward, code: return x + self.b
            getitem_1: "f32[3, 3]" = executorch_call_delegate[0];  executorch_call_delegate = None
            return (getitem_1,)

Graph signature: ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='x'), target=None, persistent=None)], output_specs=[OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='getitem_1'), target=None)])
```

Delegate (consumes the buffer mutation):
```
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, b_b: "f32[3, 3]", x: "f32[3, 3]"):
             # File: /data/users/angelayi/fbsource/buck-out/v2/gen/fbcode/389acaeb40d57230/executorch/exir/backend/test/__test_partitioner__/test_partitioner#link-tree/executorch/exir/backend/test/test_partitioner.py:631 in forward, code: self.b.add_(x)
            aten_add_tensor: "f32[3, 3]" = executorch_exir_dialects_edge__ops_aten_add_Tensor(b_b, x);  b_b = None

             # File: /data/users/angelayi/fbsource/buck-out/v2/gen/fbcode/389acaeb40d57230/executorch/exir/backend/test/__test_partitioner__/test_partitioner#link-tree/executorch/exir/backend/test/test_partitioner.py:632 in forward, code: return x + self.b
            aten_add_tensor_1: "f32[3, 3]" = executorch_exir_dialects_edge__ops_aten_add_Tensor(x, aten_add_tensor);  x = None
            return (aten_add_tensor, aten_add_tensor_1)

Graph signature: ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.BUFFER: 3>, arg=TensorArgument(name='b_b'), target='b', persistent=True), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='x'), target=None, persistent=None)], output_specs=[OutputSpec(kind=<OutputKind.BUFFER_MUTATION: 3>, arg=TensorArgument(name='aten_add_tensor'), target='b'), OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='aten_add_tensor_1'), target=None)])
```

Differential Revision: D60838243

Author: angelayi

exir/backend/test/TARGETS

@@ -290,6 +290,7 @@ python_unittest(
         "//executorch/exir/backend/test/demos/rpc:executor_backend_register",
     ],
     deps = [
+        ":op_partitioner_demo",
         "//caffe2:torch",
         "//executorch/exir:lib",
         "//executorch/exir/backend:backend_details",

exir/backend/test/test_partitioner.py

@@ -26,6 +26,7 @@
 from executorch.exir.backend.test.demos.rpc.executor_backend_preprocess import (
     ExecutorBackend,
 )
+from executorch.exir.backend.test.op_partitioner_demo import AddAttributePartitionerDemo
 from executorch.exir.backend.utils import get_delegates, tag_constant_data
 
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -619,3 +620,90 @@ def partition(
             and node.target == torch.ops.aten.copy_.default
         ]
         self.assertEqual(len(copy_node), 1)
+
+    def test_buffer_mutation1(self):
+        class TestModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("b", torch.ones(3, 3))
+
+            def forward(self, x):
+                self.b.add_(x)
+                return x + self.b
+
+        model_inputs = (torch.ones(3, 3),)
+        orig_res = TestModule()(*model_inputs)
+        edge_program = exir.to_edge(torch.export.export(TestModule(), model_inputs))
+        lowered = edge_program.to_backend(AddAttributePartitionerDemo())
+
+        self.assertTrue(
+            torch.allclose(lowered.exported_program().module()(*model_inputs), orig_res)
+        )
+
+        self.assertEqual(
+            len(lowered.exported_program().graph_signature.buffers_to_mutate),
+            0,
+        )
+        lowered_module_nodes = get_delegates(lowered.exported_program().graph)
+        self.assertEqual(len(lowered_module_nodes), 1)
+        lowered_module_node = lowered_module_nodes[0]
+
+        # get call delegate node
+        call_delegate_node = list(lowered_module_node.users.keys())[0]
+        self.assertEqual(len(call_delegate_node.args), 2)
+
+        lower_module = getattr(
+            lowered.exported_program().graph_module, lowered_module_node.name
+        )
+        delegated_ep = lower_module.original_module
+
+        self.assertEqual(len(delegated_ep.state_dict), 1)
+        self.assertEqual(len(delegated_ep.graph_signature.buffers_to_mutate), 1)
+        self.assertEqual(len(delegated_ep.graph_signature.buffers), 1)
+
+    def test_buffer_mutation2(self):
+        SHAPE = (2, 3)
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("state_1", torch.zeros(SHAPE, dtype=torch.float32))
+
+            def forward(self, x):
+                add = self.state_1.add_(x)
+                return add
+
+        model = Model()
+        model.eval()
+
+        example_inputs = (torch.randn(SHAPE),)
+        exir_program_aten = torch.export.export(model, example_inputs)
+        edge_program_manager = exir.to_edge(exir_program_aten)
+        lowered = edge_program_manager.to_backend(AddAttributePartitionerDemo())
+
+        self.assertTrue(
+            torch.allclose(
+                lowered.exported_program().module()(*example_inputs), example_inputs[0]
+            )
+        )
+
+        self.assertEqual(
+            len(lowered.exported_program().graph_signature.buffers_to_mutate),
+            0,
+        )
+        lowered_module_nodes = get_delegates(lowered.exported_program().graph)
+        self.assertEqual(len(lowered_module_nodes), 1)
+        lowered_module_node = lowered_module_nodes[0]
+
+        # get call delegate node
+        call_delegate_node = list(lowered_module_node.users.keys())[0]
+        self.assertEqual(len(call_delegate_node.args), 2)
+
+        lower_module = getattr(
+            lowered.exported_program().graph_module, lowered_module_node.name
+        )
+        delegated_ep = lower_module.original_module
+
+        self.assertEqual(len(delegated_ep.state_dict), 1)
+        self.assertEqual(len(delegated_ep.graph_signature.buffers_to_mutate), 1)
+        self.assertEqual(len(delegated_ep.graph_signature.buffers), 1)

exir/lowered_backend_module.py

@@ -488,8 +488,16 @@ def _get_new_signature(  # noqa: C901
         else {}
     )
 
+    toplevel_output_node_to_sig: Dict[str, OutputSpec] = (
+        {
+            output_spec.arg.name: output_spec
+            for output_spec in old_signature.output_specs
+        }
+        if not is_submodule
+        else {}
+    )
+
     for node in gm.graph.nodes:
-        is_tagged = tag is None or node.meta.get("delegation_tag", None) == tag
         if node.op == "placeholder":
 
             if node.name not in input_node_to_sig:
@@ -507,7 +515,7 @@ def _get_new_signature(  # noqa: C901
             if not isinstance(orig_input_spec.arg, TensorArgument):
                 input_specs.append(orig_input_spec)
 
-            elif is_tagged:
+            elif node.meta.get("delegation_tag", None) == tag:
                 input_specs.append(orig_input_spec)
 
                 if orig_input_spec.kind == InputKind.USER_INPUT:
@@ -551,11 +559,55 @@ def _get_new_signature(  # noqa: C901
                 )
 
         if node.op == "output":
-            output_nodes = pytree.tree_leaves((node.args, node.kwargs))
+            buffer_mutation_idxs: Dict[int, OutputSpec] = {}
+            for user in call_module_node.users.keys():
+                if user.name in toplevel_output_node_to_sig:
+                    assert (
+                        user.op == "call_function" and user.target == operator.getitem
+                    ), f"Invalid user {user}, node.op is {user.op} and node.target is {user.target}"
+                    getitem_idx = user.args[1]
+                    assert isinstance(
+                        getitem_idx, int
+                    ), f"Invalid getitem type: {type(getitem_idx)}"
+                    buffer_mutation_idxs[getitem_idx] = toplevel_output_node_to_sig[
+                        user.name
+                    ]
+
+            for i, output_node in enumerate(node.args[0]):
+                if i in buffer_mutation_idxs:
+                    assert isinstance(output_node, torch.fx.Node)
+                    orig_output_spec = buffer_mutation_idxs[i]
+
+                    if (
+                        orig_output_spec.kind == OutputKind.BUFFER_MUTATION
+                        and orig_output_spec.target in new_state_dict
+                    ):
+                        # If the delegate wants to consume the buffer, then
+                        # the delegate should also consume the buffer
+                        # mutation (output spec would be a BUFFER_MUTATION).
+                        # Otherwise the delegate will just return the result
+                        # of the mutation as a USER_OUTPUT.
+                        output_specs.append(
+                            OutputSpec(
+                                kind=OutputKind.BUFFER_MUTATION,
+                                arg=TensorArgument(name=output_node.name),
+                                target=orig_output_spec.target,
+                            )
+                        )
+                        output_specs_to_delete[orig_output_spec.arg.name] = (
+                            orig_output_spec
+                        )
 
-            for output_node in output_nodes:
+                    else:
+                        output_specs.append(
+                            OutputSpec(
+                                kind=OutputKind.USER_OUTPUT,
+                                arg=TensorArgument(name=output_node.name),
+                                target=None,
+                            )
+                        )
 
-                if not isinstance(output_node, torch.fx.Node):
+                elif not isinstance(output_node, torch.fx.Node):
                     output_specs.append(
                         OutputSpec(
                             kind=OutputKind.USER_OUTPUT,
@@ -774,7 +826,7 @@ def get_lowered_backend_modules(
     return lowered_programs
 
 
-def _unsafe_adjust_original_program(
+def _unsafe_adjust_original_program(  # noqa: C901
     original_program: ExportedProgram,
     call_delegate_node: torch.fx.Node,
     input_specs_to_delete: Dict[str, InputSpec],
@@ -830,3 +882,50 @@ def _unsafe_adjust_original_program(
             del original_program._constants[input_spec.target]
         else:
             raise RuntimeError(f"Invalid input spec {input_spec} received")
+
+    # Delete buffer mutations from the output which were consumed by the delegate
+    toplevel_output_node = None
+    for node in reversed(original_program.graph.nodes):
+        if node.op == "output":
+            toplevel_output_node = node
+            break
+
+    assert toplevel_output_node is not None
+    assert (
+        len(toplevel_output_node.args) == 1
+    ), f"Invalid output node: {toplevel_output_node} with args {toplevel_output_node.args}"
+
+    new_output_args = [
+        arg
+        for arg in toplevel_output_node.args[0]
+        if not isinstance(arg, torch.fx.Node) or arg.name not in output_specs_to_delete
+    ]
+    toplevel_output_node.args = (tuple(new_output_args),)
+
+    # Delete the buffer mutation getitem nodes
+    getitem_idxs: List[int] = []
+    user_nodes = list(call_delegate_node.users.keys())
+    for user in user_nodes:
+        if user.name in output_specs_to_delete:
+            assert (
+                user.op == "call_function" and user.target == operator.getitem
+            ), f"Invalid user {user}, node.op is {node.op} and node.target is {node.target}"
+            user_idx = user.args[1]
+            assert isinstance(user_idx, int), f"Invalid getitem type: {type(user_idx)}"
+            getitem_idxs.append(user_idx)
+            original_program.graph.erase_node(user)
+
+    getitem_idxs.sort(reverse=True)
+
+    # Adjust all the getitem indices after the deleted getitems
+    user_nodes = list(call_delegate_node.users.keys())
+    for user in user_nodes:
+        assert user.op == "call_function" and user.target == operator.getitem
+        user_idx = user.args[1]
+        assert isinstance(user_idx, int)
+        for i, idx in enumerate(getitem_idxs):
+            if user_idx > idx:
+                user.args = (user.args[0], user_idx - (len(getitem_idxs) - i))
+                break
+
+    original_program._validate()