New exported program pass manager and exported program passes (#16986)

Differential Revision: D91725222

Pull Request resolved: #16986

Author: DrJessop

backends/arm/test/tester/test_pipeline.py

@@ -48,7 +48,7 @@
 from executorch.backends.arm.vgf.compile_spec import VgfCompileSpec
 from executorch.backends.test.harness.stages import StageType
 from executorch.exir.pass_base import ExportPass
-from torch._export.pass_base import PassType
+from executorch.exir.pass_manager import PassType
 from torch.export.graph_signature import InputKind, OutputKind
 from torchao.quantization.pt2e.quantizer import QuantizationSpec
 

backends/qualcomm/_passes/recompose_pad_maxpool2d.py

@@ -13,12 +13,8 @@
 from executorch.exir.pass_base import ExportPass, PassResult
 from executorch.exir.passes import dead_code_elimination_pass
 
-from torch._subclasses.fake_tensor import FakeTensorMode
-
-
-def add_fake_tensor_to_node(padding_node, input_shape, padding_args, dtype):
-    fake_mode = FakeTensorMode()
 
+def add_fake_tensor_to_node(padding_node, input_shape, padding_args, dtype, fake_mode):
     with fake_mode:
         batch, channels, height, width = input_shape
         pad_left, pad_right, pad_top, pad_bottom = padding_args
@@ -114,6 +110,7 @@ def call(self, graph_module: torch.fx.GraphModule):  # noqa C901
                         input_node.meta["val"].shape,
                         padding,
                         input_node.meta["val"].dtype,
+                        input_node.meta["val"].fake_mode,
                     )
                     if quant_attrs:
                         padding_node.meta["quant_attrs"] = node.meta["quant_attrs"]

backends/qualcomm/_passes/utils.py

@@ -137,7 +137,23 @@ def copy_nn_module_stack(src, target):
         target.meta["nn_module_stack"] = value
 
 
