Extend tranform() method params with PassManager. (#13140)

Summary:
Pull Request resolved: #13140

There is a user request to pass a custom PassManager to the  EdgeProgramManager::transform(). This diff extends method API appropriately.

Differential Revision: D79292048

Author: eigen-k

exir/program/_program.py

@@ -240,8 +240,27 @@ 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
 
@@ -1493,19 +1512,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_or_pass_manager: 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_or_pass_manager: Can either 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
@@ -1517,19 +1540,32 @@ def transform(
         """
         compile_config = compile_config or self.compile_config
         new_programs: Dict[str, ExportedProgram] = {}
-        if isinstance(passes, dict):
+        if isinstance(passes_or_pass_manager, dict):
+            # Transform only those methods which match the keys in the provided dictionary.
+            passes_dict : Dict[str, Sequence[PassType]] = passes_or_pass_manager
             for name, program in self._edge_programs.items():
-                if name in passes.keys():
-                    new_programs[name] = _transform(program, *passes[name])
+                if name in passes_dict.keys():
+                    new_programs[name] = _transform(program, *passes_dict[name])
                     EXIREdgeDialectVerifier(edge_compile_config=compile_config)(
                         new_programs[name].graph_module
                     )
                 else:
                     new_programs[name] = copy.deepcopy(program)
 
-        else:  # apply passes to every method
+        if isinstance(passes_or_pass_manager, Sequence):
+            # Transform all methods with the provided passes.
+            passes_seq : Sequence[PassType] = passes_or_pass_manager
+            for name, program in self._edge_programs.items():
+                new_programs[name] = _transform(program, *passes_seq)
+                EXIREdgeDialectVerifier(edge_compile_config=compile_config)(
+                    new_programs[name].graph_module
+                )
+
+        if isinstance(passes_or_pass_manager, PassManager):
+            # Transform all methods with the provided PassManager.
+            pass_manager : PassManager = passes_or_pass_manager
             for name, program in self._edge_programs.items():
-                new_programs[name] = _transform(program, *passes)
+                new_programs[name] = _transform_with_pass_manager(program, pass_manager)
                 EXIREdgeDialectVerifier(edge_compile_config=compile_config)(
                     new_programs[name].graph_module
                 )