Fix creating graph signature in delegation (#2580)

Summary: Pull Request resolved: #2580

Reviewed By: JacobSzwejbka

Differential Revision: D55225522

fbshipit-source-id: 3f665ecebaa555879dc0adb4e2295f6522f4c6af

Author: angelayi

exir/lowered_backend_module.py

@@ -27,6 +27,7 @@
 
 from torch._subclasses import FakeTensor
 from torch.export.exported_program import (
+    ConstantArgument,
     ExportedProgram,
     ExportGraphSignature,
     InputKind,
@@ -422,7 +423,7 @@ def arrange_graph_placeholders(
 
 
 # TODO Don't regenerate new signature manually.
-def _get_new_signature(
+def _get_new_signature(  # noqa: C901
     original_program: ExportedProgram,
     gm: torch.fx.GraphModule,
     tag: Optional[str] = None,
@@ -431,6 +432,18 @@ def _get_new_signature(
     Dict[str, Union[torch.Tensor, torch.nn.Parameter]],
     Dict[str, Union[torch.Tensor, torch.ScriptObject]],
 ]:
+    """
+    Args:
+        tag: If tag is None, this means that we are constructing the graph
+        signature for the toplevel graph, after delegation. We need to do this
+        because sometimes delegates will swallow some parameters/buffers, so we
+        need to update the graph signature/state dict to reflect these changes.
+        Otherwise, if tag is not None, this means we are constructing the graph
+        signature for the delegated modules. In this case, we need to look
+        through the input nodes and see which ones were originally
+        parameters/buffers, and lower them down to the delegate.
+    """
+
     old_signature = original_program.graph_signature
 
     input_specs = []
@@ -441,84 +454,104 @@ def _get_new_signature(
     new_state_dict = {}
     new_constants = {}
 
-    non_persistent_buffers = set(old_signature.non_persistent_buffers)
+    input_tensor_node_to_sig = {
+        input_spec.arg.name: input_spec
+        for input_spec in old_signature.input_specs
+        if isinstance(input_spec.arg, TensorArgument)
+    }
 
     for node in gm.graph.nodes:
-        is_tagged = node.meta.get("delegation_tag", None) == tag
+        is_tagged = tag is None or node.meta.get("delegation_tag", None) == tag
         if node.op == "placeholder":
-            if node.name in old_signature.inputs_to_parameters and is_tagged:
-                parameter_name = old_signature.inputs_to_parameters[node.name]
-                # add param to graph signature
-                input_specs.append(
-                    InputSpec(
-                        kind=InputKind.PARAMETER,
-                        arg=TensorArgument(name=node.name),
-                        target=parameter_name,
-                    )
-                )
 
-                # add param to state_dict
-                new_state_dict[parameter_name] = original_program.state_dict[
-                    parameter_name
-                ]
-            elif node.name in old_signature.inputs_to_buffers and is_tagged:
-                buffer_name = old_signature.inputs_to_buffers[node.name]
-                persistent = buffer_name not in non_persistent_buffers
-                # add buffer to graph signature
+            if node.name not in input_tensor_node_to_sig:
+                assert tag is not None
                 input_specs.append(
                     InputSpec(
-                        kind=InputKind.BUFFER,
+                        kind=InputKind.USER_INPUT,
                         arg=TensorArgument(name=node.name),
-                        target=buffer_name,
-                        persistent=persistent,
+                        target=None,
                     )
                 )
+                continue
 
-                # add param to new_state_dict
-                if persistent:
-                    new_state_dict[buffer_name] = original_program.state_dict[
-                        buffer_name
-                    ]
-                else:
-                    new_constants[buffer_name] = original_program.constants[buffer_name]
-            elif (
-                node.name in old_signature.inputs_to_lifted_tensor_constants
-                and is_tagged
-            ):
-                constant_name = old_signature.inputs_to_lifted_tensor_constants[
-                    node.name
-                ]
-                # add constant to graph signature
-                input_specs.append(
-                    InputSpec(
-                        kind=InputKind.CONSTANT_TENSOR,
-                        arg=TensorArgument(name=node.name),
-                        target=constant_name,
+            orig_input_spec = input_tensor_node_to_sig[node.name]
+
+            if not isinstance(orig_input_spec.arg, TensorArgument):
+                input_specs.append(orig_input_spec)
+
+            elif is_tagged:
+                input_specs.append(orig_input_spec)
+
+                if orig_input_spec.kind in (InputKind.PARAMETER, InputKind.BUFFER):
+                    new_state_dict[orig_input_spec.target] = (
+                        original_program.state_dict[orig_input_spec.target]
                     )
-                )
+                elif orig_input_spec.kind in (
+                    InputKind.CONSTANT_TENSOR,
+                    InputKind.CUSTOM_OBJ,
+                ):
+                    new_constants[orig_input_spec.target] = original_program.constants[
+                        orig_input_spec.target
+                    ]
 
-                # add constant to new_constants
-                new_constants[constant_name] = original_program.constants[constant_name]
             else:
-                # not param, buffer, or lifted_tensor_constant then user input
                 input_specs.append(
                     InputSpec(
                         kind=InputKind.USER_INPUT,
                         arg=TensorArgument(name=node.name),
                         target=None,
                     )
                 )
+
         if node.op == "output":
-            for output in pytree.tree_leaves((node.args, node.kwargs)):
-                if not isinstance(output, torch.fx.Node):
-                    continue
-                output_specs.append(
-                    OutputSpec(
-                        kind=OutputKind.USER_OUTPUT,
-                        arg=TensorArgument(name=output.name),
-                        target=None,
-                    )
-                )
+            output_nodes = pytree.tree_leaves((node.args, node.kwargs))
+
+            if tag is not None:
+                # We are constructing output_specs for the delegate outputs.
+                # These don't have any buffer mutations.
+
+                for output_node in output_nodes:
+                    if not isinstance(output_node, torch.fx.Node):
+                        output_specs.append(
+                            OutputSpec(
+                                kind=OutputKind.USER_OUTPUT,
+                                arg=ConstantArgument(output_node),
+                                target=None,
+                            )
+                        )
+                    else:
+                        output_specs.append(
+                            OutputSpec(
+                                kind=OutputKind.USER_OUTPUT,
+                                arg=TensorArgument(name=output_node.name),
+                                target=None,
+                            )
+                        )
+
+            else:
+                # We are reconstruting the toplevel module which contains
+                # delegates. Delegation should not change the number of outputs
+                # in the toplevel module, and it does not touch the mutated buffers
+
+                assert len(old_signature.output_specs) == len(output_nodes)
+                for prev_output_spec, output_node in zip(
+                    old_signature.output_specs, output_nodes
+                ):
+                    if not isinstance(output_node, torch.fx.Node):
+                        assert isinstance(prev_output_spec.arg, ConstantArgument)
+                        output_specs.append(
+                            OutputSpec(
+                                kind=OutputKind.USER_OUTPUT,
+                                arg=ConstantArgument(output_node),
+                                target=None,
+                            )
+                        )
+
+                    else:
+                        new_output_spec = copy.deepcopy(prev_output_spec)
+                        new_output_spec.arg.name = output_node.name
+                        output_specs.append(new_output_spec)
 
     return new_signature, new_state_dict, new_constants