Support graphs which return get_attr nodes directly as output

Summary:
X-link: pytorch/pytorch#107610

Currently serializing graphs which return get_attr's directly as output fails. This diff adds support for that only in EXIR serializer while we still support unlifted params.

Reviewed By: angelayi

Differential Revision: D48258552

fbshipit-source-id: 741e6166db86492fdff1ef0b1647601ee351ff77

Author: tarun292

exir/serde/serialize.py

@@ -34,6 +34,7 @@
     CompileSpec,
     LoweredBackendModule as SerdeLoweredBackendModule,
 )
+from torch._export.serde.serialize import SerializeError
 from torch.fx.experimental import symbolic_shapes
 
 log: logging.Logger = logging.getLogger(__name__)
@@ -452,6 +453,24 @@ def deserialize_metadata(self, metadata: Dict[str, str]) -> Dict[str, Any]:
 
         return res
 
+    def deserialize_graph_output(self, output: schema.Argument) -> torch.fx.Node:
+        if isinstance(output.value, schema.TensorArgument):
+            if output.value.name in self.state_dict:  # TODO(T157676982)
+                val = self.state_dict[output.value.name]
+                setattr(self.module, output.value.name, val)
+                node = self.graph.create_node(
+                    "get_attr",
+                    output.value.name,
+                    name=output.value.name,
+                )
+                node.meta = {"val": ""}
+                return node
+            return self.serialized_name_to_node[output.value.name]
+        elif isinstance(output.value, (schema.SymIntArgument, schema.SymBoolArgument)):
+            return self.serialized_name_to_node[output.value.as_name]
+        else:
+            raise SerializeError(f"Unable to deserialize output node {output}")
+
     # pyre-ignore
     def deserialize_alloc_inputs(self, serialized_inputs: List[schema.NamedArgument]):
         def deserialize_alloc_spec(serialized_alloc_spec: str) -> memory.AllocSpec:

exir/tests/test_serde.py

@@ -175,3 +175,20 @@ def forward(self, x):
         edge = exir.capture(m, inputs, exir.CaptureConfig()).to_edge()
         edge_new = deserialize(*serialize(edge.exported_program))
         self.check_ep(edge, edge_new, inputs)
+
+    # Get rid of this test once parameters are lifted by default.
+    def test_return_get_attr_as_outputs(self) -> None:
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.a = torch.ones([1, 1])
+
+            def forward(self, x):
+                return self.a
+
+        m = Model()
+        inputs = (torch.ones([1, 1]),)
+
+        edge = exir.capture(m, inputs, exir.CaptureConfig(pt2_mode=True)).to_edge()
+        edge_new = deserialize(*serialize(edge.exported_program))
+        self.check_ep(edge, edge_new, inputs)