Move verification to after to_backend (#12630)

Summary:

Some operators require preservation because they are intended to be consumed by a backend. These operators can contain view and mutation, as they won't be part of the graph after to_backend.

If there are still view and mutation ops after to_backend, verification should throw an error.

This diff:
1. Removes verification check from _generated_edge_program, which is called by to_edge and to_edge_transform_and_lower
2. Add verification check in to_edge
3. to_edge_transform_and_lower already has a verification check after to_backend; add config.preserve ops there.

Remove preserve_ops_view test; after to_edge, there should be no view ops.

https://www.internalfb.com/code/fbsource/[f677b61422b9927eabc36d3b15857b06186cf7ef]/fbcode/executorch/exir/program/_program.py?lines=835

Reviewed By: metascroy

Differential Revision: D78535519

Author: lucylq

exir/program/_program.py

@@ -791,15 +791,13 @@ def edge_to_executorch_passes(
 
 
 def _generate_edge_program(
-    name: str,
     config: EdgeCompileConfig,
     program: ExportedProgram,
     core_aten_ops_exception_list: Optional[List[torch._ops.OpOverload]] = None,
     preserve_ops: Optional[List[torch._ops.OpOverload]] = None,
 ) -> ExportedProgram:
     """
     Args:
-        name: The name of the program.
         config: The configuration for the edge program.
         program: The exported program to be converted to an edge program.
         core_aten_ops_exception_list: A list of aten ops that are missing decompositions to core aten.
@@ -816,18 +814,6 @@ def _generate_edge_program(
     # Remove unused parameters
     program = remove_unused_parameters_pass(program)
 
-    if config._check_ir_validity:
-        try:
-            EXIRATenDialectVerifier(
-                edge_compile_config=config,
-                class_only=False,
-                core_aten_ops_exception_list=core_aten_ops_exception_list,
-                preserve_ops=preserve_ops,
-            )(gm)
-        except ExportError as e:
-            logging.info(f"Input program {name} is not in ATen dialect.")
-            raise e
-
     pre_op_replace_passes, post_op_replace_passes = _get_aten_to_edge_passes(config)
 
     passes = []
@@ -1144,7 +1130,6 @@ def _gen_edge_manager_for_partitioners(
         edge_programs[name] = program
 
         edge_programs[name] = _generate_edge_program(
-            name,
             config,
             program,
             preserve_ops=list(ops_set_to_not_decompose_by_program.get(name, [])),
@@ -1288,11 +1273,12 @@ def to_edge_transform_and_lower(
                 generate_error=True,
             )
 
+        preserve_ops = config.preserve_ops + list(ops_set_to_not_decompose)
         if config._check_ir_validity:
             EXIREdgeDialectVerifier(
                 edge_compile_config=config,
                 class_only=True,
-                preserve_ops=list(ops_set_to_not_decompose),
+                preserve_ops=preserve_ops,
             )()(program.graph_module)
 
     return edge_manager
@@ -1337,8 +1323,18 @@ def to_edge(
                 table.pop(op, None)
         program = program.run_decompositions(table)
         edge_programs[name] = _generate_edge_program(
-            name, config, program, preserve_ops=preserve_ops
+            config, program, preserve_ops=preserve_ops
         )
+        if config._check_ir_validity:
+            try:
+                EXIREdgeDialectVerifier(
+                    edge_compile_config=config,
+                    class_only=True,
+                    preserve_ops=preserve_ops,
+                )()(edge_programs[name].graph_module)
+            except ExportError as e:
+                logging.info(f"Input program {name} is not in aten dialect.")
+                raise e
 
     return EdgeProgramManager(edge_programs, constant_methods, config)
 

exir/verification/test/test_verifier.py

@@ -161,17 +161,3 @@ def forward(self, input, label):
         edge_verifier = EXIREdgeDialectVerifier()
 
         edge_verifier(edge.exported_program())
-
-    def test_verifier_preserve_ops_view(self) -> None:
-        class TestExpand(nn.Module):
-            def __init__(self):
-                super().__init__()
-
-            def forward(self, x):
-                return x.expand(2, 2, 2, 2)
-
-        model = TestExpand()
-        config = EdgeCompileConfig(preserve_ops=[torch.ops.aten.expand.default])
-        export_model = export(model, (torch.randn(2, 2, 2, 2),), strict=True)
-        with self.assertRaises(RuntimeError):
-            to_edge(export_model, compile_config=config)

exir/verification/verifier.py

@@ -145,7 +145,7 @@ def check_valid_op(self, op):
                     # which may affect memory planning.
                     if op.is_view:
                         raise RuntimeError(
-                            f"Cannot preserve operator {op} because it is a view or mutation."
+                            f"Cannot preserve operator {op} because it is a view."
                         )
                     if op._schema.is_mutable:
                         logging.warning(