-def merge_decomposed_graph(
+def _unify_fake_mode(node: torch.fx.Node, fake_mode) -> None:
+    val = node.meta.get("val")
+    if val is None:
+        return
+    if isinstance(val, FakeTensor) and val.fake_mode is not fake_mode:
+        node.meta["val"] = fake_mode.from_tensor(val)
+    elif isinstance(val, (list, tuple)):
+        unified = []
+        for v in val:
+            if isinstance(v, FakeTensor) and v.fake_mode is not fake_mode:
+                unified.append(fake_mode.from_tensor(v))
+            else:
+                unified.append(v)
+        node.meta["val"] = type(val)(unified)
+
+
+def merge_decomposed_graph(  # noqa: C901
     remap: Dict[str, torch.fx.Node],
     target_node: torch.fx.Node,
     target_graph: torch.fx.GraphModule,
@@ -148,6 +164,16 @@ def merge_decomposed_graph(
         [torch.fx.Node, torch.fx.Node, Dict[str, torch.fx.Node]], None
     ] = None,
 ) -> None:
+    target_fake_mode = None
+    target_val = target_node.meta.get("val")
+    if isinstance(target_val, FakeTensor):
+        target_fake_mode = target_val.fake_mode
+    elif isinstance(target_val, (list, tuple)):
+        for v in target_val:
+            if isinstance(v, FakeTensor):
+                target_fake_mode = v.fake_mode
+                break
+
     def default_output_process(node):
         for user in node.users.copy():
             # remap
@@ -170,10 +196,13 @@ def default_output_process(node):
                 # replace node map from string to graph node
                 remap[decomposed_node] = remap.pop(decomposed_node.name)
             else:
-                remap[decomposed_node] = target_graph.node_copy(
+                copied = target_graph.node_copy(
                     decomposed_node,
                     arg_transform=lambda x, remap=remap: remap[x],
                 )
+                if target_fake_mode is not None:
+                    _unify_fake_mode(copied, target_fake_mode)
+                remap[decomposed_node] = copied
 
 
 def is_float_tensor(node: torch.fx.Node) -> bool:

exir/BUCK

@@ -259,14 +259,26 @@ fbcode_target(_kind = runtime.python_library,
     ],
 )
 
+fbcode_target(_kind = runtime.python_library,
+    name = "_program_utils",
+    srcs = [
+        "_program_utils.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+    ],
+)
+
 fbcode_target(_kind = runtime.python_library,
     name = "pass_manager",
     srcs = [
         "pass_manager.py",
     ],
     deps = [
         "fbsource//third-party/pypi/typing-extensions:typing-extensions",
+        ":_program_utils",
         ":error",
+        ":pass_base",
         "//caffe2:torch",
     ],
 )

exir/_program_utils.py

@@ -0,0 +1,104 @@
+# 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-unsafe
+
+import torch
+from torch.export.exported_program import (
+    ConstantArgument,
+    ExportGraphSignature,
+    InputSpec,
+    OutputSpec,
+)
+
+
+def _get_updated_range_constraints(gm):
+    def get_shape_env(gm):
+        vals = [
+            node.meta["val"]
+            for node in gm.graph.nodes
+            if node.meta.get("val", None) is not None
+        ]
+        from torch._guards import detect_fake_mode  # type: ignore[21]
+
+        fake_mode = detect_fake_mode(vals)
+        if fake_mode is not None:
+            return fake_mode.shape_env
+        for v in vals:
+            if isinstance(v, torch.SymInt):
+                return v.node.shape_env
+
+    shape_env = get_shape_env(gm)
+    if shape_env is None:
+        return {}
+    range_constraints = {
+        shape_env.replacements.get(k, k): v for k, v in shape_env.var_to_range.items()
+    }
+    # Only when we have an unbacked symint, and it's used as constructor inputs,
+    # runtime_var_to_range will make a difference compated to var_to_range.
+    # e.g. [2, oo) -> [0, oo)
+    for k, v in shape_env.var_to_range.items():
+        if k not in shape_env.replacements:
+            range_constraints[k] = v
+    return range_constraints
+
+
+def _get_updated_graph_signature(
+    old_signature: ExportGraphSignature,
+    new_gm: torch.fx.GraphModule,
+) -> ExportGraphSignature:
+    """
+    Update the graph signature's user_input/user_outputs.
+    """
+    new_input_specs = []
+    i = 0
+    for node in new_gm.graph.nodes:
+        if node.op != "placeholder":
+            continue
+
+        assert i < len(
+            old_signature.input_specs
+        ), "Number of inputs changed after transformation"
+        old_input_spec = old_signature.input_specs[i]
+        arg = (
+            old_input_spec.arg
+            if isinstance(old_input_spec.arg, ConstantArgument)
+            # pyre-fixme[20]: Argument `class_fqn` expected.
+            else type(old_input_spec.arg)(node.name)
+        )
+        new_input_specs.append(
+            InputSpec(
+                old_input_spec.kind,
+                arg,
+                old_input_spec.target,
+                persistent=old_input_spec.persistent,
+            )
+        )
+        i += 1
+
+    output_node = new_gm.graph.output_node()
+    assert output_node.op == "output"
+
+    new_output_specs = []
+    for i, node in enumerate(output_node.args[0]):
+        assert i < len(
+            old_signature.output_specs
+        ), "Number of outputs changed after transformation"
+        old_output_spec = old_signature.output_specs[i]
+        arg = (
+            old_output_spec.arg
+            if isinstance(old_output_spec.arg, ConstantArgument)
+            # pyre-fixme[20]: Argument `class_fqn` expected.
+            else type(old_output_spec.arg)(node.name)
+        )
+        new_output_specs.append(
+            OutputSpec(old_output_spec.kind, arg, old_output_spec.target)
+        )
+
+    new_signature = ExportGraphSignature(
+        input_specs=new_input_specs, output_specs=new_output_specs
+    )
+    return new_signature

exir/pass_base.py

@@ -6,10 +6,11 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
-
 import operator
 import traceback
+from abc import ABC, abstractmethod
 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
@@ -182,6 +182,58 @@ class ExportPassBaseError(RuntimeError):
     pass
 
 
+@dataclass(frozen=True)
+class ExportedProgramPassResult:
+    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
+        """
+
+
 class _ExportPassBase(PassBase):
     """
     Interpreter-based pass class to help users maintain the IR spec while writing