passes on edge dialect will be done via tranform (#236)

Summary:
Pull Request resolved: #236

For graph transform in edge dialect, let's use `transform`. Also we don't expect any transform from aten to edge as discussed {F1081902441}

Reviewed By: angelayi

Differential Revision: D48911694

fbshipit-source-id: 3f895b6e8266511b38c0521066bea5c681080a30

Author: cccclai

backends/xnnpack/test/TARGETS

@@ -104,6 +104,7 @@ python_unittest(
         "//executorch/backends/xnnpack/passes:xnnpack_passes",
         "//executorch/backends/xnnpack/utils:xnnpack_utils",
         "//executorch/exir:lib",
+        "//executorch/exir:pass_base",
         "//executorch/exir/backend/canonical_partitioners:canonical_partitioner_lib",
         "//executorch/exir/dialects:lib",
     ],
@@ -130,7 +131,7 @@ python_unittest(
         "//caffe2:torch",
         "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
         "//executorch/backends/xnnpack/test/tester:tester",
-        "//executorch/backends/xnnpack/utils:xnnpack_utils",
+        "//executorch/exir:lib",
         "//pytorch/vision:torchvision",
     ],
 )

backends/xnnpack/test/test_xnnpack_utils.py

@@ -18,6 +18,7 @@
     XnnpackQuantizedPartitioner,
 )
 from executorch.backends.xnnpack.utils.configs import (
+    get_transform_passes,
     get_xnnpack_edge_compile_config,
     get_xnnpack_executorch_backend_config,
 )
@@ -323,7 +324,9 @@ def quantize_and_test_model_with_quantizer(
             config=exir.CaptureConfig(enable_aot=True, _unlift=True),
         )
 
