Move verification to after to_backend

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.

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, [])),
@@ -1271,6 +1256,19 @@ def to_edge_transform_and_lower(
         edge_manager = edge_manager.to_backend(method_to_partitioner)
 
     for name, program in edge_manager._edge_programs.items():
+        # Check ir validity after to_backend, which should consume any ops
+        # that contain view or mutation.
+        if config._check_ir_validity:
+            try:
+                EXIREdgeDialectVerifier(
+                    edge_compile_config=config,
+                    class_only=True,
+                    preserve_ops=config.preserve_ops,
+                )()(program.graph_module)
+            except ExportError as e:
+                logging.info(f"Input program {name} is not in aten dialect.")
+                raise e
+
         ops_set_to_not_decompose: Set[torch._ops.OpOverload] = set()
         partitioners = partitioner.get(name, [])
         for curr_partitioner in partitioners:
@@ -1337,8 +1335,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/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(