[export] Implement cpp deserializer. (pytorch#136398)

Differential Revision: D63206258

This diff introduces a mechanism to generate a json-compatible deserializer in cpp using nlohmann json (already being used by AOTI).

Why we need this? Because there will be a lot of cases where people don't want to use Python to load the graph (e.g. cpp runtime), and instead they can use this header to deserialize the JSON graph.

Every time we call update_schema.py to update the schema, the header will be auto generated and included into the source files.
Pull Request resolved: pytorch#136398
Approved by: https://github.com/angelayi

Author: zhxchen17

.gitattributes

@@ -5,3 +5,4 @@
 .github/scripts/gql_mocks.json linguist-generated=true
 third_party/LICENSES_BUNDLED.txt linguist-generated=true
 tools/build/bazel/requirements.txt linguist-generated=true
+torch/csrc/utils/generated_serialization_types.h linguist-generated=true

build_variables.bzl

@@ -844,6 +844,7 @@ libtorch_python_core_sources = [
     "torch/csrc/fx/node.cpp",
     "torch/csrc/mps/Module.cpp",
     "torch/csrc/mtia/Module.cpp",
+    "torch/csrc/export/pybind.cpp",
     "torch/csrc/inductor/aoti_package/pybind.cpp",
     "torch/csrc/inductor/aoti_runner/pybind.cpp",
     "torch/csrc/inductor/aoti_eager/kernel_holder.cpp",

scripts/export/update_schema.py

@@ -29,7 +29,7 @@
 
     commit = schema_check.update_schema()
 
-    if os.path.exists(args.prefix + commit.path):
+    if os.path.exists(args.prefix + commit.yaml_path):
         if commit.result["SCHEMA_VERSION"] < commit.base["SCHEMA_VERSION"]:
             raise RuntimeError(
                 f"Schema version downgraded from {commit.base['SCHEMA_VERSION']} to {commit.result['SCHEMA_VERSION']}."
@@ -55,17 +55,28 @@
             + f"Reason: {reason}"
         )
 
-    header = (
-        "# @" + "generated by " + os.path.basename(__file__).rsplit(".", 1)[0] + ".py"
+    first_line = (
+        "@" + "generated by " + os.path.basename(__file__).rsplit(".", 1)[0] + ".py"
     )
-    header += f"\n# checksum<<{commit.checksum_result}>>"
-    payload = dump(commit.result, Dumper=Dumper, sort_keys=False)
+    checksum = f"checksum<<{commit.checksum_result}>>"
+    yaml_header = "# " + first_line
+    yaml_header += "\n# " + checksum
+    yaml_payload = dump(commit.result, Dumper=Dumper, sort_keys=False)
 
-    content = header + "\n" + payload
+    cpp_header = "// " + first_line
+    cpp_header += "\n// " + checksum
+    cpp_header += "\n// clang-format off"
+    cpp_header += "\n" + commit.cpp_header
+    cpp_header += "\n// clang-format on"
+    cpp_header += "\n"
+
+    yaml_content = yaml_header + "\n" + yaml_payload
 
     if args.dry_run:
-        print(content)
-        print("\nWill write the above schema to" + args.prefix + commit.path)
+        print(yaml_content)
+        print("\nWill write the above schema to" + args.prefix + commit.yaml_path)
     else:
-        with open(args.prefix + commit.path, "w") as f:
-            f.write(content)
+        with open(args.prefix + commit.yaml_path, "w") as f:
+            f.write(yaml_content)
+        with open(args.prefix + commit.cpp_header_path, "w") as f:
+            f.write(cpp_header)

test/export/test_cpp_serdes.py

@@ -0,0 +1,60 @@
+# Owner(s): ["oncall: export"]
+
+
+import torch
+from torch._export.serde.serialize import deserialize, serialize
+
+
+try:
+    from . import test_export, testing
+except ImportError:
+    import test_export  # @manual=fbcode//caffe2/test:test_export-library
+    import testing  # @manual=fbcode//caffe2/test:test_export-library
+
+from torch.export import export
+
+
+test_classes = {}
+
+
+def mocked_cpp_serdes_export(*args, **kwargs):
+    ep = export(*args, **kwargs)
+    try:
+        payload = serialize(ep)
+    except Exception:
+        return ep
+    cpp_ep = torch._C._export.deserialize_exported_program(payload.exported_program)
+    loaded_json = torch._C._export.serialize_exported_program(cpp_ep)
+    payload.exported_program = loaded_json.encode()
+    loaded_ep = deserialize(payload)
+    return loaded_ep
+
+
+def make_dynamic_cls(cls):
+    cls_prefix = "CppSerdes"
+
+    test_class = testing.make_test_cls_with_mocked_export(
+        cls,
+        cls_prefix,
+        "_cpp_serdes",
+        mocked_cpp_serdes_export,
+        xfail_prop="_expected_failure_cpp_serdes",
+    )
+
+    test_classes[test_class.__name__] = test_class
+    # REMOVING THIS LINE WILL STOP TESTS FROM RUNNING
+    globals()[test_class.__name__] = test_class
+    test_class.__module__ = __name__
+
+
+tests = [
+    test_export.TestExport,
+]
+for test in tests:
+    make_dynamic_cls(test)
+del test
+
+if __name__ == "__main__":
+    from torch._dynamo.test_case import run_tests
+
+    run_tests()

test/export/test_export.py

@@ -2951,6 +2951,7 @@ def forward(self, x):
         export(N(), inputs, dynamic_shapes=dynamic_shapes)
 
     @testing.expectedFailureSerDer  # no unbacked bindings after deserialization?
+    @testing.expectedFailureCppSerDes  # no unbacked bindings after deserialization?
     @testing.expectedFailureSerDerNonStrict
     def test_unbacked_bindings_for_divisible_u_symint(self):
         with torch.library._scoped_library("mylib", "FRAGMENT") as lib:
@@ -3673,6 +3674,7 @@ def forward(self, arg1, arg2, *args, kw1, kw2, **kwargs):
         self._test_export_same_as_eager(kw_func, args, kwargs)
 
     @testing.expectedFailureSerDer  # we don't save placeholder metadata
+    @testing.expectedFailureCppSerDes  # we don't save placeholder metadata
     @testing.expectedFailureSerDerNonStrict
     @testing.expectedFailureNonStrict
     @testing.expectedFailureTrainingIRToRunDecompNonStrict  # source_fn_stack failure
@@ -8078,6 +8080,7 @@ def forward(self, x):
         export(f, (inputs,), dynamic_shapes=dynamic_shapes)
 
     @testing.expectedFailureRetraceabilityNonStrict
+    @testing.expectedFailureCppSerDes  # dynamic shape serialization
     def test_disable_forced_specializations_ok(self):
         # check that we don't force specialization, and defer to runtime asserts
         # with allow_complex_guards_as_runtime_asserts=True to successfully export
@@ -8198,6 +8201,7 @@ def forward(self, w, x, y, z):
 
     # TODO requires_grad doesn't seem to work with serialization.
     @testing.expectedFailureSerDer
+    @testing.expectedFailureCppSerDes
     @testing.expectedFailureSerDerNonStrict
     def test_preserve_requires_grad_placeholders(self):
         class Module(torch.nn.Module):
@@ -8536,6 +8540,7 @@ def forward(self, x, y):
             ep.graph_module.code
         )
 
+    @testing.expectedFailureCppSerDes
     def test_slice_with_floordiv(self):
         # slice operation emits runtime assert s0//2 <= s1
         class M1(torch.nn.Module):
@@ -9105,6 +9110,7 @@ def test_dynamic_shapes_serdes_user_errors(self):
             _load_dynamic_shapes(spec, from_dict=True)
 
     @testing.expectedFailureSerDer  # TODO(pianpwk): PowByNatural valuerange deserialization
+    @testing.expectedFailureCppSerDes  # TODO(pianpwk): PowByNatural valuerange deserialization
     @testing.expectedFailureSerDerNonStrict
     @testing.expectedFailureRetraceabilityNonStrict
     def test_dim_dynamic(self):

test/export/test_schema.py

@@ -106,11 +106,13 @@ def test_schema_check(self):
         commit = _Commit(
             result=src,
             checksum_result="",
-            path="",
+            yaml_path="",
             additions=additions,
             subtractions=subtractions,
             base=dst,
             checksum_base="",
+            cpp_header="",
+            cpp_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [4, 1])
@@ -138,11 +140,13 @@ def test_schema_check(self):
         commit = _Commit(
             result=src,
             checksum_result="",
-            path="",
+            yaml_path="",
             additions=additions,
             subtractions=subtractions,
             base=dst,
             checksum_base="",
+            cpp_header="",
+            cpp_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [4, 1])
@@ -173,11 +177,13 @@ def test_schema_check(self):
         commit = _Commit(
             result=src,
             checksum_result="",
-            path="",
+            yaml_path="",
             additions=additions,
             subtractions=subtractions,
             base=dst,
             checksum_base="",
+            cpp_header="",
+            cpp_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [3, 3])
@@ -231,11 +237,13 @@ def test_schema_check(self):
         commit = _Commit(
             result=src,
             checksum_result="",
-            path="",
+            yaml_path="",
             additions=additions,
             subtractions=subtractions,
             base=dst,
             checksum_base="",
+            cpp_header="",
+            cpp_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [3, 3])
@@ -259,11 +267,13 @@ def test_schema_check(self):
         commit = _Commit(
             result=src,
             checksum_result="",
-            path="",
+            yaml_path="",
             additions=additions,
             subtractions=subtractions,
             base=dst,
             checksum_base="",
+            cpp_header="",
+            cpp_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [3, 3])
@@ -294,11 +304,13 @@ def test_schema_check(self):
         commit = _Commit(
             result=src,
             checksum_result="",
-            path="",
+            yaml_path="",
             additions=additions,
             subtractions=subtractions,
             base=dst,
             checksum_base="",
+            cpp_header="",
+            cpp_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [3, 3])
@@ -326,11 +338,13 @@ def test_schema_check(self):
         commit = _Commit(
             result=src,
             checksum_result="",
-            path="",
+            yaml_path="",
             additions=additions,
             subtractions=subtractions,
             base=dst,
             checksum_base="",
+            cpp_header="",
+            cpp_header_path="",
         )
         next_version, _ = check(commit)
         self.assertEqual(next_version, [4, 1])

test/export/testing.py

@@ -284,3 +284,8 @@ def expectedFailureSerDerPreDispatch(fn):
 def expectedFailurePreDispatchRunDecomp(fn):
     fn._expected_failure_pre_dispatch = True
     return fn
+
+
+def expectedFailureCppSerDes(fn):
+    fn._expected_failure_cpp_serdes = True
+    return fn

torch/_C/__init__.pyi.in

@@ -64,6 +64,7 @@ from torch.utils._python_dispatch import TorchDispatchMode
 
 from . import (
     _aoti,
+    _export,
     _cpu,
     _dynamo,
     _functorch,

torch/_C/_export.pyi

@@ -0,0 +1,10 @@
+# Defined in torch/csrc/export/pybind.cpp
+
+class CppExportedProgram: ...
+
+def deserialize_exported_program(
+    serialized_program: str,
+) -> CppExportedProgram: ...
+def serialize_exported_program(
+    cpp_exported_program: CppExportedProgram,
+) -> str: ...