move memory format pass into to_edge (pytorch#1891)

Summary:
Pull Request resolved: pytorch#1891

This diff enable memory foramt pass by moving it into to_edge.
Also introduced `_skip_dim_order` in edge compile config for gradually enable the pass. Currently we set its default as True, only enable it in `test_memory_format_ops_pass` for evaluation.
Will graduatlly enable it in our system, and finally remove it from EdgeCompileConfig

Reviewed By: larryliu0820

Differential Revision: D53567636

fbshipit-source-id: 969dea799d785d825e796386afcdfa7144a47ba9

Author: Gasoonjia

exir/capture/_config.py

@@ -35,6 +35,8 @@ class EdgeCompileConfig:
     # TODO(larryliu): remove this
     _use_edge_ops: bool = True
     _skip_type_promotion: bool = False
+    # TODO(gasoonjia): set it as False by default, and remove it in the long term
+    _skip_dim_order: bool = True
 
 
 @compatibility(is_backward_compatible=False)

exir/passes/__init__.py

@@ -464,7 +464,7 @@ def dead_code_elimination_pass(graph_module: torch.fx.GraphModule) -> PassResult
 
 # Passes to convert a graph module from ATen to Edge IR
 
-pre_op_replace_passes = PassManager(
+base_pre_op_replace_passes: List[Callable[[torch.nn.Module], PassResult]] = PassManager(
     passes=[
         # ReplaceSymSizeOpPass need to be run before other passes which inherits
         # from ExportPass. ExportPass can not handle OpOverloadPacket in its
@@ -479,7 +479,9 @@ def dead_code_elimination_pass(graph_module: torch.fx.GraphModule) -> PassResult
     ]
 ).passes
 
-post_op_replace_passes = PassManager(
+base_post_op_replace_passes: List[
+    Callable[[torch.nn.Module], PassResult]
+] = PassManager(
     passes=[
         dead_code_elimination_pass,
         DebugHandleGeneratorPass(),

exir/program/_program.py

@@ -20,10 +20,11 @@
 from executorch.exir.error import ExportError
 from executorch.exir.pass_manager import PassType
 from executorch.exir.passes import (
+    base_post_op_replace_passes,
+    base_pre_op_replace_passes,
     EdgeToBackendOpsPass,
+    MemoryFormatOpsPass,
     OpReplacePass,
-    post_op_replace_passes,
-    pre_op_replace_passes,
 )
 from executorch.exir.passes.remove_graph_asserts_pass import RemoveGraphAssertsPass
 from executorch.exir.passes.remove_mixed_type_operators import RemoveMixedTypeOperators
@@ -459,6 +460,23 @@ def dump_exported_program(self) -> ExportedProgram:
         return self.exported_program
 
 
+def _get_aten_to_edge_passes(config: EdgeCompileConfig):
+    # 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.
+    # It might be regenerated in SpecPropPass so it may not be visiable. However debug handle will be lost.
+
+    pre_op_replace_passes = base_pre_op_replace_passes + (
+        [] if config._skip_type_promotion else [RemoveMixedTypeOperators()]
+    )
+
+    post_op_replace_passes = (
+        [] if config._skip_dim_order else [MemoryFormatOpsPass()]
+    ) + base_post_op_replace_passes
+
+    return pre_op_replace_passes, post_op_replace_passes
+
+
 def _to_edge(ep, config: EdgeCompileConfig) -> "ExirExportedProgram":
     if config._check_ir_validity:
         try:
@@ -486,15 +504,9 @@ def _to_edge(ep, config: EdgeCompileConfig) -> "ExirExportedProgram":
             ),
             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.
-    # It might be regenerated in SpecPropPass so it may not be visiable. However debug handle will be lost.
+    pre_op_replace_passes, post_op_replace_passes = _get_aten_to_edge_passes(config)
 
-    passes = pre_op_replace_passes + (
-        [] if config._skip_type_promotion else [RemoveMixedTypeOperators()]
-    )
-    new_ep = copy.deepcopy(ep).transform(*passes)
+    new_ep = copy.deepcopy(ep).transform(*pre_op_replace_passes)
     if dialect == "ATEN":
         new_ep.exported_program = lift_constant_tensor_pass(new_ep.exported_program)
 
@@ -824,17 +836,13 @@ def to_edge(
                 logging.info(f"Input program {name} is not in ATen dialect.")
                 raise e
 
-        # 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.
-        # It might be regenerated in SpecPropPass so it may not be visiable. However debug handle will be lost.
+        pre_op_replace_passes, post_op_replace_passes = _get_aten_to_edge_passes(config)
+
         passes = []
         passes.append(
             ReplaceViewOpsWithViewCopyOpsPass()
         )  # TODO move inside aten_to_edge passes after all users are migrated off v1 capture
         passes.extend(pre_op_replace_passes)
-        if not config._skip_type_promotion:
-            passes.append(RemoveMixedTypeOperators())
         if config._use_edge_ops:
             passes.append(OpReplacePass())
 

exir/tests/test_memory_format_ops_pass.py

@@ -75,23 +75,23 @@ class TestSet:
                 edge_op_str
             ).run(before.graph_module.code)
 
-            ep = to_edge(
-                before, compile_config=EdgeCompileConfig(_check_ir_validity=False)
-            )  # Only replacing edge_ops
-
-            # Run the pass
-            # TODO move this in to_edge passes, make to_dim_copy pass verifier
-            after = ep.transform([MemoryFormatOpsPass()], check_ir_validity=False)
+            # TODO(gasoonjia): make to_dim_copy pass verifier
+            epm = to_edge(
+                before,
+                compile_config=EdgeCompileConfig(
+                    _check_ir_validity=False, _skip_dim_order=False
+                ),
+            )
 
             # check op strings
             FileCheck().check_not(aten_op_str).check_count(
                 edge_op_str, 1, exactly=True
-            ).run(after.exported_program().graph_module.code)
+            ).run(epm.exported_program().graph_module.code)
 
             # check EdgeOp and the new BackendOp should behave the same
             expected = before(*test_set.sample_input)
-            actual = after.exported_program()(*test_set.sample_input)
+            actual = epm.exported_program()(*test_set.sample_input)
             self.assertTrue(torch.allclose(actual, expected))
 
             # TODO - more
-            after.to_executorch()
+            epm.to_executorch()