EdgeProgramManager passes (pytorch#16986)

Summary:
Pull Request resolved: pytorch#16986

- Adds support to run to run passes on ExportedPrograms and EdgeProgramManager
- EdgeProgramManager transform behaves basically like a pass manager

Reviewed By: larryliu0820, ethansfng

Differential Revision: D91725222

Author: Andrew Grebenisan

exir/BUCK

@@ -259,6 +259,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 = "pass_manager",
     srcs = [
@@ -267,6 +278,7 @@ fbcode_target(_kind = runtime.python_library,
     deps = [
         "fbsource//third-party/pypi/typing-extensions:typing-extensions",
         ":error",
+        ":pass_base",
         "//caffe2:torch",
     ],
 )

exir/edge_program_manager_pass_base.py

@@ -0,0 +1,313 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import copy
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Sequence, TYPE_CHECKING, Union
+
+import torch
+from torch.export import ExportedProgram
+from torch.fx.passes.infra.pass_base import PassResult
+
+if TYPE_CHECKING:
+    from executorch.exir.program._program import EdgeProgramManager
+
+
+@dataclass(frozen=True)
+class ExportedProgramPassResult:
+    """Result of running a pass on an ExportedProgram."""
+
+    exported_program: ExportedProgram
+    modified: bool
+
+
+class ExportedProgramPassBase(ABC):
+    """
+    Base interface for implementing passes that operate on ExportedProgram.
+    """
+
+    def __call__(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
+        """
+        Runs the precondition check, the pass itself, and the postcondition check.
+        """
+
+        self.requires(exported_program)
+        res = self.call(exported_program)
+        self.ensures(exported_program)
+        return res
+
+    @abstractmethod
+    def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
+        """
+        The pass that is run through the given exported program. To implement a
+        pass, it is required to implement this function.
+
+        Args:
+            exported_program: The exported program we will run a pass on
+        """
+
+    def requires(self, exported_program: ExportedProgram) -> None:  # noqa: B027
+        """
+        This function will be called before the pass is run and will check that
+        the given exported program contains the preconditions needed to run the
+        pass. It is not required to implement this function.
+
+        Args:
+            exported_program: The exported program we will run checks on
+        """
+
+    def ensures(self, exported_program: ExportedProgram) -> None:  # noqa: B027
+        """
+        This function will be called after the pass is run and will check that
+        the given exported program contains the postconditions needed to run the
+        pass. It is not required to implement this function.
+
+        Args:
+            exported_program: The exported program we will run checks on
+        """
+
+
+@dataclass(frozen=True)
+class EdgeProgramManagerPassResult:
+    """Result of running a pass on an EdgeProgramManager."""
+
+    edge_program_manager: "EdgeProgramManager"
+    modified: bool
+
+
+class EdgeProgramManagerPassBase(ABC):
+    """
+    Base interface for implementing passes that operate on EdgeProgramManager.
+
+    This is the highest-level pass abstraction. Passes at this level can:
+    - Transform individual ExportedPrograms within the manager
+    - Modify constant methods
+    - Split one program into multiple programs
+    - Add or remove programs from the manager
+
+    Lower-level passes (ExportedProgramPassBase, GraphModule callables) can be
+    lifted to this level using the provided wrapper classes.
+    """
+
+    def __call__(
+        self, epm: "EdgeProgramManager"
+    ) -> EdgeProgramManagerPassResult:
+        """
+        Runs the precondition check, the pass itself, and the postcondition check.
+        """
+        self.requires(epm)
+        res = self.call(epm)
+        self.ensures(res.edge_program_manager)
+        return res
+
+    @abstractmethod
+    def call(
+        self, epm: "EdgeProgramManager"
+    ) -> EdgeProgramManagerPassResult:
+        """
+        The pass that is run on the given EdgeProgramManager. To implement a
+        pass, it is required to implement this function.
+
+        Args:
+            epm: The EdgeProgramManager to transform
+        """
+
+    def requires(self, epm: "EdgeProgramManager") -> None:  # noqa: B027
+        """
+        This function will be called before the pass is run and will check that
+        the given EdgeProgramManager contains the preconditions needed to run the
+        pass. It is not required to implement this function.
+
+        Args:
+            epm: The EdgeProgramManager we will run checks on
+        """
+
+    def ensures(self, epm: "EdgeProgramManager") -> None:  # noqa: B027
+        """
+        This function will be called after the pass is run and will check that
+        the given EdgeProgramManager contains the postconditions needed to run the
+        pass. It is not required to implement this function.
+
+        Args:
+            epm: The EdgeProgramManager we will run checks on
+        """
+
+
+class GraphModuleBackedExportedProgramPassWrapper(ExportedProgramPassBase):
+    """
+    Wrapper that adapts a GraphModule pass to work as an ExportedProgramPassBase.
+
+    This wrapper takes a pass that operates on GraphModule and makes it compatible
+    with ExportedProgramPassBase by extracting the graph module, running the pass,
+    and updating the ExportedProgram in-place.
+    """
+
+    def __init__(
+        self,
+        graph_module_pass: Callable[[torch.fx.GraphModule], PassResult],
+    ) -> None:
+        super().__init__()
+        self._pass = graph_module_pass
+
+    def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
+        from executorch.exir.program._program import (
+            _get_updated_graph_signature,
+            _get_updated_range_constraints,
+        )
+
+        result = self._pass(exported_program.graph_module)
+
+        if result.modified:
+            # Cannot use _update_exported_program_graph_module because it
+            # runs verification, and it is not the responsibility of the
+            # pass to run verification. EdgeProgram manager can
+            # optionally run verification after a pass.
+            result.graph_module.recompile()
+            exported_program = copy.copy(exported_program)  # bypasses __init__ and _validate()
+
+            exported_program._graph_module = result.graph_module
+            exported_program._graph_signature = _get_updated_graph_signature(
+                exported_program.graph_signature, result.graph_module
+            )
+            exported_program._range_constraints = _get_updated_range_constraints(
+                result.graph_module
+            )
+            exported_program._module_call_graph = copy.deepcopy(
+                exported_program._module_call_graph
+            )
+            exported_program._graph_module.meta.update(exported_program.graph_module.meta)
+
+
+        return ExportedProgramPassResult(exported_program, result.modified)
+
+
+class ExportedProgramToEdgeProgramManagerPassWrapper(EdgeProgramManagerPassBase):
+    """
+    Adapts an ExportedProgramPassBase to run on every method in an EdgeProgramManager.
+
+    This wrapper takes a pass that operates on a single ExportedProgram and applies it
+    to every method in the EdgeProgramManager, collecting results into a new EPM.
+    This is where the iteration over methods lives -- not in the pass manager, and not
+    in EdgeProgramManager.transform().
+    """
+
+    def __init__(self, ep_pass: ExportedProgramPassBase) -> None:
+        super().__init__()
+        self._pass = ep_pass
+
+    def call(
+        self, epm: "EdgeProgramManager"
+    ) -> EdgeProgramManagerPassResult:
+        new_epm = copy.copy(epm)
+        new_epm._edge_programs = dict(epm._edge_programs)
+
+        overall_modified = False
+        for name, program in epm._edge_programs.items():
+            result = self._pass(program)
+            new_epm._edge_programs[name] = result.exported_program
+            overall_modified = overall_modified or result.modified
+
+        new_epm._config_methods = epm._config_methods
+        return EdgeProgramManagerPassResult(new_epm, overall_modified)
+
+
+PassType = Union[
+    EdgeProgramManagerPassBase,
+    ExportedProgramPassBase,
+    Callable[[torch.fx.GraphModule], Optional[PassResult]],
+]
+
+# Passes that operate on a single method (ExportedProgram or GraphModule level).
+# Excludes EdgeProgramManagerPassBase, which operates on the whole EdgeProgramManager.
+# Use this for per-method pass specifications (e.g. Dict[str, Sequence[MethodPassType]]).
+MethodPassType = Union[
+    ExportedProgramPassBase,
+    Callable[[torch.fx.GraphModule], Optional[PassResult]],
+]
+
+
+def _get_pass_name(fn: PassType) -> str:
+    """Unwraps wrapper chain to get the underlying pass name."""
+    import inspect
+
+    if isinstance(fn, ExportedProgramToEdgeProgramManagerPassWrapper):
+        return _get_pass_name(fn._pass)
+    if isinstance(fn, GraphModuleBackedExportedProgramPassWrapper):
+        return _get_pass_name(fn._pass)
+    return fn.__name__ if inspect.isfunction(fn) else type(fn).__name__
+
+
+def wrap_passes(
+    passes: Sequence[PassType],
+) -> list[EdgeProgramManagerPassBase]:
+    """
+    Wraps a list of mixed-level passes up to the EdgeProgramManager level.
+
+    Accepts passes at three levels:
+    - EdgeProgramManagerPassBase: used as-is
+    - ExportedProgramPassBase: wrapped with ExportedProgramToEdgeProgramManagerPassWrapper
+    - GraphModule callables: wrapped with GraphModuleBackedExportedProgramPassWrapper
+      then ExportedProgramToEdgeProgramManagerPassWrapper
+
+    Args:
+        passes: A sequence of passes at any level.
+
+    Returns:
+        A list of EdgeProgramManagerPassBase passes.
+    """
+    from torch.fx.passes.infra.pass_manager import pass_result_wrapper
+
+    wrapped: list[EdgeProgramManagerPassBase] = []
+    for fn in passes:
+        if isinstance(fn, EdgeProgramManagerPassBase):
+            wrapped.append(fn)
+        elif isinstance(fn, ExportedProgramPassBase):
+            wrapped.append(
+                ExportedProgramToEdgeProgramManagerPassWrapper(fn)
+            )
+        else:
+            assert callable(fn)
+            ep_pass = GraphModuleBackedExportedProgramPassWrapper(
+                pass_result_wrapper(fn)
+            )
+            wrapped.append(
+                ExportedProgramToEdgeProgramManagerPassWrapper(ep_pass)
+            )
+    return wrapped
+
+
+class MethodFilteredEdgeProgramManagerPass(EdgeProgramManagerPassBase):
+    """
+    Applies different passes to different methods in an EdgeProgramManager.
+
+    Converts the Dict[str, Sequence[MethodPassType]] pattern (previously handled inline
+    in EdgeProgramManager.transform) into a proper pass. Used by
+    to_edge_transform_and_lower to handle the dict case.
+    """
+
+    def __init__(self, passes_dict: Dict[str, Sequence[MethodPassType]]) -> None:
+        super().__init__()
+        self._passes_dict = passes_dict
+
+    def call(
+        self, epm: "EdgeProgramManager"
+    ) -> EdgeProgramManagerPassResult:
+        from executorch.exir.program._program import _transform
+
+        new_epm = copy.copy(epm)
+        new_epm._edge_programs = dict(epm._edge_programs)
+
+        overall_modified = False
+        for name, program in epm._edge_programs.items():
+            if name in self._passes_dict:
+                new_program = _transform(program, *self._passes_dict[name])
+                new_epm._edge_programs[name] = new_program
+                overall_modified = True
+
+        return EdgeProgramManagerPassResult(new_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")
 
@@ -47,6 +46,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: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 [])