Enable verifier [2/n] (#1155)

Summary:
Pull Request resolved: #1155

X-link: pytorch/pytorch#113075

Turn on verifier check for exportec program ctor. Note that this effectively detect a large surface of spec violations, so we also spend some time fixing them one by one in this diff.

Reviewed By: angelayi

Differential Revision: D51014944

fbshipit-source-id: 9a9e689ec9c8561a9c5cd4b6bd8836809ae5ff30

Author: zhxchen17

backends/example/example_backend.py

@@ -32,9 +32,10 @@ def preprocess(
         print("entering  the lowerable parts in ExampleBackend.preprocess....")
 
         copy_edge_program = copy.deepcopy(edge_program)
-        copy_edge_program._transform(
-            PermuteMemoryFormatsPass(),
-            MergeToDimPass(),
-        )
-        processed_bytes = str(copy_edge_program.graph)
+        graph_module = copy_edge_program.graph_module
+        graph_module_res = PermuteMemoryFormatsPass()(graph_module)
+        assert graph_module_res is not None
+        graph_module_res = MergeToDimPass()(graph_module_res.graph_module)
+        assert graph_module_res is not None
+        processed_bytes = str(graph_module_res.graph_module.graph)
         return PreprocessResult(bytes(processed_bytes, encoding="utf8"))

exir/backend/backend_api.py

@@ -326,4 +326,6 @@ def to_backend(
         copy.deepcopy(edge_program.range_constraints),
         copy.deepcopy(edge_program.equality_constraints),
         copy.deepcopy(edge_program.module_call_graph),
+        None,
+        edge_program.verifier,
     )

exir/capture/_capture.py

@@ -22,7 +22,7 @@
     flatten_output,
     Value,
 )
-from executorch.exir.verification.verifier import EXIRATenDialectVerifier
+from executorch.exir.verification.verifier import EXIRATenDialectVerifierBase
 from torch import _guards
 from torch._dispatch.python import enable_python_dispatcher
 from torch._dynamo.eval_frame import Constraint
@@ -79,6 +79,16 @@ def _capture_legacy_do_not_use(f, args) -> ExirExportedProgram:
     assert output_node.op == "output"
     user_outputs = [arg.name for arg in output_node.args[0]]
 
+    for n in graph_module.graph.nodes:
+        if n.op == "call_function" and "val" not in n.meta:
+            try:
+                args, kwargs = pytree.tree_map_only(
+                    torch.fx.Node, lambda x: x.meta["val"], (n.args, n.kwargs)
+                )
+                n.meta["val"] = n.target(*args, **kwargs)
+            except Exception:
+                n.meta["val"] = None
+
     ep = HackedUpExportedProgramDONOTUSE(
         graph_module,
         graph_module.graph,
@@ -112,6 +122,7 @@ def _capture_legacy_do_not_use(f, args) -> ExirExportedProgram:
             )
         ],
         None,
+        EXIRATenDialectVerifierBase,
     )
     return ExirExportedProgram(ep, False)
 
@@ -281,6 +292,7 @@ def convert_to_fake(x):
         for arg in output_node.args[0]
     ]
 
+    graph_module.graph.eliminate_dead_code()
     ep = ExportedProgram(
         graph_module,
         graph_module.graph,
@@ -300,7 +312,7 @@ def convert_to_fake(x):
             )
         ],
         None,
-        EXIRATenDialectVerifier,
+        EXIRATenDialectVerifierBase,
     )
     return ExirExportedProgram(ep, False)
 

exir/emit/test/test_emit.py

@@ -146,13 +146,12 @@ def test_basic_end_to_end(self) -> None:
             self.assertIn(op.overload, {"out", "unary_out"})
 
         self.assertEqual(ops[0].name, "aten::sin")
-        self.assertEqual(ops[1].name, "aten::max")
 
         self.assertEqual(len(exec_plan.inputs), 1)
         self.assertEqual(len(exec_plan.outputs), 1)
 
         self.assertEqual(exec_plan.inputs[0], 0)
-        self.assertEqual(exec_plan.outputs[0], 2)
+        self.assertEqual(exec_plan.outputs[0], 1)
 
     def test_nested_return(self) -> None:
         def f(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor]]:

exir/lowered_backend_module.py

