EdgeProgramManager passes (pytorch#16986)

Summary:

- Adds support to run to run passes on ExportedPrograms and EdgeProgramManager
- Creates an EdgeProgramManagerPassManager

Reviewed By: larryliu0820

Differential Revision: D91725222

Author: Andrew Grebenisan

backends/nxp/tests/ir/edge_passes/test_edge_passes.py

@@ -317,7 +317,9 @@ def test_remove_additional_quantize_dequantize_nodes_pass(self):
             example_input,
         )
 
-        edge_program_manager = edge_program_manager.transform(NeutronEdgePassManager())
+        edge_program_manager = edge_program_manager.transform(
+            NeutronEdgePassManager().passes
+        )
 
         compile_spec = generate_neutron_compile_spec(target)
         partitioner = NeutronPartitioner(

examples/nxp/aot_neutron_compile.py

@@ -351,7 +351,7 @@ def get_model_and_inputs_from_name(model_name: str, use_random_dataset: bool):
         )
 
     edge_program_manager = edge_program_manager.transform(
-        NeutronEdgePassManager([RemoveAdditionalQDQClustersPass()])
+        [RemoveAdditionalQDQClustersPass()]
     )
 
     logging.debug(f"Lowered graph:\n{edge_program_manager.exported_program().graph}")

exir/BUCK

@@ -267,7 +267,9 @@ fbcode_target(_kind = runtime.python_library,
     deps = [
         "fbsource//third-party/pypi/typing-extensions:typing-extensions",
         ":error",
+        ":pass_base",
         "//caffe2:torch",
+        "//executorch/exir/program:edge_program_manager_pass_base",
     ],
 )
 

exir/pass_manager.py

@@ -5,17 +5,49 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
-
-from typing import Callable, List, Optional, Union
+import inspect
+import logging
+from typing import Callable, List, Optional, TYPE_CHECKING, TypeAlias, Union
 
 import torch
 import torch.fx.passes.infra.pass_manager as fx
 import torch.utils._pytree as pytree
 from executorch.exir.error import ExportError, ExportErrorType
+
 from torch.fx.passes.infra.pass_base import PassResult
-from typing_extensions import TypeAlias
 
-PassType: TypeAlias = Callable[[torch.fx.GraphModule], Optional[PassResult]]
+if TYPE_CHECKING:
+    from executorch.exir.program._program import EdgeProgramManager
+    from executorch.exir.program.edge_program_manager_pass_base import (
+        EdgeProgramManagerPassBase,
+        EdgeProgramManagerPassResult,
+        ExportedProgramPassBase,
+    )
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+PassType: TypeAlias = Union[
+    "EdgeProgramManagerPassBase",
+    "ExportedProgramPassBase",
+    Callable[[torch.fx.GraphModule], Optional[PassResult]],
+]
+
+
+def _get_pass_name(fn: PassType) -> str:
+    from executorch.exir.program.edge_program_manager_pass_base import (
+        ExportedProgramToEdgeProgramManagerPassWrapper,
+        GraphModuleBackedExportedProgramPassWrapper,
+        GraphModuleToEdgeProgramManagerPassWrapper,
+    )
+
+    if isinstance(fn, ExportedProgramToEdgeProgramManagerPassWrapper):
+        return _get_pass_name(fn._pass)
+    if isinstance(fn, GraphModuleToEdgeProgramManagerPassWrapper):
+        return _get_pass_name(fn._inner._pass)
+    if isinstance(fn, GraphModuleBackedExportedProgramPassWrapper):
+        return _get_pass_name(fn._pass)
+    return fn.__name__ if inspect.isfunction(fn) else type(fn).__name__
 
 
 class PassManager(fx.PassManager):
@@ -27,6 +59,8 @@ class PassManager(fx.PassManager):
         * **passes**: A list of callable passes
         * **params**: An instance of PassManagerParams containing the result of the
             flags set in the constructor.
+
+    Note: This class is deprecated. Please use EdgeProgramManagerPassManager instead.
     """
 
     def __init__(
@@ -41,6 +75,9 @@ def __init__(
             enable_debug_pass: set to true to enable the debug passes
             run_checks_after_each_pass: whether to run checks and linting after each pass
         """
+        logger.warning(
+            "PassManager is deprecated. Please use EdgeProgramManagerPassManager instead."
+        )
 
         # Flatten the passes to a list of callables
         passes = passes if passes else []
@@ -76,3 +113,138 @@ def check(self, module: torch.nn.Module) -> None:
                     ExportErrorType.NOT_SUPPORTED,
                     f"call_method `{node}` is not supported except for backend delegate.",
                 )
