1# Copyright (c) Meta Platforms, Inc. and affiliates. 2# All rights reserved. 3# 4# This source code is licensed under the BSD-style license found in the 5# LICENSE file in the root directory of this source tree. 6 7# pyre-strict 8 9import json 10import os 11import tempfile 12 13import pkg_resources 14from executorch.devtools.etdump.schema_flatcc import ETDumpFlatCC 15 16from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass 17 18from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile 19 20# The prefix of schema files used for etdump 21ETDUMP_FLATCC_SCHEMA_NAME = "etdump_schema_flatcc" 22SCALAR_TYPE_SCHEMA_NAME = "scalar_type" 23 24 25def _write_schema(d: str, schema_name: str) -> None: 26 schema_path = os.path.join(d, "{}.fbs".format(schema_name)) 27 with open(schema_path, "wb") as schema_file: 28 schema_file.write( 29 pkg_resources.resource_string(__name__, "{}.fbs".format(schema_name)) 30 ) 31 32 33def _serialize_from_etdump_to_json(etdump: ETDumpFlatCC) -> str: 34 return json.dumps(etdump, cls=_DataclassEncoder, indent=4) 35 36 37""" 38ETDump FlatCC Schema Implementations 39""" 40 41 42# from json to etdump 43def _deserialize_from_json_to_etdump_flatcc(etdump_json: bytes) -> ETDumpFlatCC: 44 etdump_json = json.loads(etdump_json) 45 return _json_to_dataclass(etdump_json, ETDumpFlatCC) 46 47 48def _convert_to_flatcc(etdump_json: str) -> bytes: 49 with tempfile.TemporaryDirectory() as d: 50 # load given and common schema 51 _write_schema(d, ETDUMP_FLATCC_SCHEMA_NAME) 52 _write_schema(d, SCALAR_TYPE_SCHEMA_NAME) 53 54 schema_path = os.path.join(d, "{}.fbs".format(ETDUMP_FLATCC_SCHEMA_NAME)) 55 json_path = os.path.join(d, "{}.json".format(ETDUMP_FLATCC_SCHEMA_NAME)) 56 with open(json_path, "wb") as json_file: 57 json_file.write(etdump_json.encode("ascii")) 58 59 _flatc_compile(d, schema_path, json_path) 60 output_path = os.path.join(d, "{}.etdp".format(ETDUMP_FLATCC_SCHEMA_NAME)) 61 with open(output_path, "rb") as output_file: 62 return output_file.read() 63 64 65def _convert_from_flatcc(etdump_flatbuffer: bytes, size_prefixed: bool = True) -> bytes: 66 with tempfile.TemporaryDirectory() as d: 67 _write_schema(d, ETDUMP_FLATCC_SCHEMA_NAME) 68 _write_schema(d, SCALAR_TYPE_SCHEMA_NAME) 69 70 schema_path = os.path.join(d, "{}.fbs".format(ETDUMP_FLATCC_SCHEMA_NAME)) 71 bin_path = os.path.join(d, "schema.bin") 72 with open(bin_path, "wb") as bin_file: 73 bin_file.write(etdump_flatbuffer) 74 additional_args = [] 75 if size_prefixed: 76 additional_args = ["--size-prefixed"] 77 _flatc_decompile(d, schema_path, bin_path, additional_args) 78 output_path = os.path.join(d, "schema.json") 79 with open(output_path, "rb") as output_file: 80 return output_file.read() 81 82 83def serialize_to_etdump_flatcc( 84 etdump: ETDumpFlatCC, 85) -> bytes: 86 """ 87 Given an ETdump python object this function will return a serialized object 88 that can then be written to a file using the FlatCC schema. 89 Args: 90 etdump: ETDump python object that the user wants to serialize. 91 Returns: 92 Serialized etdump binary blob using the FlatCC schema 93 """ 94 return _convert_to_flatcc(_serialize_from_etdump_to_json(etdump)) 95 96 97def deserialize_from_etdump_flatcc( 98 data: bytes, size_prefixed: bool = True 99) -> ETDumpFlatCC: 100 """ 101 Given an etdump binary blob (constructed using the FlatCC schema) this function will deserialize 102 it and return the FlatCC python object representation of etdump. 103 Args: 104 data: Serialized etdump binary blob. 105 Returns: 106 Deserialized ETDump python object. 107 """ 108 return _deserialize_from_json_to_etdump_flatcc( 109 _convert_from_flatcc(data, size_prefixed) 110 ) 111