1#!/usr/bin/env fbpython 2# Copyright (c) Meta Platforms, Inc. and affiliates. 3# All rights reserved. 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 8""" 9This is a tool to create sample BundledProgram flatbuffer for executor_runner. 10This file and generate_linear_out_model.py share the same model. 11To update the bundled program flatbuffer file, just simply run: 12buck2 run executorch/test/models:generate_linear_out_bundled_program 13Then commit the updated file (if there are any changes). 14""" 15 16import subprocess 17from typing import List 18 19import torch 20from executorch.devtools import BundledProgram 21from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite 22from executorch.devtools.bundled_program.serialize import ( 23 serialize_from_bundled_program_to_flatbuffer, 24) 25from executorch.exir import ExecutorchBackendConfig, to_edge 26 27from executorch.exir.passes import MemoryPlanningPass, ToOutVarPass 28from executorch.exir.print_program import pretty_print 29 30from executorch.test.models.linear_model import LinearModel 31from torch.export import export 32 33 34def main() -> None: 35 model = LinearModel() 36 37 trace_inputs = (torch.ones(2, 2, dtype=torch.float),) 38 39 # Trace to FX Graph. 40 exec_prog = to_edge(export(model, trace_inputs)).to_executorch( 41 config=ExecutorchBackendConfig( 42 memory_planning_pass=MemoryPlanningPass(), 43 to_out_var_pass=ToOutVarPass(), 44 ) 45 ) 46 # Emit in-memory representation. 47 pretty_print(exec_prog.executorch_program) 48 49 # Serialize to flatbuffer. 50 exec_prog.executorch_program.version = 0 51 52 # Create test sets 53 method_test_cases: List[MethodTestCase] = [] 54 for _ in range(10): 55 x = [ 56 torch.rand(2, 2, dtype=torch.float), 57 ] 58 method_test_cases.append(MethodTestCase(inputs=x, expected_outputs=model(*x))) 59 method_test_suites = [ 60 MethodTestSuite(method_name="forward", test_cases=method_test_cases) 61 ] 62 63 bundled_program = BundledProgram(exec_prog, method_test_suites) 64 65 bundled_program_flatbuffer = serialize_from_bundled_program_to_flatbuffer( 66 bundled_program 67 ) 68 69 fbsource_base_path = ( 70 subprocess.run(["hg", "root"], stdout=subprocess.PIPE).stdout.decode().strip() 71 ) 72 with open( 73 f"{fbsource_base_path}/fbcode/executorch/test/models/linear_out_bundled_program.bpte", 74 "wb", 75 ) as file: 76 file.write(bundled_program_flatbuffer) 77 78 79if __name__ == "__main__": 80 main() 81