+
+
+class EdgeProgramManagerPassManager:
+    """
+    Runs multiple passes on an EdgeProgramManager.
+
+    This PassManager accepts passes at three levels of abstraction:
+    - EdgeProgramManagerPassBase: operates on the full EdgeProgramManager
+    - ExportedProgramPassBase: operates on individual ExportedPrograms
+    - Callable[[GraphModule], PassResult]: operates on GraphModules
+
+    Lower-level passes are automatically wrapped to operate at the EPM level.
+    The iteration over methods within an EPM is handled by the wrapper passes,
+    not by this pass manager.
+    """
+
+    def __init__(
+        self,
+        passes: Optional[Union[List[PassType], List[List[PassType]]]] = None,
+        constraints: Optional[List[Callable[[Callable, Callable], bool]]] = None,
+        run_checks_after_each_pass: bool = False,
+        suppress_exported_program_pre_verification: bool = True,
+        steps: int = 1,
+    ) -> None:
+        from executorch.exir.program.edge_program_manager_pass_base import (
+            EdgeProgramManagerPassBase,
+            ExportedProgramPassBase,
+            ExportedProgramToEdgeProgramManagerPassWrapper,
+            GraphModuleToEdgeProgramManagerPassWrapper,
+        )
+
+        wrapped_passes: List[EdgeProgramManagerPassBase] = []
+        for fn in pytree.tree_flatten(passes)[0] if passes else []:
+            if isinstance(fn, EdgeProgramManagerPassBase):
+                wrapped_passes.append(fn)
+            elif isinstance(fn, ExportedProgramPassBase):
+                wrapped_passes.append(
+                    ExportedProgramToEdgeProgramManagerPassWrapper(fn)
+                )
+            else:
+                wrapped_passes.append(
+                    GraphModuleToEdgeProgramManagerPassWrapper(fn)
+                )
+
+        if suppress_exported_program_pre_verification:
+            logger.warning(
+                "Pre-verification of exported program is suppressed. This means that the exported program may pass validation prior to running the pass manager."
+            )
+
+        self.passes: List["EdgeProgramManagerPassBase"] = wrapped_passes
+        self.constraints = constraints
+        self.run_checks_after_each_pass = run_checks_after_each_pass
+        self.suppress_check_failures = suppress_exported_program_pre_verification
+        self.steps = steps
+        self._validated = False
+
+    def solve_constraints(self) -> None:
+        """Placeholder for constraint solving. Marks the pass manager as validated."""
+        self._validated = True
+
+    def check(self, epm: "EdgeProgramManager") -> None:
+        """
+        Runs exported program validation on each method in the EdgeProgramManager.
+        """
+        for name, program in epm._edge_programs.items():  # noqa: B007
+            if not self.suppress_check_failures:
+                program.validate()
+
+            module = program.graph_module
+            module.recompile()
+            module.graph.lint()
+
+            for node in module.graph.nodes:
+                if node.op == "call_method":
+                    raise ExportError(
+                        ExportErrorType.NOT_SUPPORTED,
+                        f"call_method `{node}` is not supported except for backend delegate.",
+                    )
+
+    def __call__(
+        self, epm: "EdgeProgramManager"
+    ) -> "EdgeProgramManagerPassResult":
+        """
+        Runs passes on an EdgeProgramManager.
+
+        Args:
+            epm: The EdgeProgramManager to transform.
+
+        Returns:
+            EdgeProgramManagerPassResult containing the transformed EPM and whether
+            or not it was modified.
+        """
+        from executorch.exir.program.edge_program_manager_pass_base import (
+            EdgeProgramManagerPassResult,
+        )
+
+        if not self._validated:
+            self.solve_constraints()
+
+        self.check(epm)
+
+        overall_modified = False
+
+        for _ in range(self.steps):
+            step_modified = False
+
+            for i, fn in enumerate(self.passes):
+                try:
+                    result = fn(epm)
+                    if result.modified:
+                        logger.debug(
+                            "EPM after pass '%s'",
+                            _get_pass_name(fn),
+                        )
+
+                    epm = result.edge_program_manager
+                    step_modified = step_modified or result.modified
+
+                    if self.run_checks_after_each_pass:
+                        self.check(epm)
+
+                except Exception as e:
+                    prev_names = [_get_pass_name(p) for p in self.passes[:i]]
+                    msg = (
+                        f"An error occurred when running the '{_get_pass_name(fn)}' pass "
+                        f"after the following passes: {prev_names}\n"
+                        f"Original error: {e}"
+                    )
+                    raise type(e)(msg) from e
+
+            overall_modified = overall_modified or step_modified
+            if not step_modified:
+                break
+
+        return EdgeProgramManagerPassResult(epm, overall_modified)

exir/program/BUCK

@@ -1,6 +1,5 @@
 load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target")
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
-load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
 oncall("executorch")
 
@@ -15,6 +14,17 @@ fbcode_target(_kind = runtime.python_library,
     ],
 )
 
+fbcode_target(_kind = runtime.python_library,
+    name = "edge_program_manager_pass_base",
+    srcs = [
+        "edge_program_manager_pass_base.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/exir:pass_base",
+    ],
+)
+
 fbcode_target(_kind = runtime.python_library,
     name = "program",
     srcs = [
@@ -47,6 +57,7 @@ fbcode_target(_kind = runtime.python_library,
         "//executorch/exir/passes:spec_prop_pass",
         "//executorch/exir/passes:weights_to_outputs_pass",
         "//executorch/exir/passes:convert_constant_dim_order_pass",
+        "//executorch/exir/program:edge_program_manager_pass_base",
         "//executorch/exir/verification:verifier",
         "//executorch/extension/flat_tensor/serialize:serialize",
     ] +  (["//executorch/exir/program/fb:logger"] if not runtime.is_oss else [])