-        edge_program = captured_program.to_edge(get_xnnpack_edge_compile_config())
+        edge_program = captured_program.to_edge(
+            get_xnnpack_edge_compile_config()
+        ).transform(*get_transform_passes())
         delegated_module = self.lower_module_and_test_output(
             module=edge_program,
             sample_inputs=example_inputs,

backends/xnnpack/utils/TARGETS

@@ -6,6 +6,7 @@ python_library(
     deps = [
         "//caffe2:torch",
         "//executorch/exir:lib",
+        "//executorch/exir:pass_manager",
         "//executorch/exir/backend/canonical_partitioners:canonical_partitioner_lib",
         "//executorch/exir/dialects:lib",
     ],

backends/xnnpack/utils/configs.py

@@ -4,24 +4,28 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Optional
+from typing import List, Optional
 
 import executorch.exir as exir
 from executorch.exir import CaptureConfig
 from executorch.exir.backend.canonical_partitioners.duplicate_dequant_node_pass import (
     DuplicateDequantNodePass,
 )
+from executorch.exir.pass_manager import PassType
 
 ### XNNPACK Configs ###
-def get_xnnpack_edge_compile_config(additional_passes=None) -> exir.EdgeCompileConfig:
-    additional_passes = additional_passes if additional_passes else []
-    passes = additional_passes + [DuplicateDequantNodePass()]
+def get_xnnpack_edge_compile_config() -> exir.EdgeCompileConfig:
     return exir.EdgeCompileConfig(
-        passes=passes,
         _check_ir_validity=False,
     )
 
 
+def get_transform_passes(additional_passes=None) -> List[PassType]:
+    additional_passes = additional_passes if additional_passes else []
+    passes = additional_passes + [DuplicateDequantNodePass()]
+    return passes
+
+
 def get_xnnpack_executorch_backend_config(
     additional_passes=None,
 ) -> exir.ExecutorchBackendConfig:

backends/xnnpack/utils/utils.py

@@ -10,6 +10,7 @@
 import torch
 
 from executorch.backends.xnnpack.utils.configs import (
+    get_transform_passes,
     get_xnnpack_capture_config,
     get_xnnpack_edge_compile_config,
 )
@@ -25,11 +26,15 @@ def capture_graph_for_xnnpack(
     enable_aot: Optional[bool] = None,
     unlift: Optional[bool] = None,
 ) -> exir.ExirExportedProgram:
-    return exir.capture(
-        module,
-        inputs,
-        get_xnnpack_capture_config(enable_aot=enable_aot, unlift=unlift),
-    ).to_edge(get_xnnpack_edge_compile_config())
+    return (
+        exir.capture(
+            module,
+            inputs,
+            get_xnnpack_capture_config(enable_aot=enable_aot, unlift=unlift),
+        )
+        .to_edge(get_xnnpack_edge_compile_config())
+        .transform(*get_transform_passes())
+    )
 
 
 ### XNNPACK Utils ###

examples/backend/xnnpack_examples.py

@@ -91,13 +91,16 @@
         capture_config=CaptureConfig(enable_aot=True),
         edge_compile_config=EdgeCompileConfig(
             # TODO(T162080278): Duplicated Dequant nodes will be in quantizer spec
-            _check_ir_validity=False if args.quantize else True,
-            passes=[DuplicateDequantNodePass()],
+            _check_ir_validity=False
+            if args.quantize
+            else True,
         ),
     )
     logging.info(f"Exported graph:\n{edge.exported_program.graph}")
 
-    edge.exported_program = to_backend(edge.exported_program, partitioner)
+    edge.exported_program = to_backend(
+        edge.transform(DuplicateDequantNodePass()).exported_program, partitioner
+    )
     logging.info(f"Lowered graph:\n{edge.exported_program.graph}")
 
     exec_prog = edge.to_executorch()

exir/capture/_config.py

@@ -29,7 +29,6 @@ class CaptureConfig:
 @compatibility(is_backward_compatible=False)
 @dataclass
 class EdgeCompileConfig:
-    passes: List[PassType] = field(default_factory=list)
     # TODO(qihan): remove ability to opt out
     _check_ir_validity: bool = True
     # TODO(larryliu): remove this

exir/emit/test/test_emit.py

@@ -179,7 +179,8 @@ def f(x: torch.Tensor) -> torch.Tensor:
 
         program = (
             exir.capture(f, (torch.randn(100),), exir.CaptureConfig())
-            .to_edge(exir.EdgeCompileConfig(passes=[ConstPropPass()]))
+            .to_edge()
+            .transform(ConstPropPass())
             .to_executorch()
             .program
         )
@@ -641,7 +642,8 @@ def forward(self, x):
         x = (torch.randn(1, 1, 2, 2),)
         program = (
             exir.capture(M(), x, exir.CaptureConfig())
-            .to_edge(exir.EdgeCompileConfig(passes=[ConstPropPass()]))
+            .to_edge()
+            .transform(ConstPropPass())
             .to_executorch()
             .program
         )

exir/program/_program.py

@@ -283,7 +283,7 @@ def _to_edge(ep, config: EdgeCompileConfig) -> "ExirExportedProgram":
         aten_to_edge_passes.passes[:-2]
         + op_replace_pass
         + aten_to_edge_passes.passes[-2:]
-    ) + config.passes
+    )
     new_ep = copy.deepcopy(ep).transform(*passes)
     if config._check_ir_validity:
         EXIREdgeDialectVerifier(check_edge_ops=config._use_edge_ops)(

exir/tests/test_memory_planning.py

@@ -462,19 +462,20 @@ def quantize(self, eager_model: nn.Module) -> nn.Module:
     def test_asr_joiner(self) -> None:
         eager_model = self.quantize(ASRJoiner())
         inputs = eager_model.get_random_inputs()
-        edge_program = exir.capture(
-            eager_model,
-            inputs,
-            exir.CaptureConfig(
-                enable_dynamic_shape=True,
-            ),
-        ).to_edge(
-            exir.EdgeCompileConfig(
-                passes=[
-                    ConstPropPass(),
-                ],
-                _check_ir_validity=False,
+        edge_program = (
+            exir.capture(
+                eager_model,
+                inputs,
+                exir.CaptureConfig(
+                    enable_dynamic_shape=True,
+                ),
+            )
+            .to_edge(
+                exir.EdgeCompileConfig(
+                    _check_ir_validity=False,
+                )
             )
+            .transform(ConstPropPass())
         )
         with validation_disabled():
             backend_module = to_backend(