Fix delegates

Summary:
With Yidi's migration in D46729844, the typing for the delegates is incorrect/inconsistent, with to_backend and preprocess taking in ExirExportedProgram sometimes, and other times just taking in an fx.GraphModule (pyre-strict is not enabled for the backend/ folder, so it is not caught by typechecking). This diff updates the delegate APIs to take in ExportedProgram.

https://docs.google.com/document/d/1vpvjnGk1TWdnVbgrzuTBJTOFx2D06GTaMoZIrSipWAc/edit

Reviewed By: mergennachin

Differential Revision: D47252888

fbshipit-source-id: 647e26d49dcafd2d08b1a91c1565515c31159ef8

Author: angelayi

backends/backend_api.py

@@ -12,24 +12,24 @@
 from executorch.backends.partitioner import Partitioner, TPartitioner
 from executorch.backends.utils import is_identical_graph
 from executorch.exir import (
-    attach_export_graph_metadata,
     CallSpec,
-    ExirExportedProgram,
     ExportGraphSignature,
-    get_exir_meta,
     MultiMethodExirExportedProgram,
-    pytree,
 )
 
 from executorch.exir.delegate import (
     create_submodule_from_nodes,
     executorch_call_delegate,
     get_lowered_module_name,
     LoweredBackendModule,
-    patch_lowered_functions,
 )
-from executorch.exir.graph_module import get_control_flow_submodules
+from executorch.exir.graph_module import (
+    attach_export_graph_metadata,
+    ExirMetadata,
+    get_control_flow_submodules,
+)
 from executorch.exir.pass_base import ExportPass
+from torch._export.exported_program import ExportedProgram
 
 
 @singledispatch
@@ -39,7 +39,7 @@ def to_backend(args):
 
     def to_backend(
         backend_id: str,
-        edge_graph_module: torch.fx.GraphModule,
+        edge_graph_module: ExportedProgram,
         compile_specs: List[CompileSpec],
     ) -> LoweredBackendModule:
 
@@ -58,23 +58,24 @@ def to_backend(
 @to_backend.register
 def _(
     backend_id: str,
-    edge_graph_module: torch.fx.GraphModule,
+    edge_program: ExportedProgram,
     compile_specs: List[CompileSpec],
 ) -> LoweredBackendModule:
     """
     Add overloaded implementations for to_backend:
     def to_backend(
         backend_id: str,
-        edge_graph_module: torch.fx.GraphModule,
+        edge_program: ExportedProgram,
         compile_specs: List[CompileSpec],
     ) -> LoweredBackendModule:
-    Requires the passed in Module in Edge dialect to be executed in the backend identified
-    by backend_id. The forward method of the given edge_graph_module will be
-    targeted for execution.
+    Requires the passed in exported program in Edge dialect to be executed in
+    the backend identified by backend_id. The forward method of the given
+    edge_graph_module will be targeted for execution.
 
     Args:
         backend_id: The backend identifier.
-        edge_graph_module: A module in Edge dialect to target for lowering to the backend.
+        exported_program: An exported program in Edge dialect to target for
+        lowering to the backend.
         compile_specs: A list of backend-specific objects with static
             metadata to configure the "compilation" process (e.g. it could be
             another dictionary itself).
@@ -83,7 +84,7 @@ def to_backend(
         LoweredBackendModule: A Module that has been lowered to the target backend.
         Internally, the lowered Module contains these special attributes:
         backend_id (str: backend id), __processed_module__ (str: a compiled module)
-        compile_spec, original_module (original exported module)
+        compile_spec, original_module (original exported program)
 
     Raises:
         NotImplementedError: The backend is not implemented (e.g. it was not found).
@@ -93,18 +94,17 @@ def to_backend(
     # All backend implementation are final, so we don't need to consider nested subclasses.
     for cls in BackendDetails.__subclasses__():
         if backend_id == cls.__name__:
-            copied_graph_module = copy.deepcopy(edge_graph_module)
+            copied_edge_program = copy.deepcopy(edge_program)
             processed_bytes = cls.preprocess(
-                copied_graph_module,
+                copied_edge_program,
                 compile_specs,
             )
             lowered_module = LoweredBackendModule(
-                edge_graph_module,
+                edge_program,
                 backend_id,
                 processed_bytes,
                 compile_specs,
             )
-            patch_lowered_functions(lowered_module)
             return lowered_module
     raise NotImplementedError(f"Backend {backend_id} was not found.")
 
@@ -156,9 +156,26 @@ def _partition_and_lower(
         )
         logging.debug(f"Partitioned graph module: {tagged_graph_module}")
 
+        # TODO(T158558782): Update the metadata once we migrate to torch.export
+        submodule_program = ExportedProgram(
+            submodule,
+            submodule.graph,
+            ExportGraphSignature([], [], [], [], {}, {}, {}, None),
+            CallSpec(None, None),
+            {},
+            {},
+            [],
+        )
+        meta = ExirMetadata(
+            in_spec=None,
+            out_spec=None,
+            update_spec=0,
+        )
+        attach_export_graph_metadata(submodule_program.graph_module, meta)
+
         lowered_submodule = to_backend(
             delegation_spec.backend_id,
-            submodule,
+            submodule_program,
             delegation_spec.compile_specs,
         )
 
@@ -199,22 +216,22 @@ def _partition_and_lower(
 
 @to_backend.register
 def _(
-    edge_graph_module: torch.fx.GraphModule,
+    edge_program: ExportedProgram,
     partitioner: Type[TPartitioner],
-) -> torch.fx.GraphModule:
+) -> ExportedProgram:
     """
     Add overloaded implementations for to_backend:
     def to_backend(
-        edge_graph_module: torch.fx.GraphModule,
+        edge_program: ExportedProgram,
         partitioner: Type[TPartitioner],
-    ) -> torch.fx.GraphModule
+    ) -> ExportedProgram:
 
     Returns a semantically-equivalent program to the one given as input (represented
     as a graph module in Edge dialect), but with portions of the program targeted for
     delegation as determined by the partitioner.
 
     Args:
