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 add a new pass manager which operates on exported programs. This is done to ensure backwards compatibility, while allowing _program transformations to use either pass manager

Differential Revision: D91725222

Author: Andrew Grebenisan

exir/pass_base.py

@@ -6,14 +6,16 @@
 # 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,
     Dict,
+    Final,
     List,
     MutableMapping,
     Optional,
@@ -27,16 +29,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 +158,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 +443,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: Final[ExportPassType] = pass_type
 
     def _fx(
         self,
@@ -651,6 +701,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):