xref: /aosp_15_r20/external/executorch/examples/arm/executor_runner/pte_to_header.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3# Copyright 2023-2024 Arm Limited and/or its affiliates.
4#
5# This source code is licensed under the BSD-style license found in the
6# LICENSE file in the root directory of this source tree.
7
8import binascii
9import os
10from argparse import ArgumentParser, ArgumentTypeError
11
12# Also see: https://git.mlplatform.org/ml/ethos-u/ml-embedded-evaluation-kit.git/tree/scripts/py/gen_model_cpp.py
13
14bytes_per_line = 32
15hex_digits_per_line = bytes_per_line * 2
16
17
18def input_file_path(path):
19    if os.path.exists(path):
20        return path
21    else:
22        raise ArgumentTypeError(f"input filepath:{path} does not exist")
23
24
25parser = ArgumentParser()
26parser.add_argument(
27    "-p",
28    "--pte",
29    help="ExecuTorch .pte model file",
30    type=input_file_path,
31    required=True,
32)
33parser.add_argument(
34    "-d",
35    "--outdir",
36    help="Output dir for model header",
37    type=str,
38    required=False,
39    default=".",
40)
41
42parser.add_argument(
43    "-o",
44    "--outfile",
45    help="Output filename for model header",
46    type=str,
47    required=False,
48    default="model_pte.h",
49)
50parser.add_argument(
51    "-s",
52    "--section",
53    help="Section attribute for the data array",
54    type=str,
55    required=False,
56    default="network_model_sec",
57)
58
59if __name__ == "__main__":
60    args = parser.parse_args()
61    outfile = os.path.join(args.outdir, args.outfile)
62    attr = f'__attribute__((section("{args.section}"), aligned(16))) char '
63
64    with open(args.pte, "rb") as fr, open(outfile, "w") as fw:
65        data = fr.read()
66        hexstream = binascii.hexlify(data).decode("utf-8")
67        fw.write(attr + "model_pte[] = {")
68
69        for i in range(0, len(hexstream), 2):
70            if 0 == (i % hex_digits_per_line):
71                fw.write("\n")
72            fw.write("0x" + hexstream[i : i + 2] + ", ")
73
74        fw.write("};\n")
75
76        print(
77            f"Input: {args.pte} with {len(data)} bytes. Output: {outfile} with {fw.tell()} bytes. Section: {args.section}."
78        )
79