ExportedProgram passes (#16986)

Summary:

- Adds support to _ExportPassBase to run passes on ExportedPrograms. Ensures that we can only run either call or call_exported_program, not both
- Updates exir.PassManager to run on either exported programs or graph modules. Updates PassManagers which extend it to override the correct interfaces (minimal change, just a method renaming)
- If we run the pass manager with an exported program, supports both graph module and exported program passes, but will always return an ExportedProgramPassResult
- Updates transform to always call pass manager with an ExportedProgram, thus getting back an ExportedProgramPassResult.

Differential Revision: D91725222

Author: Andrew Grebenisan

backends/arm/_passes/arm_pass_manager.py

@@ -8,8 +8,10 @@
 import logging
 from collections import defaultdict
 from collections.abc import Sequence
+from typing import Union
 
 import executorch.backends.arm.tosa.dialect  # noqa: unused
+import torch
 from executorch.backends.arm._passes import (
     AccumulateIndexPutPass,
     AnnotateOutputDimOrderPass,
@@ -133,11 +135,10 @@
     TosaSpecification,
 )
 from executorch.exir import ExportedProgram
-from executorch.exir.pass_base import ExportPass
+from executorch.exir.pass_base import ExportedProgramPassResult, ExportPass
 from executorch.exir.pass_manager import PassManager
 from torch.fx import GraphModule
 from torch.fx.passes.infra.pass_base import PassResult
-from torch.nn.modules import Module
 
 logger = logging.getLogger(__name__)
 
@@ -452,9 +453,11 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
 
         return self._transform(graph_module)
 
-    def __call__(self, module: Module) -> PassResult:
+    def __call__(
+        self, module_or_program: Union[torch.fx.GraphModule, ExportedProgram]
+    ) -> Union[PassResult, ExportedProgramPassResult]:
         try:
-            return super().__call__(module)
+            return super().__call__(ExportedProgramPassResult)
         except Exception as e:
             first_exception = e.__cause__ or e.__context__ or e
             import re

backends/nxp/aten_passes/neutron_aten_pass_manager.py

@@ -6,7 +6,6 @@
 from typing import Callable
 
 import torch
-
 from executorch.backends.nxp.aten_passes.convert_unsqueeze_to_view import (
     ConvertUnsqueezeToViewPass,
 )
@@ -36,7 +35,7 @@
 )
 from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
 from executorch.exir.pass_manager import PassManager
-from torch import nn
+from torch.fx import GraphModule
 from torch.fx.passes.infra.pass_base import PassResult
 
 PassType = type[Callable[[torch.fx.GraphModule], PassResult]]
@@ -70,7 +69,7 @@ def __init__(
         passes: list[PassType] = passes or _get_default_passes(neutron_target_spec)
         super().__init__(passes)
 
-    def __call__(self, module: nn.Module) -> PassResult:
+    def run_passes(self, module: GraphModule) -> PassResult:
         pass_result: PassResult = super().__call__(module)
 
         graph_module = pass_result.graph_module

backends/nxp/edge_passes/neutron_edge_pass_manager.py

@@ -11,7 +11,7 @@
 from executorch.backends.nxp.edge_passes.remove_as_strided_copy_nodes import (
     RemoveUselessAsStridedCopyNodes,
 )
-from torch.fx.passes.infra.pass_manager import PassManager
+from executorch.exir.pass_manager import PassManager
 
 
 class NeutronEdgePassManager(PassManager):

exir/pass_base.py

@@ -6,10 +6,11 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
-
+import enum
 import operator
 import traceback
 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,12 +157,60 @@ class ExportPassBaseError(RuntimeError):
     pass
 
 
+@dataclass(frozen=True)
+class ExportedProgramPassResult:
+    exported_program: ExportedProgram
+    modified: bool
+
+
+class ExportPassType(enum.Enum):
+    EXPORTED_PROGRAM = enum.auto()
+    GRAPH_MODULE = enum.auto()
+
+
 class _ExportPassBase(PassBase):
     """
     Interpreter-based pass class to help users maintain the IR spec while writing
     transformations.
     """
 
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        # Required to ensure that subclasses of ExportPassBase only call
+        # the correct call method based on the pass_type
+        super().__init_subclass__(**kwargs)
+        if "call" in cls.__dict__:
+            original_call = cls.__dict__["call"]
+
+            def wrapped_call(
+                self: "_ExportPassBase", graph_module: torch.fx.GraphModule
+            ) -> PassResult:
+                if self.pass_type != ExportPassType.GRAPH_MODULE:
+                    raise ExportPassBaseError(
+                        f"Cannot call 'call' on a pass with pass_type={self.pass_type}. "
+                        f"Expected pass_type=ExportPassType.GRAPH_MODULE. "
+                        f"Use 'call_exported_program' for passes with ExportPassType.EXPORTED_PROGRAM."
+                    )
+
+                return original_call(self, graph_module)
+
+            cls.call = wrapped_call
+
+        if "call_exported_program" in cls.__dict__:
+            original_call_ep = cls.__dict__["call_exported_program"]
+
+            def wrapped_call_exported_program(
+                self: "_ExportPassBase", exported_program: ExportedProgram
+            ) -> ExportedProgramPassResult:
+                if self.pass_type != ExportPassType.EXPORTED_PROGRAM:
+                    raise ExportPassBaseError(
+                        f"Cannot call 'call_exported_program' on a pass with pass_type={self.pass_type}. "
+                        f"Expected pass_type=ExportPassType.EXPORTED_PROGRAM. "
+                        f"Use 'call' for passes with ExportPassType.GRAPH_MODULE."
+                    )
+                return original_call_ep(self, exported_program)
+
+            cls.call_exported_program = wrapped_call_exported_program
+
     @staticmethod
     def _create_dummy_node_metadata() -> NodeMetadata:
         return NodeMetadata({"stack_trace": "".join(traceback.format_stack(limit=1))})
@@ -394,14 +442,15 @@ def run_node(self, n: torch.fx.Node) -> Argument:
             self.callback.node_debug_str = n.format_node()
             return super().run_node(n)
 
-    def __init__(self) -> None:
+    def __init__(self, pass_type: ExportPassType = ExportPassType.GRAPH_MODULE) -> None:
         self.interpreter = torch.fx.Interpreter(
             torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph())
         )
         self.tracer = self.ExportTracer(self, CodeGen())  # pyre-ignore
         self.fake_tensor_mode: Optional[FakeTensorMode] = None
         self._initialized = True
         self.node_debug_str: Optional[str] = None
+        self.pass_type = pass_type
 
     def _fx(
         self,
@@ -651,6 +700,11 @@ def call(self, graph_module: fx.GraphModule) -> PassResult:
 
         return result
 
+    def call_exported_program(
+        self, exported_program: ExportedProgram
+    ) -> ExportedProgramPassResult:
+        raise NotImplementedError("call_exported_program is not implemented.")
+
 
 class ExportPass(_ExportPassBase):
     class ExportTracer(_ExportPassBase.ExportTracer):