xref: /aosp_15_r20/external/executorch/backends/qualcomm/serialization/qc_schema_serialize.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1# Copyright (c) Qualcomm Innovation Center, Inc.
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
7import json
8import os
9import tempfile
10
11import pkg_resources
12from executorch.backends.qualcomm.serialization.qc_schema import QnnExecuTorchOptions
13from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass
14from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile
15
16
17def _convert_to_flatbuffer(obj, schema: str):
18    obj_json = json.dumps(obj, cls=_DataclassEncoder)
19    with tempfile.TemporaryDirectory() as d:
20        schema_path = os.path.join(d, f"{schema}.fbs")
21        with open(schema_path, "wb") as schema_file:
22            schema_file.write(pkg_resources.resource_string(__name__, f"{schema}.fbs"))
23        json_path = os.path.join(d, f"{schema}.json")
24        with open(json_path, "wb") as json_file:
25            json_file.write(obj_json.encode("ascii"))
26
27        _flatc_compile(d, schema_path, json_path)
28        output_path = os.path.join(d, f"{schema}.bin")
29        with open(output_path, "rb") as output_file:
30            return output_file.read()
31
32
33def _convert_to_object(flatbuffers: bytes, obj_type, schema: str):
34    with tempfile.TemporaryDirectory() as d:
35        json_path = os.path.join(d, f"{schema}.json")
36        schema_path = os.path.join(d, f"{schema}.fbs")
37        bin_path = os.path.join(d, f"{schema}.bin")
38        with open(schema_path, "wb") as schema_file:
39            schema_file.write(pkg_resources.resource_string(__name__, f"{schema}.fbs"))
40        with open(bin_path, "wb") as bin_file:
41            bin_file.write(flatbuffers)
42
43        _flatc_decompile(d, schema_path, bin_path, ["--raw-binary"])
44        with open(json_path, "rb") as output_file:
45            return _json_to_dataclass(json.load(output_file), obj_type)
46
47
48def option_to_flatbuffer(qnn_executorch_options: QnnExecuTorchOptions) -> bytes:
49    return _convert_to_flatbuffer(qnn_executorch_options, "qc_compiler_spec")
50
51
52def flatbuffer_to_option(flatbuffers: bytes) -> QnnExecuTorchOptions:
53    return _convert_to_object(flatbuffers, QnnExecuTorchOptions, "qc_compiler_spec")
54