xref: /aosp_15_r20/external/executorch/devtools/bundled_program/serialize/__init__.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
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
9# TODO(T138924864): Refactor to unify the serialization for bundled program and executorch program.
10
11import json
12import os
13import tempfile
14
15import executorch.devtools.bundled_program.schema as bp_schema
16
17# @manual=fbsource//third-party/pypi/setuptools:setuptools
18import pkg_resources
19from executorch.devtools.bundled_program.core import BundledProgram
20
21from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass
22from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile
23
24# The prefix of schema files used for bundled program
25BUNDLED_PROGRAM_SCHEMA_NAME = "bundled_program_schema"
26SCALAR_TYPE_SCHEMA_NAME = "scalar_type"
27
28
29def write_schema(d: str, schema_name: str) -> None:
30    schema_path = os.path.join(d, "{}.fbs".format(schema_name))
31    with open(schema_path, "wb") as schema_file:
32        schema_file.write(
33            pkg_resources.resource_string(__name__, "{}.fbs".format(schema_name))
34        )
35
36
37def serialize_from_bundled_program_to_json(
38    bundled_program: bp_schema.BundledProgram,
39) -> str:
40    return json.dumps(bundled_program, cls=_DataclassEncoder)
41
42
43def deserialize_from_json_to_bundled_program(
44    program_json: bytes,
45) -> bp_schema.BundledProgram:
46    program_json = json.loads(program_json)
47    return _json_to_dataclass(program_json, bp_schema.BundledProgram)
48
49
50def convert_to_flatbuffer(program_json: str) -> bytes:
51    with tempfile.TemporaryDirectory() as d:
52        # load given and common schema
53        write_schema(d, BUNDLED_PROGRAM_SCHEMA_NAME)
54        write_schema(d, SCALAR_TYPE_SCHEMA_NAME)
55
56        schema_path = os.path.join(d, "{}.fbs".format(BUNDLED_PROGRAM_SCHEMA_NAME))
57        json_path = os.path.join(d, "{}.json".format(BUNDLED_PROGRAM_SCHEMA_NAME))
58        with open(json_path, "wb") as json_file:
59            json_file.write(program_json.encode("ascii"))
60        _flatc_compile(d, schema_path, json_path)
61        output_path = os.path.join(d, "{}.bpte".format(BUNDLED_PROGRAM_SCHEMA_NAME))
62        with open(output_path, "rb") as output_file:
63            return output_file.read()
64
65
66def convert_from_flatbuffer(program_flatbuffer: bytes) -> bytes:
67    with tempfile.TemporaryDirectory() as d:
68        write_schema(d, BUNDLED_PROGRAM_SCHEMA_NAME)
69        write_schema(d, SCALAR_TYPE_SCHEMA_NAME)
70
71        schema_path = os.path.join(d, "{}.fbs".format(BUNDLED_PROGRAM_SCHEMA_NAME))
72        bin_path = os.path.join(d, "schema.bin")
73        with open(bin_path, "wb") as bin_file:
74            bin_file.write(program_flatbuffer)
75        _flatc_decompile(d, schema_path, bin_path)
76        output_path = os.path.join(d, "schema.json")
77        with open(output_path, "rb") as output_file:
78            return output_file.read()
79
80
81# from bundled program to flatbuffer
82def serialize_from_bundled_program_to_flatbuffer(
83    bundled_program: BundledProgram,
84) -> bytes:
85    """
86    Serialize a BundledProgram into FlatBuffer binary format.
87
88    Args:
89        bundled_program (BundledProgram): The `BundledProgram` variable to be serialized.
90
91    Returns:
92        The serialized FlatBuffer binary data in bytes.
93    """
94
95    bundled_program_in_schema = bundled_program.serialize_to_schema()
96
97    return convert_to_flatbuffer(
98        serialize_from_bundled_program_to_json(bundled_program_in_schema)
99    )
100
101
102# From flatbuffer to bundled program in schema.
103# Please notice here the bundled program is the one in our schema (bp_schema.BundledProgram),
104# not the bundled program user interact with (core.bundled_program).
105# However there're two concerns for current design:
106# 1. the misalignment of serialization input and deserialization out, which may confuse our user.
107# 2. the mis-exposion of schema.bundled_program. all classes in schema should not directly
108#    interact with user, but the deserialization api returns one.
109# TODO(T170042248): Solve the above issues.
110def deserialize_from_flatbuffer_to_bundled_program(
111    flatbuffer: bytes,
112) -> bp_schema.BundledProgram:
113    """
114    Deserialize a FlatBuffer binary format into a BundledProgram.
115
116    Args:
117        flatbuffer (bytes): The FlatBuffer binary data in bytes.
118
119    Returns:
120        A `BundledProgram` instance.
121    """
122    return deserialize_from_json_to_bundled_program(convert_from_flatbuffer(flatbuffer))
123