Serialize debug handle

Summary: Unfortunately we can't serialize the debug handle for placeholder/output nodes for now because they aren't stored as nodes...

Reviewed By: zhxchen17

Differential Revision: D47311042

fbshipit-source-id: 296e3cb213160319cb37cbf6285f216b36f16c07

Author: angelayi

exir/serde/serialize.py

@@ -41,6 +41,15 @@ def handle_call_function(self, node: torch.fx.Node) -> None:
 
         super().handle_call_function(node)
 
+    def serialize_metadata(self, node: torch.fx.Node) -> Dict[str, str]:
+        meta = super().serialize_metadata(node)
+
+        if "debug_handle" in node.meta:
+            debug_handle = node.meta["debug_handle"]
+            meta["debug_handle"] = str(debug_handle)
+
+        return meta
+
     def serialize_alloc_inputs(
         self, inputs  # pyre-ignore
     ) -> List[schema.NamedArgument]:
@@ -198,6 +207,14 @@ def fake_op(x):
 
         super().deserialize_node(serialized_node, target)
 
+    def deserialize_metadata(self, metadata: Dict[str, str]) -> Dict[str, Any]:
+        res = super().deserialize_metadata(metadata)
+
+        if debug_handle := metadata.get("debug_handle"):
+            res["debug_handle"] = int(debug_handle)
+
+        return res
+
     # pyre-ignore
     def deserialize_alloc_inputs(self, serialized_inputs: List[schema.NamedArgument]):
         def deserialize_alloc_spec(serialized_alloc_spec: str) -> memory.AllocSpec: