Add support for backend dialect in serde

Summary: Adding support for backend dialect ops in EXIR serde.

Reviewed By: larryliu0820

Differential Revision: D48106247

fbshipit-source-id: 60424599ec59263c3a4e8e52086eb9a6c8e1f660

Author: tarun292

exir/serde/TARGETS

@@ -16,6 +16,7 @@ python_library(
         "//executorch/exir:memory",
         "//executorch/exir/backend:compile_spec_schema",
         "//executorch/exir/dialects:lib",
+        "//executorch/exir/dialects/backend:lib",
         "//executorch/exir/dialects/edge:lib",
     ],
 )

exir/serde/serialize.py

@@ -24,7 +24,8 @@
 from executorch.exir.backend.compile_spec_schema import (
     CompileSpec as delegate_CompileSpec,
 )
-from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.dialects._ops import _DialectNamespace, ops as exir_ops
+from executorch.exir.dialects.backend._ops import BackendOpOverload
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.lowered_backend_module import (
     LoweredBackendModule as ExirLoweredBackendModule,
@@ -50,20 +51,27 @@ def serialize_operator(
         target: Union[
             str,
             EdgeOpOverload,
+            BackendOpOverload,
             torch._ops.OpOverload,
             torch._ops.HigherOrderOperator,
         ],
     ) -> str:
         if isinstance(target, str):
             return target
-        elif target.__module__.startswith("executorch.exir.dialects"):
+        elif target.__module__.startswith("executorch.exir.dialects.edge"):
             # TODO(zhxchen17) Maybe provide a function name helper in FX.
             # From torch.fx.node._get_qualified_name
             module = target.__module__.replace(
                 "executorch.exir.dialects.edge._ops",
                 "executorch.exir.dialects.edge.ops",
             )
             return f"{module}.{target.__name__}"
+        elif target.__module__.startswith("executorch.exir.dialects.backend"):
+            module = target.__module__.replace(
+                "executorch.exir.dialects.backend._ops",
+                "executorch.exir.dialects.backend.ops",
+            )
+            return f"{module}.{target.__name__}"
 
         return super().serialize_operator(target)
 
@@ -337,8 +345,7 @@ def __init__(self, state_dict: Dict[str, torch.Tensor]) -> None:
         self.state_dict: Dict[str, Any] = state_dict  # TODO(T157676982)
 
     def deserialize_operator(self, serialized_target: str) -> str:
-        if serialized_target.startswith("executorch.exir.dialects.edge.ops"):
-            module = exir_ops.edge
+        def find_operator(module: _DialectNamespace, serialized_target: str) -> str:
             serialized_target_names = serialized_target.split(".")[5:]
 
             target = module
@@ -349,6 +356,11 @@ def deserialize_operator(self, serialized_target: str) -> str:
                     target = getattr(target, name)
             return target
 
+        if serialized_target.startswith("executorch.exir.dialects.edge.ops"):
+            return find_operator(exir_ops.edge, serialized_target)
+        elif serialized_target.startswith("executorch.exir.dialects.backend.ops"):
+            return find_operator(exir_ops.backend, serialized_target)
+
         return super().deserialize_operator(serialized_target)
 
     # pyre-ignore