Seperate lowered backend module as a seperate file

Summary: Previously we define `LoweredBackendModule` and have `executorch_call_delegate` implemented in `delegate.py`. As a preparation for emit `LoweredBackendModule` in D47803806, we may need to separate them to two files due to some circular dependency issues.

Reviewed By: angelayi

Differential Revision: D47809768

fbshipit-source-id: 18101f5be3819935ecd7ece2daa71670121a81fb

Author: cccclai

backends/TARGETS

@@ -103,5 +103,6 @@ runtime.python_library(
     ],
     deps = [
         "//caffe2:torch",
+        "//executorch/exir:lowered_backend_module",
     ],
 )

backends/backend_api.py

@@ -17,13 +17,13 @@
     MultiMethodExirExportedProgram,
 )
 
-from executorch.exir.delegate import (
+from executorch.exir.delegate import executorch_call_delegate, get_lowered_module_name
+
+from executorch.exir.graph_module import get_control_flow_submodules
+from executorch.exir.lowered_backend_module import (
     create_submodule_from_nodes,
-    executorch_call_delegate,
-    get_lowered_module_name,
     LoweredBackendModule,
 )
-from executorch.exir.graph_module import get_control_flow_submodules
 from executorch.exir.pass_base import ExportPass
 from torch._export.exported_program import ExportedProgram
 

backends/test/TARGETS

@@ -100,6 +100,7 @@ python_unittest(
         "//executorch/exir:delegate",
         "//executorch/exir:graph_module",
         "//executorch/exir:lib",
+        "//executorch/exir:lowered_backend_module",
         "//executorch/exir:print_program",
         "//executorch/exir:schema",
         "//executorch/exir/dialects:lib",
@@ -147,6 +148,7 @@ python_unittest(
         "//executorch/exir:delegate",
         "//executorch/exir:graph_module",
         "//executorch/exir:lib",
+        "//executorch/exir:lowered_backend_module",
         "//executorch/kernels/portable:custom_ops_generated_lib",
         "//executorch/kernels/quantized:custom_ops_generated_lib",
         "//executorch/runtime/executor/test:test_backend_compiler_lib",

backends/test/test_backends.py

@@ -25,9 +25,10 @@
 from executorch.backends.test.qnn_backend_demo import QnnBackend
 from executorch.exir import EdgeCompileConfig, multi_method_program_to_executorch
 
-from executorch.exir.delegate import executorch_call_delegate, get_lowered_submodules
+from executorch.exir.delegate import executorch_call_delegate
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.graph_module import get_control_flow_submodules
+from executorch.exir.lowered_backend_module import get_lowered_submodules
 from executorch.exir.print_program import print_program
 from executorch.exir.schema import (
     BackendDelegate,

backends/test/test_backends_nested.py

@@ -17,10 +17,10 @@
     AddOperatorSupport,
     MatmulOperatorSupport,
 )
+from executorch.exir.delegate import executorch_call_delegate
 
-from executorch.exir.delegate import executorch_call_delegate, get_lowered_submodules
 from executorch.exir.graph_module import _get_submodule, get_control_flow_submodules
-
+from executorch.exir.lowered_backend_module import get_lowered_submodules
 from functorch.experimental import control_flow
 from torch.fx.passes.operator_support import any_chain, OperatorSupportBase
 

backends/utils.py

@@ -1,9 +1,9 @@
 from typing import Iterable, List, Tuple
 
 import torch
-
-from executorch.exir.delegate import create_submodule_from_nodes
 from executorch.exir.dialects._ops import ops as exir_ops
+
+from executorch.exir.lowered_backend_module import create_submodule_from_nodes
 from torch.fx.passes.utils.source_matcher_utils import SourcePartition
 
 T_QuantPerTensor = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default

exir/TARGETS

@@ -108,6 +108,17 @@ python_library(
         "delegate.py",
     ],
     deps = [
+        "//caffe2:torch",
+    ],
+)
+
+python_library(
+    name = "lowered_backend_module",
+    srcs = [
+        "lowered_backend_module.py",
+    ],
+    deps = [
+        ":delegate",
         ":graph_module",
         ":tracer",
         "//caffe2:torch",

exir/delegate.py

@@ -1,101 +1,30 @@
 # pyre-strict
 
-import types
-from typing import List, Tuple
+from __future__ import annotations
+
+from typing import Any
 
 import torch
 import torch.utils._pytree as pytree
-
-from executorch.backends.compile_spec_schema import CompileSpec
-from executorch.exir.graph_module import _get_submodule
-from executorch.exir.tracer import Value
-from torch._export.exported_program import ExportedProgram
 from torch._functorch.eager_transforms import (
     _unwrap_all_tensors_from_functional,
     _wrap_all_tensors_to_functional,
 )
 from torch._ops import HigherOrderOperator
-from torch._subclasses import FakeTensor
 from torch._subclasses.fake_tensor import FakeTensorMode
 from torch.fx.experimental.proxy_tensor import (
     disable_proxy_modes_tracing,
     get_proxy_slot,
     ProxyTorchDispatchMode,
     track_tensor_tree,
 )
-from torch.fx.passes.utils.fuser_utils import (
-    erase_nodes,
-    fuse_as_graphmodule,
-    insert_subgm,
-    legalize_graph,
-    NodeList,
-    topo_sort,
-)
 from torch.utils._python_dispatch import (
     _get_current_dispatch_mode,
     _pop_mode_temporarily,
 )
-
 from torch.utils._pytree import tree_flatten
 
 
-class LoweredBackendModule(torch.nn.Module):
-    """
-    A subclass of nn.Module that is generated for modules containing
-    delegated functions. This is can be created by calling `to_backend`.
-
-    Private Attributes:
-        * **backend_id**: The backend's name
-        * **processed_bytes**: The delegate blobs created from backend.preprocess
-        * **compile_specs**: A list of backend-specific objects with static
-            metadata to configure the "compilation" process.
-        * **original_module**: The original EXIR module
-    """
-
-    _backend_id: str
-    _processed_bytes: bytes
-    _compile_specs: List[CompileSpec]
-    _original_module: ExportedProgram
-
-    def __init__(
-        self,
-        edge_program: ExportedProgram,
-        backend_id: str,
-        processed_bytes: bytes,
-        compile_specs: List[CompileSpec],
-    ) -> None:
-        super().__init__()
-        self._original_module = edge_program
-        self._backend_id = backend_id
-        self._processed_bytes = processed_bytes
-        self._compile_specs = compile_specs
-
-    @property
-    def backend_id(self) -> str:
-        return self._backend_id
-
-    @property
-    def processed_bytes(self) -> bytes:
-        return self._processed_bytes
-
-    @property
-    def compile_specs(self) -> List[CompileSpec]:
-        return self._compile_specs
-
-    @property
-    def original_module(self) -> ExportedProgram:
-        return self._original_module
-
-    # Used to patch each delegated function with a call_delegate call
-    # @staticmethod
-    def forward(
-        self,
-        *args: Value,
-        **kwargs: Tuple[Value, ...],
-    ) -> Value:
-        return executorch_call_delegate(self, *args)
-
-
 executorch_call_delegate = HigherOrderOperator(
     "executorch_call_delegate", _deprecated_global_ns=True
 )
@@ -108,6 +37,7 @@ def forward(
 # pyre-ignore
 executorch_call_delegate.fallthrough(torch._C.DispatchKey.AutocastCPU)
 
+LOWERED_BACKEND_MODULE_TYPE = "LoweredBackendModule"
 
 # pyre-ignore
 def trace_call_delegate(proxy_mode, func_overload, lowered_module, *args):
@@ -117,7 +47,7 @@ def _unwrap_proxy(e):
             return e
         return get_proxy_slot(e, proxy_mode.tracer, e, lambda e: e.proxy)
 
-    if not isinstance(lowered_module, LoweredBackendModule):
+    if not is_lowered_module(lowered_module):
         raise ValueError(
             "executorch_call_delegate()'s first argument must be a LoweredBackendModule"
         )
@@ -235,8 +165,18 @@ def call_delegate_functionalize(interpreter, lowered_module, *args):
         return _wrap_all_tensors_to_functional(res, level=interpreter.level())
 
 
+# pyre-ignore: Missing parameter annotation [2]: Parameter `obj` must have a type other than `Any`.Pyre
+def is_lowered_module(obj: Any) -> bool:
+    """
+    This function is added to avoid using isinstance(obj, LoweredBackendModule) as it will import LoweredBackendModule, which may cause a circular import.
+    """
+    return type(obj).__name__ == LOWERED_BACKEND_MODULE_TYPE
+
+
 def get_lowered_module_name(
-    root: torch.nn.Module, lowered_module: LoweredBackendModule
+    root: torch.nn.Module,
+    # pyre-ignore: Undefined or invalid type [11]: Annotation `LoweredBackendModule` is not defined as a type.
+    lowered_module: LOWERED_BACKEND_MODULE_TYPE,  # noqa
 ) -> str:
     """
     Adds the given lowered_module into the given root module and returns the
@@ -254,110 +194,3 @@ def get_lowered_module_name(
 
     root.add_module(qualname, lowered_module)
     return qualname
-
-
-# TODO(zhxchen17) Try ExportPass
-def _fixup_output_node(gm: torch.fx.GraphModule) -> None:
-    for node in reversed(gm.graph.nodes):
-        if node.op == "output":
-            with gm.graph.inserting_before(node):
-                assert len(node.args) == 1
-                outputs = node.args[0]
-                if isinstance(outputs, torch.fx.Node):
-                    val = outputs.meta.get("val")
-                    if isinstance(val, list):
-                        # If a list is returned, in some cases it is represented as a
-                        # singular node, like `split_copy_tensor` but EXIR will return a
-                        # opened-up list like `[getitem1, getitem2]`
-                        outputs = [
-                            torch.fx.Proxy(outputs)[i].node for i in range(len(val))
-                        ]
-            returns, out_spec = pytree.tree_flatten(outputs)
-            node.args = (returns,)
-            return
-
-
-def create_submodule_from_nodes(
-    gm: torch.fx.GraphModule,
-    node_list: NodeList,
-    tag: str,
-    skip_legalize_graph: bool = False,
-) -> Tuple[torch.fx.GraphModule, torch.fx.Node]:
-    """
-    Modifies the given graph module in-place to separate out the given nodes
-    into a submodule. The given node_list should form a fully connected
-    subgraph.
-
-    Args:
-        gm: The graph module that we want to partition
-        node_list: A list of nodes that belong in the partition
-
-    Returns:
-        The submodule that has been partitioned, the call_module node in the
-        toplevel graph module calling the submodule
-    """
-    sorted_nodes = topo_sort(node_list)
-
-    submodule_name = "fused_" + tag
-    sub_gm, orig_inputs, orig_outputs = fuse_as_graphmodule(
-        gm, sorted_nodes, submodule_name
-    )
-
-    _fixup_output_node(sub_gm)
-
-    gm = insert_subgm(gm, sub_gm, orig_inputs, orig_outputs)
-    if len(orig_outputs) == 1 and isinstance(orig_outputs[0].meta["val"], FakeTensor):
-        # If the original output is a single tensor, it has been
-        # pytree.tree_flatten-ed to be a singleton list, so we want to replace
-        # all uses with a getitem call to the 0th index of the result
-        for node in gm.graph.nodes:
-            if node.op == "call_module":
-                with gm.graph.inserting_after(node):
-                    proxy_out = torch.fx.Proxy(node)[0].node  # type: ignore[index]
-                    node.replace_all_uses_with(proxy_out, propagate_meta=True)
-                    # Reset the args since it was overwritten in the previous line
-                    proxy_out.args = (node, 0)
-
-    erase_nodes(gm, sorted_nodes)
-
-    # Topological sort original gm with newly created sub_gm
-    # TODO : T153794167 Get rid of support for skipping legalize graph in create_submodule_from_nodes
-    # once we transition to using fuse_by_partitions.
-    if not skip_legalize_graph:
-        legalize_graph(gm)
-
-    # Get the call_module node
-    submodule_node = None
-    for node in gm.graph.nodes:
-        if node.op == "call_module" and node.target == submodule_name:
-            submodule_node = node
-        elif node.op == "call_module":
-            raise RuntimeError(
-                f"The submodule created with nodes {node_list} did not form \
-                one fully contained subgraph. Check that these nodes form a \
-                fully contained graph. Partitioned graph: {gm.graph}."
-            )
-
-    assert (
-        submodule_node is not None
-    ), f"No submodule was created with the nodes {node_list} in the graph {gm.graph}"
-
-    return sub_gm, submodule_node
-
-
-def get_lowered_submodules(
-    graph_module: torch.fx.GraphModule,
-) -> List[Tuple[str, LoweredBackendModule, torch.fx.Node]]:
-    """
-    Returns a list of lowered modules that are in the given graph (does not look
-    into submodules). Specifically, the returned value is a list containing a
-    tuple of (name of the lowered module that's stored in the graph module, the
-    lowered module itself, and the fx node that called this lowered module).
-    """
-    lowered_submodules = []
-    for node in graph_module.graph.nodes:
-        if node.op == "call_function" and node.target == executorch_call_delegate:
-            name, module, node = _get_submodule(graph_module, node, 0)
-            assert isinstance(module, LoweredBackendModule)
-            lowered_submodules.append((name, module, node))
-    return lowered_submodules

exir/emit/_emitter.py

@@ -35,7 +35,7 @@
 import torch.fx
 from executorch.exir import delegate
 from executorch.exir.common import add_cursor_to_graph
-from executorch.exir.delegate import LoweredBackendModule
+from executorch.exir.delegate import executorch_call_delegate, is_lowered_module
 from executorch.exir.dialects.backend._ops import BackendOpOverload
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.error import ExportError, ExportErrorType, InternalError
@@ -877,7 +877,7 @@ def _add_debug_handle(self, emitter_id: int, target: _Target) -> None:
         """
         # If it's a delegate call, collect the list of debug handles that are consumed by this
         # delegate call and store it in the debug handle map.
-        if target == delegate.executorch_call_delegate:
+        if target == executorch_call_delegate:
             debug_handle_list = []
             for node in self.node.graph.nodes:
                 if (
@@ -912,7 +912,8 @@ def _emit_argument(
 
     def _emit_delegate(
         self,
-        lowered_module: LoweredBackendModule,
+        # pyre-ignore: Undefined or invalid type [11]: Annotation `LoweredBackendModule` is not defined as a type.
+        lowered_module: "LoweredBackendModule",  # noqa
         args: Tuple[_Argument, ...],
         kwargs: Dict[str, _Argument],
     ) -> _EmitterValue:
@@ -1201,9 +1202,9 @@ def call_function(
         elif target is torch.ops.map_impl:
             return self._emit_control_flow(target, args, kwargs)
 
-        elif target == delegate.executorch_call_delegate:
+        elif target == executorch_call_delegate:
             lowered_module = args[0]
-            assert isinstance(lowered_module, LoweredBackendModule)
+            assert is_lowered_module(lowered_module)
             v = self._emit_delegate(lowered_module, args[1:], kwargs)
             self._add_debug_handle(len(self.chain.instructions) - 1, target)
             return v