emitter handles lifted graphs

Summary: just needed some extra logic for placeholder nodes. Later can add some assertions in getattr node once lifted is all that exists

Reviewed By: mcr229

Differential Revision: D47888934

fbshipit-source-id: 361019d4e59fc60e146950d98406132ffed00f83

Author: JacobSzwejbka

exir/emit/_emit_program.py

@@ -149,6 +149,13 @@ def emit_program(
 
     # emit each entry point in order according to name.
     for name, exported_program in sorted(methods.items()):
+        if (
+            exported_program.graph_signature.buffers_to_mutate
+        ):  # see if we are mutating any state
+            raise ExportError(
+                ExportErrorType.INVALID_INPUT_TYPE,
+                "Buffers cannot be modified in executorch.",
+            )
         # create empty state
         emitter_state = _EmitterState(
             values=[],

exir/emit/_emitter.py

@@ -39,7 +39,6 @@
 import executorch.extension.pytree as ex_pytree
 import torch
 import torch.fx
-from executorch.exir import delegate
 from executorch.exir.common import add_cursor_to_graph
 from executorch.exir.delegate import executorch_call_delegate, is_lowered_module
 from executorch.exir.dialects.backend._ops import BackendOpOverload
@@ -1270,6 +1269,7 @@ def __init__(
     ) -> None:
         super().__init__(exported_program.graph_module, emitter_state, program_state)
         self.name = name
+        self.exported_program = exported_program
 
         self.inputs: List[int] = []
         self.outputs: List[int] = []
@@ -1301,13 +1301,30 @@ def placeholder(
         https://pytorch.org/docs/stable/fx.html#torch.fx.Graph.placeholder
         """
         spec = self.node.meta["spec"]
+        const_tensor = False
+        if isinstance(target, str) and (
+            target in self.exported_program.graph_signature.inputs_to_parameters
+            or target in self.exported_program.graph_signature.inputs_to_buffers
+        ):
+
+            fqn = (
+                self.exported_program.graph_signature.inputs_to_parameters[target]
+                if target in self.exported_program.graph_signature.inputs_to_parameters
+                else self.exported_program.graph_signature.inputs_to_buffers[target]
+            )
+            spec = TensorSpec.from_tensor(
+                self.exported_program.state_dict[fqn], const=True
+            )
+            const_tensor = True
         evalue = (
             self._tensor_spec_to_evalue(spec)
             if isinstance(spec, TensorSpec)
             else self._constant_to_evalue(spec, None)
         )
         value = self._emit_evalue(evalue)
-        self.inputs.append(value.id)
+        if not const_tensor:
+            self.inputs.append(value.id)
+
         return value
 
     def output(

exir/emit/test/TARGETS

@@ -18,5 +18,6 @@ python_unittest(
         "//executorch/exir/passes:const_prop_pass",
         "//executorch/exir/tests:lib",
         "//executorch/exir/tests:models",
+        "//executorch/extension/pybindings:portable",  # @manual
     ],
 )

exir/emit/test/test_emit.py

@@ -15,11 +15,11 @@
 import executorch.exir.schema as schema
 import executorch.exir.tests.models as models
 import torch
-from executorch.exir import CaptureConfig, ExecutorchProgram
+from executorch.exir import CaptureConfig, EdgeCompileConfig, ExecutorchProgram
 from executorch.exir.emit import emit_program
 from executorch.exir.error import InternalError
 from executorch.exir.passes.const_prop_pass import ConstPropPass
-from executorch.exir.print_program import print_program  # noqa
+from executorch.exir.print_program import pretty_print, print_program  # noqa
 from executorch.exir.schema import (
     Bool,
     EValue,
@@ -36,10 +36,19 @@
     Tensor,
 )
 from executorch.exir.tests.common import register_additional_test_aten_ops
-from executorch.exir.tests.models import Mul
+from executorch.exir.tests.models import (
+    Emformer,
+    FeedForwardBlock,
+    MLP,
+    Mul,
+    ScaledDotProductAttention,
+    ScaledDotProductAttentionModularized,
+)
 from executorch.exir.tracer import ExirDynamoConfig
+from executorch.extension.pybindings.portable import (  # pyre-ignore
+    _load_for_executorch_from_buffer,
+)
 from functorch.experimental import control_flow
-from torch.fx.passes.infra.pass_base import PassResult
 
 
 class TestEmit(unittest.TestCase):
@@ -555,6 +564,39 @@ def f(x: torch.Tensor, n: torch.Tensor) -> torch.Tensor:
             len(program.execution_plan[0].chains[0].instructions[0].instr_args.args), 4
         )
 
+    def test_lifted(self) -> None:
+        def test_model(eager_module):
+            inputs = eager_module.get_random_inputs()
+            eager_output = eager_module.forward(*inputs)
+            capture_config = exir.CaptureConfig(
+                pt2_mode=True,
+                enable_functionalization=True,
+                enable_dynamic_shape=True,
+                enable_aot=True,
+                _unlift=False,
+            )
+
+            aten_dialect = exir.capture(
+                eager_module,
+                eager_module.get_random_inputs(),
+                capture_config,
+            )
+
+            edge_dialect = aten_dialect.to_edge()
+
+            executorch_dialect = edge_dialect.to_executorch()
+
+            pretty_print(executorch_dialect.program)
+
+            executorch_module = _load_for_executorch_from_buffer(
+                executorch_dialect.buffer
+            )
+            et_output = executorch_module.forward(inputs)
+            self.assertTrue(torch.allclose(eager_output, et_output[0], atol=1e-04))
+
+        test_model(MLP())
+        # test_model(Emformer()) cannot run without bernoulli.out being added
+
     def test_emit_multiple_out(self) -> None:
         def f(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
             return torch.topk(x, 5)