@@ -210,6 +210,9 @@ def program(self, emit_stacktrace: bool = False) -> Program:
         delegate_node.meta["spec"] = tuple(
             [make_spec(node.meta["val"]) for node in original_output_nodes]
         )
+        delegate_node.meta["val"] = tuple(
+            [node.meta["val"] for node in original_output_nodes]
+        )
 
         # The getitem nodes that are going to be inserted to the lowered graph module
         getitem_nodes = []
@@ -218,6 +221,7 @@ def program(self, emit_stacktrace: bool = False) -> Program:
                 operator.getitem,
                 args=(delegate_node, i),
             )
+            getitem_node.meta["val"] = delegate_node.meta["val"][i]
             getitem_nodes.append(getitem_node)
         lowered_exported_program.graph.output(getitem_nodes)
 
@@ -264,6 +268,8 @@ def program(self, emit_stacktrace: bool = False) -> Program:
             range_constraints=lowered_exported_program.range_constraints,
             equality_constraints=lowered_exported_program.equality_constraints,
             module_call_graph=lowered_exported_program.module_call_graph,
+            example_inputs=None,
+            verifier=lowered_exported_program.verifier,
         )
         exported_program = exported_program._transform(
             SpecPropPass(), MemoryPlanningPass("greedy")
@@ -466,6 +472,7 @@ def create_exported_program_from_submodule(
         range_constraints=copy.deepcopy(owning_program.range_constraints),
         equality_constraints=[],
         module_call_graph=[],
+        verifier=owning_program.verifier,
     )
 
 

exir/program/_program.py

@@ -32,6 +32,7 @@
 from executorch.exir.verification.verifier import (
     EXIRATenDialectVerifier,
     EXIREdgeDialectVerifier,
+    get_aten_verifier,
 )
 from torch._export import ExportedProgram
 from torch._export.passes import ReplaceViewOpsWithViewCopyOpsPass
@@ -44,6 +45,21 @@
 Val = Any
 
 
+def _copy_module(new_prog, new_gm):
+    new_prog.meta.update(new_gm.meta)
+    new_prog.graph = new_gm.graph
+    submodules = [name for name, _ in new_prog.named_children()]
+    for name in submodules:
+        delattr(new_prog, name)
+    for name, mod in new_gm.named_children():
+        setattr(new_prog, name, mod)
+    for node in new_gm.graph.nodes:
+        if node.op == "get_attr":
+            t = getattr(new_gm, node.target, None)
+            if isinstance(t, torch.Tensor):
+                setattr(new_prog, node.target, t)
+
+
 def lift_constant_tensor_pass(ep):
     """
     Takes an ExportedProgram and returns the ExportedProgram modified in-place,
@@ -131,6 +147,7 @@ def __init__(
         equality_constraints,
         module_call_graph,
         example_inputs,
+        verifier,
     ):
         super().__init__(
             root,
@@ -141,8 +158,8 @@ def __init__(
             equality_constraints,
             module_call_graph,
             example_inputs,
+            verifier,
         )
-        self._dialect = "HACKED_ATEN"
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
         import torch._export.error as error
@@ -245,9 +262,15 @@ def to_executorch(
         if not self.after_to_edge_passes:
             raise RuntimeError("Must run to_edge before to_executorch.")
         config = config or ExecutorchBackendConfig()
-        ep = self.exported_program
-        new_prog = ep._transform(*edge_to_executorch_passes(config))
-        new_prog = ExirExportedProgram(new_prog, self.after_to_edge_passes)
+        new_gm = self.exported_program.graph_module
+        for p in edge_to_executorch_passes(config):
+            new_gm_res = p(new_gm)
+            assert new_gm_res is not None
+            new_gm = new_gm_res.graph_module
+        new_prog = ExirExportedProgram(
+            copy.deepcopy(self.exported_program), self.after_to_edge_passes
+        )
+        _copy_module(new_prog.exported_program.graph_module, new_gm)
         executorch_prog = ExecutorchProgram(
             new_prog,
             emit_stacktrace=config.emit_stacktrace,
@@ -256,9 +279,7 @@ def to_executorch(
             constant_tensor_alignment=config.constant_tensor_alignment,
             delegate_alignment=config.delegate_alignment,
         )
-        executorch_prog.graph_module.meta.update(
-            new_prog.exported_program.graph_module.meta
-        )
+        executorch_prog.graph_module.meta.update(new_gm.meta)
         executorch_prog.graph_module.meta.update(
             self.exported_program.graph_module.meta
         )
@@ -356,6 +377,22 @@ def _to_edge(ep, config: EdgeCompileConfig) -> "ExirExportedProgram":
             )
             raise
 
+    dialect = ep.exported_program.dialect
+    if dialect == "ATEN":
+        ep = ExirExportedProgram(
+            ExportedProgram(
+                ep.exported_program.graph_module,
+                ep.exported_program.graph_module.graph,
+                ep.exported_program.graph_signature,
+                ep.exported_program.state_dict,
+                ep.exported_program.range_constraints,
+                ep.exported_program.equality_constraints,
+                ep.exported_program.module_call_graph,
+                ep.exported_program.example_inputs,
+                verifier=get_aten_verifier(enable=config._check_ir_validity),
+            ),
+            False,
+        )
     # TODO: the last two passes for aten_to_edge need to be eliminated_dead_code -> debug_handle_generator. After enable
     # use_edge_op it can be moved to aten_to_edge_passes before eliminated_dead_code pass. Also ExportPass doesn't play
     # well with node.meta, meaning after some passes permuting operators, we may lose some information in node.meta.
@@ -364,16 +401,23 @@ def _to_edge(ep, config: EdgeCompileConfig) -> "ExirExportedProgram":
     post_op_replace_passes = aten_to_edge_passes.passes[-2:]
 
     new_ep = copy.deepcopy(ep).transform(*pre_op_replace_passes)
-    if new_ep.exported_program.dialect == "ATEN":
+    if dialect == "ATEN":
         new_ep.exported_program = lift_constant_tensor_pass(new_ep.exported_program)
 
+    new_gm = new_ep.exported_program.graph_module
     if config._use_edge_ops:
-        new_ep = new_ep.transform(OpReplacePass())
+        new_gm_res = OpReplacePass()(new_gm)
+        assert new_gm_res is not None
+        new_gm = new_gm_res.graph_module
+
+    for p in post_op_replace_passes:
+        new_gm_res = p(new_gm)
+        assert new_gm_res is not None
+        new_gm = new_gm_res.graph_module
 
-    new_ep = new_ep.transform(*post_op_replace_passes)
     new_ep.exported_program = ExportedProgram(
-        new_ep.exported_program.graph_module,
-        new_ep.exported_program.graph,
+        new_gm,
+        new_gm.graph,
         new_ep.exported_program.graph_signature,
         new_ep.exported_program.state_dict,
         new_ep.exported_program.range_constraints,
@@ -386,11 +430,6 @@ def _to_edge(ep, config: EdgeCompileConfig) -> "ExirExportedProgram":
             class_only=True,
         ),
     )
-    if config._check_ir_validity:
-        # TODO(zhxchen17) Remove this call after we turn on verifier in ctor.
-        EXIREdgeDialectVerifier(check_edge_ops=config._use_edge_ops)(
-            new_ep.exported_program.graph_module
-        )
     new_ep.after_to_edge_passes = True
     return new_ep
 
@@ -623,8 +662,18 @@ def multi_method_program_to_executorch(
 ) -> MultiMethodExecutorchProgram:
     config = config or ExecutorchBackendConfig()
     passes = edge_to_executorch_passes(config)
+    res = {}
+    for method_name, prog in edge_dialect_program._method_to_program.items():
+        new_prog = copy.deepcopy(prog)
+        gm = prog.exported_program.graph_module
+        for p in passes:
+            gm_res = p(gm)
+            assert gm_res is not None
+            gm = gm_res.graph_module
+        _copy_module(new_prog.exported_program.graph_module, gm)
+        res[method_name] = new_prog
     return MultiMethodExecutorchProgram(
-        executorch_dialect_program=edge_dialect_program.transform(*passes),
+        executorch_dialect_program=MultiMethodExirExportedProgram(res),
         emit_stacktrace=config.emit_stacktrace,
         extract_segments=config.extract_segments,
         segment_alignment=config.segment_alignment,
@@ -674,8 +723,6 @@ def to_edge(
                 logging.info(f"Input program {name} is not in ATen dialect.")
                 raise e
 
-        op_replace_pass = [OpReplacePass()] if config._use_edge_ops else []
-
         # TODO: the last two passes for aten_to_edge need to be eliminated_dead_code -> debug_handle_generator. After enable
         # use_edge_op it can be moved to aten_to_edge_passes before eliminated_dead_code pass. Also ExportPass doesn't play
         # well with node.meta, meaning after some passes permuting operators, we may lose some information in node.meta.
@@ -686,9 +733,31 @@ def to_edge(
             ReplaceViewOpsWithViewCopyOpsPass()
         )  # TODO move inside aten_to_edge passes after all users are migrated off v1 capture
         passes.extend(aten_to_edge_passes.passes[:-2])
-        passes.extend(op_replace_pass)
-        passes.extend(aten_to_edge_passes.passes[-2:])
         edge_program = program._transform(*passes)
+        if config._use_edge_ops:
+            gm_res = OpReplacePass()(edge_program.graph_module)
+            assert gm_res is not None
+            gm = gm_res.graph_module
+        else:
+            gm = edge_program.graph_module
+        edge_program = ExportedProgram(
+            root=gm,
+            graph=gm.graph,
+            graph_signature=edge_program.graph_signature,
+            state_dict=edge_program.state_dict,
+            range_constraints=edge_program.range_constraints,
+            equality_constraints=edge_program.equality_constraints,
+            module_call_graph=edge_program.module_call_graph,
+            example_inputs=edge_program.example_inputs,
+            verifier=EXIREdgeDialectVerifier(
+                check_edge_ops=config._use_edge_ops,
+                enable=config._check_ir_validity,
+                class_only=True,
+            ),
+        )
+        passes = []
+        passes.extend(aten_to_edge_passes.passes[-2:])
+        edge_program = edge_program._transform(*passes)
         try:
             EXIREdgeDialectVerifier(
                 check_edge_ops=config._use_edge_ops, enable=config._check_ir_validity
@@ -852,7 +921,13 @@ def to_executorch(
 
         execution_programs: Dict[str, ExportedProgram] = {}
         for name, program in self._edge_programs.items():
-            new_prog = program._transform(*edge_to_executorch_passes(config))
+            new_gm = program.graph_module
+            for p in edge_to_executorch_passes(config):
+                new_gm_res = p(new_gm)
+                assert new_gm_res is not None
+                new_gm = new_gm_res.graph_module
+            new_prog = copy.deepcopy(program)
+            _copy_module(new_prog.graph_module, new_gm)
             execution_programs[name] = new_prog
 
         return ExecutorchProgramManager(

exir/serde/serialize.py

@@ -37,6 +37,7 @@
 )
 from torch._export.serde.schema import GraphSignature
 from torch._export.serde.serialize import SerializeError
+from torch._export.verifier import load_verifier
 from torch.export.exported_program import (
     ExportGraphSignature,
     ModuleCallEntry,
@@ -367,7 +368,7 @@ def serialize(
                 range_constraints=serialized_range_constraints,
                 equality_constraints=serialized_equality_constraints,
                 schema_version=schema.SCHEMA_VERSION,
-                example_inputs=None,
+                dialect=exported_program.dialect,
             ),
             export_serialize.serialize_torch_artifact(gm_serializer.state_dict),
         )
@@ -686,15 +687,23 @@ def deserialize(
             serialized_exported_program.equality_constraints
         )
 
-        return exir.ExportedProgram(
+        dummy_g = torch.fx.Graph()
+        dummy_g.output(())
+        ep = exir.ExportedProgram(
             state_dict,
-            graph_module.graph,
+            dummy_g,
             sig,
             {},  # TODO(T157676982)
             range_constraints,
             equality_constraints,
             module_call_graph,
+            None,
+            load_verifier(serialized_exported_program.dialect),
         )
+        ep.graph_module.graph = graph_module.graph
+        for name, t in state_dict.items():
+            setattr(ep.graph_module, name, t)
+        return ep
 
 
 def serialize(

exir/tests/control_flow_models.py

@@ -65,7 +65,8 @@ def __init__(self):
 
     def forward(self, inp):
         def true_branch(x):
-            return torch.inverse(x).contiguous()
+            x - 1
+            return x + 1
 
         def false_branch(x):
             return x * 2