Serialize delegate

Summary: Since call_delegate is an executorch concept, we serialize the LoweredModules as string arguments, but under the hood they are a serialized json string.

Reviewed By: tarun292

Differential Revision: D47252889

fbshipit-source-id: 4f017fb6bda8a56f21f734a102c04a78a5853fc0

Author: angelayi

exir/serde/TARGETS

@@ -8,8 +8,21 @@ python_library(
         "serialize.py",
     ],
     deps = [
+        ":schema",
         "//caffe2:torch",
+        "//executorch/backends:compile_spec_schema",
+        "//executorch/exir:delegate",
         "//executorch/exir:lib",
         "//executorch/exir:memory",
     ],
 )
+
+python_library(
+    name = "schema",
+    srcs = [
+        "schema.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+    ],
+)

exir/serde/schema.py

@@ -0,0 +1,21 @@
+# Additional schema from torch._export.serde.schema that is edge specific
+
+from dataclasses import dataclass
+from typing import List
+
+import torch._export.serde.schema as export_schema
+
+
+@dataclass
+class CompileSpec:
+    key: str
+    value: str
+
+
+@dataclass
+class LoweredBackendModule:
+    backend_id: str
+    processed_bytes: str
+    compile_specs: List[CompileSpec]
+    original_module: export_schema.ExportedProgram
+    original_state_dict: str

exir/serde/serialize.py

@@ -1,5 +1,6 @@
 # pyre-strict
 
+import base64
 import copy
 import dataclasses
 import json
@@ -8,11 +9,14 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import executorch.exir as exir
+import executorch.exir.delegate as delegate
 import executorch.exir.memory as memory
 import torch
 import torch._export.exported_program as ep
 import torch._export.serde.schema as schema
 import torch._export.serde.serialize as export_serialize
+from executorch.backends.compile_spec_schema import CompileSpec as delegate_CompileSpec
+from executorch.exir.serde.schema import CompileSpec, LoweredBackendModule
 from torch.fx.experimental import symbolic_shapes
 
 
@@ -39,6 +43,16 @@ def handle_call_function(self, node: torch.fx.Node) -> None:
             self.graph_state.nodes.append(ex_node)
             return
 
+        elif node.target is delegate.executorch_call_delegate:
+            ex_node = schema.Node(
+                target=export_serialize.serialize_operator(node.target),
+                inputs=self.serialize_call_delegate_inputs(node.args),
+                outputs=self.serialize_arbitrary_outputs(node),
+                metadata=self.serialize_metadata(node),
+            )
+            self.graph_state.nodes.append(ex_node)
+            return
+
         super().handle_call_function(node)
 
     def serialize_metadata(self, node: torch.fx.Node) -> Dict[str, str]:
@@ -138,6 +152,71 @@ def serialize_graph(self, graph_module: torch.fx.GraphModule) -> schema.Graph:
         self.original_graph_module: torch.fx.GraphModule = graph_module  # pyre-ignore
         return super().serialize_graph(graph_module)
 
+    def serialize_call_delegate_inputs(
+        self, args  # pyre-ignore
+    ) -> List[schema.NamedArgument]:
+        lowered_module_arg = args[0]
+        delegate_args = args[1:]
+
+        serialized_lowered_module = self.serialize_lowered_module(lowered_module_arg)
+        serialized_lowered_module_arg = schema.NamedArgument(
+            name=lowered_module_arg.target,
+            arg=schema.Argument.create(as_string=serialized_lowered_module),
+        )
+
+        serialized_args = [serialized_lowered_module_arg]
+        for i, arg in enumerate(delegate_args):
+            serialized_args.append(
+                schema.NamedArgument(
+                    name=f"delegate_arg_{i}", arg=self.serialize_input(arg)
+                )
+            )
+        return serialized_args
+
+    def serialize_lowered_module(self, lowered_module_arg: torch.fx.Node) -> str:
+        assert lowered_module_arg.op == "get_attr"
+        assert isinstance(lowered_module_arg.target, str)
+
+        def serialize_bytes(b: bytes) -> str:
+            # We want to serialize the bytes to string because JSON cannot
+            # serialize bytes.
+            # Since the given bytes may be serialized with any encoding, so we
+            # want to first encode with base64, and then decode it with
+            # ascii. During deserialization we can just directly decode with b64
+            # to get the original encoded bytes.
+            return base64.b64encode(b).decode("ascii")
+
+        lowered_module = getattr(
+            lowered_module_arg.graph.owning_module, lowered_module_arg.target
+        )
+        assert isinstance(lowered_module, delegate.LoweredBackendModule)
+
+        serialized_compile_spec = [
+            CompileSpec(cs.key, serialize_bytes(cs.value))
+            for cs in lowered_module.compile_specs
+        ]
+
+        (
+            serialized_original_module,
+            serialized_original_state_dict,
+        ) = ExportedProgramSerializer().serialize(lowered_module.original_module)
+
+        serialized_processed_bytes = serialize_bytes(lowered_module.processed_bytes)
+
+        serialized_lowered_module = LoweredBackendModule(
+            original_module=serialized_original_module,
+            original_state_dict=serialize_bytes(serialized_original_state_dict),
+            processed_bytes=serialized_processed_bytes,
+            compile_specs=serialized_compile_spec,
+            backend_id=lowered_module.backend_id,
+        )
+
+        json_lowered_module = json.dumps(
+            dataclasses.asdict(serialized_lowered_module),
+            cls=export_serialize.EnumEncoder,
+        )
+        return json_lowered_module
+
 
 class ExportedProgramSerializer(export_serialize.ExportedProgramSerializer):
     def serialize(
@@ -186,6 +265,27 @@ def deserialize_node(self, serialized_node: schema.Node, target: Callable) -> No
             fx_node.meta.update(self.deserialize_metadata(serialized_node.metadata))
             return
 
+        elif target is delegate.executorch_call_delegate:
+            if (
+                len(serialized_node.outputs) == 1
+                and serialized_node.outputs[0].type == "as_tensor"
+            ):
+                # If it's a single tensor return then we can use the name of the
+                # node itself
+                name = serialized_node.outputs[0].value.name
+            else:
+                # Otherwise FX will make a name for us, and we'll have `getitem`
+                # nodes pointed to that
+                name = None
+
+            args = self.deserialize_call_delegate_inputs(serialized_node.inputs)
+            fx_node = self.graph.create_node("call_function", target, args, {}, name)
+
+            self.deserialize_arbitrary_outputs(serialized_node, fx_node)
+
+            fx_node.meta.update(self.deserialize_metadata(serialized_node.metadata))
+            return
+
         elif isinstance(target, str):
             # Create a dummy fake op if the target does not exist
             # because we cannot create a call_function node w/o a
@@ -267,6 +367,49 @@ def deserialize_input(self, inp: schema.Argument) -> Any:
 
         return super().deserialize_input(inp)
 
+    # pyre-ignore
+    def deserialize_call_delegate_inputs(
+        self, serialized_inputs: List[schema.NamedArgument]
+    ):
+        serialized_lowered_module = serialized_inputs[0]
+        lowered_module_node = self.deserialize_lowered_module(serialized_lowered_module)
+        serialized_delegate_inputs = serialized_inputs[1:]
+        args = tuple(
+            self.deserialize_input(input.arg) for input in serialized_delegate_inputs
+        )
+        return (lowered_module_node,) + args
+
+    def deserialize_lowered_module(
+        self, serialized_lowered_module_arg: schema.NamedArgument
+    ) -> torch.fx.Node:
+        assert serialized_lowered_module_arg.arg.type == "as_string"
+        lowered_module_str = serialized_lowered_module_arg.arg.value
+        json_lowered_module = json.loads(lowered_module_str)
+        serialized_lowered_module = export_serialize._dict_to_dataclass(
+            LoweredBackendModule, json_lowered_module
+        )
+
+        backend_id = serialized_lowered_module.backend_id
+        processed_bytes = base64.b64decode(serialized_lowered_module.processed_bytes)
+        compile_specs = [
+            delegate_CompileSpec(key=cs.key, value=base64.b64decode(cs.value))
+            for cs in serialized_lowered_module.compile_specs
+        ]
+
+        original_module = ExportedProgramDeserializer().deserialize(
+            serialized_lowered_module.original_module,
+            base64.b64decode(serialized_lowered_module.original_state_dict),
+        )
+
+        lowered_module = delegate.LoweredBackendModule(
+            original_module,
+            backend_id,
+            processed_bytes,
+            compile_specs,
+        )
+        self.module.register_module(serialized_lowered_module_arg.name, lowered_module)
+        return self.graph.get_attr(serialized_lowered_module_arg.name)
+
 
 class ExportedProgramDeserializer(export_serialize.ExportedProgramDeserializer):
     def deserialize(

exir/tests/TARGETS

@@ -96,6 +96,9 @@ python_unittest(
     ],
     deps = [
         "//caffe2:torch",
+        "//executorch/backends:backend_api",
+        "//executorch/backends/test:backend_with_compiler_demo",
+        "//executorch/backends/test:op_partitioner_demo",
         "//executorch/exir:lib",
         "//executorch/exir/serde:serialize",
     ],

exir/tests/test_serde.py

@@ -6,6 +6,11 @@
 import executorch.exir as exir
 
 import torch
+from executorch.backends.backend_api import CompileSpec, to_backend
+from executorch.backends.test.backend_with_compiler_demo import (  # noqa
+    BackendWithCompilerDemo,
+)
+from executorch.backends.test.op_partitioner_demo import AddMulPartitionerDemo
 from executorch.exir.serde.serialize import deserialize, serialize
 from torch._export.exported_program import ExportedProgram as TorchExportedProgram
 from torch.utils import _pytree as pytree
@@ -89,3 +94,62 @@ def get_random_inputs(self):
         model = MyModel()
         inputs = model.get_random_inputs()
         self.check_serde(model, inputs)
+
+    def test_delegate(self) -> None:
+        class SinModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return torch.sin(x)
+
+        sin_module = SinModule()
+        model_inputs = (torch.ones(1),)
+        edgeir_m = exir.capture(
+            sin_module, model_inputs, exir.CaptureConfig(pt2_mode=True)
+        ).to_edge()
+        max_value = model_inputs[0].shape[0]
+        compile_specs = [CompileSpec("max_value", bytes([max_value]))]
+        lowered_sin_module = to_backend(
+            "BackendWithCompilerDemo", edgeir_m, compile_specs
+        )
+
+        class CompositeModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.lowered_linear_sin = lowered_sin_module
+
+            def forward(self, x):
+                return self.lowered_linear_sin(x)
+
+        composite_model = CompositeModule()
+        model_inputs = (torch.ones(1),)
+
+        composite_model(*model_inputs)
+
+        aten = exir.capture(
+            composite_model, model_inputs, exir.CaptureConfig(pt2_mode=True)
+        )
+        aten_new = deserialize(*serialize(aten))
+        self.check_ep(aten, aten_new, model_inputs)
+
+    def test_delegate_partitioner(self) -> None:
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, a, x, b):
+                y = torch.mm(a, x)
+                z = y + b
+                a = z - a
+                y = torch.mm(a, x)
+                z = y + b
+                return z
+
+        m = Model()
+        inputs = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
+
+        ep = exir.capture(m, inputs, exir.CaptureConfig(pt2_mode=True)).to_edge()
+        edge = to_backend(ep, AddMulPartitionerDemo)
+        edge_new = deserialize(*serialize(edge))
+        self.check_ep(edge, edge_new, inputs)