frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[export] Implement cpp deserializer. (#136398)

Browse Source 
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: https://github.com/pytorch/pytorch/pull/136398
Approved by: https://github.com/angelayi
This commit is contained in:
Zhengxu Chen
2024-11-14 16:34:56 +00:00
committed by PyTorch MergeBot
parent f98c601efe
commit 3ef2dfc1ba
15 changed files with 2656 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
scripts/export/update_schema.py
View File

@ -29,7 +29,7 @@ if __name__ == "__main__":
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 @@ if __name__ == "__main__":
+ 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)