ExportedProgram passes (#16986)

Summary:

- Add ExportedProgramPassBase, which supports running passes on both `fx.GraphModule`s and `exir.ExportedProgram`s
- Extend this class in `_ExportPassBase` so that any passes which import directly from `ExportPass` are immediately compatible with the current pass manager.
- Create `LegacyPassWrapper` to auto-wrap passes not migrated to use ExportedProgramPassBase such that it works out of the box with EdgeProgramManager.

Differential Revision: D91725222

Author: Andrew Grebenisan

exir/pass_base.py

@@ -6,10 +6,11 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
-
 import operator
 import traceback
+import warnings
 from contextlib import nullcontext
+from dataclasses import dataclass
 from typing import (
     Any,
     Callable,
@@ -27,16 +28,15 @@
 
 import torch
 from executorch.exir import memory
-
 from executorch.exir.delegate import executorch_call_delegate, is_lowered_module
-
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.error import ExportError, ExportErrorType
 from torch import fx
 from torch._dispatch.python import enable_python_dispatcher
 from torch._subclasses import FakeTensorMode, UnsupportedFakeTensorException
 from torch._subclasses.fake_tensor import FakeTensor
 from torch._subclasses.functional_tensor import FunctionalTensor, FunctionalTensorMode
+from torch.export import ExportedProgram
 from torch.fx import traceback as fx_traceback
 from torch.fx.experimental.proxy_tensor import PythonKeyTracer
 from torch.fx.graph import CodeGen
@@ -157,7 +157,64 @@ class ExportPassBaseError(RuntimeError):
     pass
 
 
-class _ExportPassBase(PassBase):
+@dataclass(frozen=True)
+class ExportedProgramPassResult:
+    exported_program: ExportedProgram
+    modified: bool
+
+
+class ExportedProgramPassBase(PassBase):
+    """
+    Base class which supports running passes on either `exir.ExportedProgram` or `fx.GraphModule`
+    types.
+    """
+
+    def call_exported_program(
+        self, exported_program: ExportedProgram
+    ) -> ExportedProgramPassResult:
+        pass_result = self.call(exported_program.graph_module)
+        modified = pass_result.modified if pass_result else False
+        if modified:
+            exported_program._graph_module = pass_result.graph_module
+        return ExportedProgramPassResult(exported_program, modified)
+
+    def call(self, graph_module: fx.GraphModule) -> Optional[PassResult]:
+        """
+        Overriden version of call which is effectively a no-op, but is required
+        since the base class method is abstract.
+        """
+        return PassResult(graph_module, False)
+
+
+class LegacyPassWrapper(ExportedProgramPassBase):
+    """
+    Wraps a legacy callable pass (Callable[[GraphModule], Optional[PassResult]])
+    to work with the new ExportedProgramPassBase infrastructure.
+
+    This provides backwards compatibility for passes that haven't been migrated yet.
+    """
+
+    def __init__(
+        self, legacy_pass: Callable[[torch.fx.GraphModule], Optional[PassResult]]
+    ) -> None:
+        self._legacy_pass = legacy_pass
+        # Preserve the original name for debugging/logging
+        self.__class__.__name__ = getattr(
+            legacy_pass, "__name__", type(legacy_pass).__name__
+        )
+
+        warnings.warn(
+            f"Callable pass '{self.__class__.__name__}' is deprecated. "
+            "Please migrate to ExportedProgramPassBase.",
+            DeprecationWarning,
+            stacklevel=3,
+        )
+
+    def call(self, graph_module: fx.GraphModule) -> Optional[PassResult]:
+        return self._legacy_pass(graph_module)
+
+
+class _ExportPassBase(ExportedProgramPassBase):
     """
     Interpreter-based pass class to help users maintain the IR spec while writing
     transformations.

exir/pass_manager.py

@@ -5,17 +5,38 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
-
-from typing import Callable, List, Optional, Union
+import copy
+from typing import Callable, cast, List, Optional, 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 executorch.exir.pass_base import (
+    ExportedProgramPassBase,
+    ExportedProgramPassResult,
+    LegacyPassWrapper,
+)
+from torch.export import ExportedProgram
 from torch.fx.passes.infra.pass_base import PassResult
-from typing_extensions import TypeAlias
 
-PassType: TypeAlias = Callable[[torch.fx.GraphModule], Optional[PassResult]]
+# Legacy type for backwards compatibility
+LegacyPassType = Callable[[torch.fx.GraphModule], Optional[PassResult]]
+
+# New union that accepts both
+PassType = Union[ExportedProgramPassBase, LegacyPassType]
+
+
+def _normalize_pass(
+    p: Union[ExportedProgramPassBase, Callable]
+) -> ExportedProgramPassBase:
+    """Normalize a pass to ExportedProgramPassBase, wrapping legacy callables."""
+    if isinstance(p, ExportedProgramPassBase):
+        return p
+    elif callable(p):
+        return LegacyPassWrapper(p)
+    else:
+        raise TypeError(f"Expected ExportedProgramPassBase or callable, got {type(p)}")
 
 
 class PassManager(fx.PassManager):
@@ -48,12 +69,21 @@ def __init__(
             fx.pass_result_wrapper(fn) for fn in pytree.tree_flatten(passes)[0]
         ]
 
+        normalized: list[ExportedProgramPassBase] = []
+        for p in flattened_passes:
+            normalized.append(_normalize_pass(p))
+
+        flattened_passes = normalized
+
         super().__init__(
             flattened_passes,
             run_checks_after_each_pass=run_checks_after_each_pass,
             suppress_check_failures=suppress_check_failures,
         )
 
+        # Improves type-checking in call_exported_program
+        self.passes = cast(list[ExportedProgramPassBase], self.passes)
+
     def check(self, module: torch.nn.Module) -> None:
         """
         Runs various checks on the given graph module to make sure it contains
@@ -76,3 +106,95 @@ def check(self, module: torch.nn.Module) -> None:
                     ExportErrorType.NOT_SUPPORTED,
                     f"call_method `{node}` is not supported except for backend delegate.",
                 )
+
+    def call(self, module: torch.nn.Module) -> PassResult:
+        """
+        Runs the passes on the given graph module.
+
+        Args:
+            module: The graph module to run the passes on
+
+        Returns:
+            A PassResult object containing the modified graph module and a
+            boolean indicating whether the module was modified
+        """
+        # Check for passes that might have ExportedProgram-specific logic
+        problematic_passes = []
+        for p in self.passes:
+            if (
+                type(p).call_exported_program
+                is not ExportedProgramPassBase.call_exported_program
+            ):
+                problematic_passes.append(type(p).__name__)
+
+        if problematic_passes:
+            fx.logger.warning(
+                f"The following passes have overridden 'call_exported_program': {problematic_passes}. "
+                "Calling PassManager.call() will only run the graph module logic via 'call()' "
+                "and may miss ExportedProgram-specific transformations. "
+                "Consider using PassManager.call_exported_program() instead.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        return super()(module)
+
+    def call_exported_program(
+        self, exported_program: ExportedProgram
+    ) -> ExportedProgramPassResult:
+        """
+        Runs the passes on the given ExportedProgram.
+
+        Args:
+            exported_program: The ExportedProgram to run the passes on
+
+        Returns:
+            A PassResult object containing the modified ExportedProgram and a
+            boolean indicating whether the ExportedProgram was modified
+        """
+        # Order the passes based on the constraints
+        if not self._validated:
+            self.solve_constraints()
+
+        # Check graph invariants
+        self.check(exported_program.graph_module)
+
+        # Run the set of passes `steps` number of times or until the graph stops
+        # changing
+        overall_modified = False
+
+        for _ in range(self.steps):
+            modified = False
+
+            # Run the set of passes on the graph module
+            for i, p in enumerate(self.passes):
+                pass_name = type(p).__name__
+                fx.logger.debug("Running pass '%s'", pass_name)
+
+                try:
+                    res = p.call_exported_program(exported_program)
+                    exported_program = res.exported_program
+                    modified = modified or res.modified
+
+                    fx.logger.debug(
+                        "Graph after pass '%s': %s",
+                        pass_name,
+                        exported_program.graph_module.graph,
+                    )
+                    exported_program.graph_module.recompile()
+
+                    # Check graph invariants
+                    if self.run_checks_after_each_pass:
+                        self.check(exported_program.graph_module)
+
+                except Exception as e:
+                    prev_pass_names = [type(p).__name__ for p in self.passes[:i]]
+                    msg = f"An error occurred when running the '{pass_name}' pass after the following passes: {prev_pass_names}"
+                    raise Exception(msg) from e  # noqa: TRY002
+
+            # If the graph no longer changes, then we can stop running these passes
+            overall_modified = overall_modified or modified
+            if not modified:
+                break
+
+        return ExportedProgramPassResult(exported_program, overall_modified)