Extend EdgeProgramManager::transform() method declaration with PassManager as a "passes" param variant.

Differential Revision: D79292048

Pull Request resolved: #13140

Author: eigen-k

exir/program/_program.py

@@ -240,8 +240,29 @@ def _transform(
         isinstance(p, (list, Verifier)) for p in passes
     ), f"Expected all passes to be of PassType, not list or Verifier. Use override_verifiers kwarg instead. Got: {list(passes)}"
 
-    pm = PassManager(list(passes))
-    res = pm(self.graph_module)
+    return _transform_with_pass_manager(
+        self, PassManager(list(passes)), override_verifiers
+    )
+
+
+def _transform_with_pass_manager(
+    self,
+    pass_manager: PassManager,
+    override_verifiers: None | list[Type[Verifier]] = None,
+) -> "ExportedProgram":
+    """
+    Transforms the program using the provided pass_manager.
+
+    Args:
+        self: The ExportedProgram instance to transform
+        pass_manager: An instance of PassManager to apply transformations.
+        override_verifiers: Optional list of verifier classes to use instead of the default verifiers.
+            This is needed if the transforms yields illegal graph that the default verifier cannot handle.
+
+    Returns:
+        ExportedProgram: A new ExportedProgram with the transformations applied, or self if no changes were made
+    """
+    res = pass_manager(self.graph_module)
     transformed_gm = res.graph_module if res is not None else self.graph_module
     assert transformed_gm is not None
 