-        torch.fx.GraphModule: Program in Edge dialect.
+        ExportedProgram: Program in Edge dialect.
 
         partitioner: An instance of the Partitioner class type, in charge with tagging
         portions of the input program for delegation. A valid partitioner must have
@@ -224,8 +241,9 @@ def to_backend(
 
 
     Returns:
-        torch.fx.GraphModule: The input program, with some portions targeted for delegation.
+        ExportedProgram: The input program, with some portions targeted for delegation.
     """
+    edge_graph_module = edge_program.graph_module
     copied_graph_module = copy.deepcopy(edge_graph_module)
     # Call the partitioner on the given graph module
     partitioner_instance: Partitioner = partitioner()
@@ -249,7 +267,8 @@ def to_backend(
         tagged_graph_module, partitioner_instance
     )
 
-    return tagged_graph_module
+    edge_program.graph_module = tagged_graph_module
+    return edge_program
 
 
 def to_backend_multiple(
@@ -287,35 +306,18 @@ def to_backend_multiple(
             + "partitioner subclass, or a partitioner subclass."
         )
 
-    method_name_to_delegated_gm = {}
+    method_name_to_delegated_program = {}
     for method_name, prog in multi_method_program.methods().items():
-        gm = prog.graph_module
         if isinstance(partitioner, dict):
             if method_name in partitioner:
-                method_name_to_delegated_gm[method_name] = to_backend(
-                    gm, partitioner[method_name]
+                method_name_to_delegated_program[method_name] = to_backend(
+                    prog, partitioner[method_name]
                 )
             else:
-                method_name_to_delegated_gm[method_name] = gm
+                method_name_to_delegated_program[method_name] = prog
         else:
-            method_name_to_delegated_gm[method_name] = to_backend(gm, partitioner)
-
-    def gm_to_program(gm: torch.fx.GraphModule):
-        ep = ExirExportedProgram(
-            gm,
-            gm.graph,
-            ExportGraphSignature([], [], [], [], {}, {}, {}, None),
-            CallSpec(None, None),
-            {},
-            {},
-            [],
-            True,
-        )
-        ep.graph_module.meta.update(gm.meta)
-        attach_export_graph_metadata(ep.graph_module, get_exir_meta(gm))
-        return ep
+            method_name_to_delegated_program[method_name] = to_backend(
+                prog, partitioner
+            )
 
-    method_name_to_delegated_program = pytree.tree_map(
-        gm_to_program, method_name_to_delegated_gm
-    )
     return MultiMethodExirExportedProgram(method_name_to_delegated_program)

backends/backend_details.py

@@ -1,11 +1,9 @@
 from abc import ABC, abstractmethod
 
-from typing import Callable, Dict, List
-
-import torch
+from typing import List
 
 from executorch.backends.compile_spec_schema import CompileSpec
-from torch.fx.node import Node
+from torch._export.exported_program import ExportedProgram
 
 
 def enforcedmethod(func):
@@ -31,7 +29,7 @@ class BackendDetails(ABC):
     enforced to implement this method.
 
     Args:
-        edge_ir_module: The original module. It will not be modified in place.
+        edge_program: The original exported program. It will not be modified in place.
         backend_debug_handle_generator: A callable to map a graph to a dictionary (key is node, value is id)
         compile_specs: List of values needed for compilation
 
@@ -45,7 +43,7 @@ class BackendDetails(ABC):
     # it's a virtual method and inheritant class needs to implement the actual function
     @abstractmethod
     def preprocess(
-        edge_ir_module: torch.fx.GraphModule,
+        edge_program: ExportedProgram,
         compile_specs: List[CompileSpec],
     ) -> bytes:
         # Users should return a compiled blob - a binary that can run the desired

backends/qnnpack/qnnpack_preprocess.py

@@ -16,6 +16,7 @@
 from executorch.backends.transforms import get_shape
 
 from executorch.exir.dialects._ops import ops as exir_ops
+from torch._export.exported_program import ExportedProgram
 
 T_Mm = exir_ops.edge.aten.mm.default
 T_Addmm = exir_ops.edge.aten.addmm.default
@@ -35,11 +36,11 @@ def _copy_buffer(storage: torch.UntypedStorage) -> bytes:
 class QnnpackBackend(BackendDetails):
     @staticmethod
     def preprocess(
-        edge_ir_module: torch.fx.GraphModule,
+        edge_program: ExportedProgram,
         compile_specs: List[CompileSpec],
     ) -> bytes:
 
-        for node in edge_ir_module.graph.nodes:
+        for node in edge_program.graph.nodes:
             # TODO(maxren): Follow this up by removing addm and mm nodes
             if node.op == "call_function":
                 # Finding the linear node
@@ -52,7 +53,7 @@ def preprocess(
                         weight = node.args[2]
                         # For linear node, bias is known
                         bias_tensor = getattr(
-                            edge_ir_module, node.args[0].target
+                            edge_program.graph_module, node.args[0].target
                         ).contiguous()
                         # t_defualt node -> dequant node
                         weight_dequant = weight.args[0]
@@ -66,7 +67,7 @@ def preprocess(
                         weight_dequant = node.args[1]
                         if len(node.args) > 2:
                             bias_tensor = getattr(
-                                edge_ir_module, node.args[2].target
+                                edge_program.graph_module, node.args[2].target
                             ).contiguous()
                     else:
                         raise RuntimeError(
@@ -89,17 +90,22 @@ def preprocess(
                     # deqaunt node -> quant node
                     weight_quant = weight_dequant.args[0]
                     # quant node -> tensor_constant
-                    weight_const = getattr(edge_ir_module, weight_quant.args[0].target)
+                    weight_const = getattr(
+                        edge_program.graph_module, weight_quant.args[0].target
+                    )
                     if (
                         weight_quant.target.__name__
                         == "quantized_decomposed.quantize_per_channel.default"
                     ):
                         # scale and zero_point are tensors
                         weight_scale = weight_quant.args[1]
-                        scale_tensor = getattr(edge_ir_module, weight_scale.target)
+                        scale_tensor = getattr(
+                            edge_program.graph_module, weight_scale.target
+                        )
                         weight_zeropoint = weight_quant.args[2]
                         zp_tensor = (
-                            getattr(edge_ir_module, weight_zeropoint.target) + 128
+                            getattr(edge_program.graph_module, weight_zeropoint.target)
+                            + 128
                         )
                         axis = weight_quant.args[3]
                         # requantize weight to uint8