Change _generate_new_graph_signature (#206)

Summary:
Pull Request resolved: #206

X-link: pytorch/pytorch#108571

Previously `_generate_new_graph_signature` had the assumption that all transformations were not in place. However, this is an incorrect assumption leading to mysterious failures when running passes doing in-place modifications.

This function is technically only needed in the case where the user output node or user input node name is changed. For example, if the user output node was "add" but a pass changes all the "add"s to "mul"s, then the output node will now be named "mul", which we have to update.

For cases where users change the number of user inputs/outputs, number of parameters/buffers, or the names of parameters/buffers it will require extra work on the user's side to update the graph signature, since there is no automatic way for us to detect where to put what.

Note: this doesn't actually change the names for the buffers_to_mutate part of the graph signature, but we're going to assume this is rare, because implementing auto-fixing for that is a little hard...

Reviewed By: digantdesai

Differential Revision: D48917505

fbshipit-source-id: 53132c750ef4b0610fa0bb0fc4c944f4e6a2afc6

Author: angelayi

exir/capture/_capture.py

@@ -204,10 +204,17 @@ def convert_to_fake(x):
     _instantiate_missing_placeholder_val_with_real_inputs(graph_module, flat_args)
     graph_module._apply(torch.Tensor.contiguous)
 
+    user_inputs = [
+        node.name for node in graph_module.graph.nodes if node.op == "placeholder"
+    ]
+    output_node = list(graph_module.graph.nodes)[-1]
+    assert output_node.op == "output"
+    user_outputs = [arg.name for arg in output_node.args[0]]
+
     ep = ExportedProgram(
         graph_module,
         graph_module.graph,
-        ExportGraphSignature([], [], [], [], {}, {}, {}, None),
+        ExportGraphSignature([], [], user_inputs, user_outputs, {}, {}, {}, None),
         CallSpec(in_spec, out_spec),
         {},
         {},

exir/lowered_backend_module.py

@@ -220,7 +220,8 @@ def program(self, emit_stacktrace: bool = False) -> Program:
         lowered_exported_program.graph_signature.user_inputs = [
             user_input
             for user_input in lowered_exported_program.graph_signature.user_inputs
-            if user_input in inputs_to_parameters or user_input in inputs_to_buffers
+            if user_input not in inputs_to_parameters
+            and user_input not in inputs_to_buffers
         ]
         lowered_exported_program.graph_signature.buffers = {}
         lowered_exported_program.graph_signature.parameters = {}