@@ -1230,7 +1251,7 @@ def collect_named_data_store_outputs(
 def to_edge_transform_and_lower(  # noqa: C901
     programs: Union[ExportedProgram, Dict[str, ExportedProgram]],
     transform_passes: Optional[
-        Union[Sequence[PassType], Dict[str, Sequence[PassType]]]
+        Union[Sequence[PassType], Dict[str, Sequence[PassType]], PassManager]
     ] = None,
     partitioner: Optional[
         Union[List[Partitioner], Dict[str, List[Partitioner]]]
@@ -1259,11 +1280,15 @@ def to_edge_transform_and_lower(  # noqa: C901
             to their corresponding ExportedPrograms. If only a single ExportedProgram is
             provided it will be assigned the name "forward".
 
-        transform_passes: The passes can either be a list of passes, or a dictionary
-            mapping method names to lists of passes. If it is just a list of passes, all methods
-            in the given EdgeProgramManager will be transformed with the provided passes. If it
-            is a dictionary, only method names specified in the dictionary will be transformed
-            with their corresponding passes.
+        transform_passes: The transform_passes can be one of:
+            1) a list of passes -
+                all methods in the given EdgeProgramManager will be transformed with the provided passes.
+            2) a dictionary -
+                only method names specified in the dictionary will be transformed
+                with their corresponding passes
+            3) an instance of a PassManager -
+                all methods in the given EdgeProgramManager will be
+                transformed with the given PassManager instance.
 
         partitioner: The partitioner can either be a Partitioner subclass instance, or a
             dictionary mapping method names to Partitioner subclass instance. If it is a
@@ -1493,19 +1518,23 @@ def exported_program(self, method_name: str = "forward") -> ExportedProgram:
     @et_logger("transform")
     def transform(
         self,
-        passes: Union[Sequence[PassType], Dict[str, Sequence[PassType]]],
+        passes: Union[Sequence[PassType], Dict[str, Sequence[PassType]], PassManager],
         compile_config: Optional[EdgeCompileConfig] = None,
     ) -> "EdgeProgramManager":
         """
         Transforms the program according to the provided passes.
 
         Args:
-            passes: The passes can either be a list of passes, or a
-                dictionary mapping method names to lists of passes. If it is
-                just a list of passes, all methods in the given EdgeProgramManager
-                will be transformed with the provided passes. If it is a
-                dictionary, only method names specified in the dictionary will be
-                transformed with their corresponding passes.
+            passes: This param can be one of:
+                1) a list of passes -
+                    all methods in the given EdgeProgramManager
+                    will be transformed with the provided passes.
+                2) a dictionary mapping method names to lists of passes -
+                    only method names specified in the dictionary will be
+                    transformed with their corresponding passes.
+                3) a PassManager instance -
+                    all methods in the given EdgeProgramManager will be
+                    transformed with the given PassManager instance.
             compile_config: Compile config to use for veriy the correctness of model
                 graph after each pass. If not specified, the compile config of the
                 calling EdgeProgramManager will be used. It will be used in as compile
@@ -1515,24 +1544,44 @@ def transform(
             EdgeProgramManager: A copy of the calling EdgeProgramManager with the
             transformations applied.
         """
+
         compile_config = compile_config or self.compile_config
         new_programs: Dict[str, ExportedProgram] = {}
+
+        # Cast passes parameter upfront.
+        passes_seq: Optional[Sequence[PassType]] = None
+        passes_dict: Optional[Dict[str, Sequence[PassType]]] = None
+        pass_manager: Optional[PassManager] = None
+
+        if isinstance(passes, Sequence):
+            passes_seq = passes
         if isinstance(passes, dict):
-            for name, program in self._edge_programs.items():
-                if name in passes.keys():
-                    new_programs[name] = _transform(program, *passes[name])
-                    EXIREdgeDialectVerifier(edge_compile_config=compile_config)(
-                        new_programs[name].graph_module
-                    )
-                else:
-                    new_programs[name] = copy.deepcopy(program)
+            passes_dict = passes
+        if isinstance(passes, PassManager):
+            pass_manager = passes
 
-        else:  # apply passes to every method
-            for name, program in self._edge_programs.items():
-                new_programs[name] = _transform(program, *passes)
-                EXIREdgeDialectVerifier(edge_compile_config=compile_config)(
-                    new_programs[name].graph_module
-                )
+        for name, program in self._edge_programs.items():
+            # If the method name is enforced, but not matched, we skip transformation.
+            if (
+                isinstance(passes, dict)
+                and passes_dict
+                and name not in passes_dict.keys()
+            ):
+                new_programs[name] = copy.deepcopy(program)
+                continue
+
+            # Depending on the passes parameter, call the corresponding transform function.
+            if passes_seq is not None:
+                new_programs[name] = _transform(program, *passes_seq)
+            elif passes_dict is not None:
+                new_programs[name] = _transform(program, *passes_dict[name])
+            elif pass_manager is not None:
+                new_programs[name] = _transform_with_pass_manager(program, pass_manager)
+
+            # Verify the correctness of model graph after each transformation.
+            EXIREdgeDialectVerifier(edge_compile_config=compile_config)(
+                new_programs[name].graph_module
+            )
 
         return EdgeProgramManager(
             new_programs, copy.deepcopy(self._config_methods), compile_config

exir/program/test/test_program.py

@@ -37,6 +37,7 @@
 from torch._export.verifier import Verifier
 from torch.export import Dim, export, ExportedProgram
 from torch.export._trace import _export
+from torch.fx.passes.infra.pass_manager import PassManager
 
 from torch.library import impl, Library
 from torch.nn import functional as F
@@ -470,6 +471,30 @@ def test_transform_dict_api(self):
             torch.ones(1) + 1,  # x + 1
         )
 
+    def test_transform_pass_manager_api(self):
+        edge_manager = to_edge(get_exported_programs(), get_config_methods())
+
+        pm = PassManager()
+        pm.add_pass(AddToMulPassEdge())
+
+        transformed_edge = edge_manager.transform(pm)
+
+        x = torch.ones(1) * 2
+        y = torch.ones(1) * 3
+
+        # x * y + x -> x * y * x
+        self.assertEqual(
+            transformed_edge.exported_program("forward").module()(x, y), x * y * x
+        )
+
+        # x + 1 -> x * 1
+        self.assertEqual(
+            transformed_edge.exported_program("foo").module()(
+                x,
+            ),
+            x * 1,
+        )
+
     def test_edge_to_backend_replaces_subgraph(self):
         edge_manager: EdgeProgramManager = to_edge(
             get_exported_programs(), get_config_methods()