[etrecord] Implement generic fallback for GraphModuleSerializer.handle_call_function (#16069)

Title says it all!

Implement the case where `node.target` is neither `torch._ops.OpOverload` or `torch._ops.HigherOrderOperator`, instead of throwing an exception.

Differential Revision: [D88216198](https://our.internmc.facebook.com/intern/diff/D88216198/)

Author: pytorchbot

exir/serde/export_serialize.py

@@ -54,6 +54,8 @@
 from torch.utils._sympy.numbers import int_oo
 from torch.utils._sympy.value_ranges import ValueRanges
 
+import executorch.exir as exir
+
 # pyre-ignore
 
 from .schema import (  # type: ignore[attr-defined]
@@ -205,6 +207,11 @@ def _reverse_map(d: Dict[Any, Enum]):
 }
 
 
+_KNOWN_FUNCTIONS = {
+    exir.memory.view,
+}
+
+
 @dataclass
 class SerializedArtifact:
     exported_program: bytes
@@ -545,6 +552,14 @@ def handle_call_function(self, node: torch.fx.Node):
                 outputs=self.serialize_hoo_outputs(node),
                 metadata=self.serialize_metadata(node),
             )
+        elif node.target in _KNOWN_FUNCTIONS:
+            ex_node = Node(
+                name=node.name,
+                target=node._pretty_print_target(node.target),
+                inputs=self.serialize_hoo_inputs(node.args, node.kwargs),
+                outputs=self.serialize_hoo_outputs(node),
+                metadata=self.serialize_metadata(node),
+            )
         else:
             raise SerializeError(f"Serializing {node.target} is not supported")
 

exir/serde/serialize.py

@@ -374,6 +374,11 @@ def serialize(
         )
 
 
+_KNOWN_FUNCTIONS_MAP = {
+    "executorch.exir.memory.view": exir.memory.view,
+}
+
+
 class GraphModuleDeserializer(export_serialize.GraphModuleDeserializer):
     def deserialize_operator(self, serialized_target: str) -> str:
         def find_operator(module: _DialectNamespace, serialized_target: str) -> str:
@@ -450,19 +455,23 @@ def deserialize_node(self, serialized_node: schema.Node, target: Callable) -> No
             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
-            # callable target
-            log.warning(
-                f"Could not find operator {target}. Returning fake operator."
-            )  # noqa: G004
-
-            # pyre-ignore
-            def fake_op(x):
-                raise NotImplementedError("Fake op is not meant to be run.")
-
-            fake_op.__name__ = target
-            target = fake_op
+            # Special handling for known functions, which are serialized as a
+            # string but are still somewhat expected in serialized graphs.
+            if target in _KNOWN_FUNCTIONS_MAP:
+                target = _KNOWN_FUNCTIONS_MAP[target]
+            else:
+                # Otherwise, create a dummy fake op if the target does not exist
+                # because we cannot create a call_function node w/o a callable
+                # target
+                log.warning(
+                    f"Could not find operator {target}. Returning fake operator."
+                )  # noqa: G004
+                # pyre-ignore
+                def fake_op(x):
+                    raise NotImplementedError("Fake op is not meant to be run.")
+
+                fake_op.__name__ = target
+                target = fake_op
 
             args = self.deserialize_inputs_no_schema(serialized_node)
             fx_node = self.graph.create_node("call_function", target, args, None, None)

exir/tests/test_serde.py

@@ -335,3 +335,18 @@ def forward(self, x):
                     node.meta.get("from_node"), node_new.meta.get("from_node")
                 ):
                     self.assertEqual(node_source.to_dict(), node_source_new.to_dict())
+
+    def test_memory_ops(self) -> None:
+        class MemoryOpsModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x, y):
+                x = exir.memory.view(x, (10, 10))
+                return x + y
+
+        inputs = (
+            torch.randn(100),
+            torch.randn(10, 10),
+        )
+        self.check_serde(MemoryOpsModule(), inputs)