def __call__(self, measure_input, tmp_dir, **kwargs): """ Wrapped build func. Parameters ---------- measure_input: MeasureInput The input of measurement tmp_dir: str The path of temporary directory to export generated library """ tic = time.time() try: filename = os.path.join( tmp_dir, "tmp_func_%0x.%s" % (getrandbits(64), self.build_func.output_format)) # TODO(tvm-team) consider linline _build_func_common func, arg_info = _build_func_common(measure_input, self.runtime, **kwargs) if self.build_func.output_format == ".model-library-format": # Late import to preserve autoTVM with USE_MICRO OFF try: from tvm import micro # pylint: disable=import-outside-toplevel except ImportError: raise ImportError("Requires USE_MICRO") micro.export_model_library_format(func, filename) else: func.export_library(filename, self.build_func) except Exception as e: # pylint: disable=broad-except return BuildResult(None, None, e, time.time() - tic) return BuildResult(filename, arg_info, None, time.time() - tic)
def test_multiple_relay_modules_same_module_name(): mod = get_conv2d_relay_module() executor = Executor("graph") runtime = Runtime("crt") target = tvm.target.target.micro("host") with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): factory1 = tvm.relay.build(mod, target, runtime=runtime, executor=executor, mod_name="mod") factory2 = tvm.relay.build(mod, target, runtime=runtime, executor=executor, mod_name="mod") temp_dir = utils.tempdir() mlf_tar_path = temp_dir.relpath("lib.tar") with pytest.raises(AssertionError, match="Multiple modules should have unique names"): micro.export_model_library_format([factory1, factory2], mlf_tar_path)
def test_multiple_relay_modules_graph(): mod = get_conv2d_relay_module() executor = Executor("graph") runtime = Runtime("crt") target = tvm.target.target.micro("host") with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): factory1 = tvm.relay.build(mod, target, runtime=runtime, executor=executor, mod_name="mod1") factory2 = tvm.relay.build(mod, target, runtime=runtime, executor=executor, mod_name="mod2") temp_dir = utils.tempdir() mlf_tar_path = temp_dir.relpath("lib.tar") micro.export_model_library_format([factory1, factory2], mlf_tar_path) with tarfile.open(mlf_tar_path, "r:*") as tf: tar_members = [ti.name for ti in tf.getmembers()] print("tar members", tar_members) assert "./metadata.json" in tar_members assert "./codegen/host/src/mod1_lib0.c" in tar_members assert "./codegen/host/src/mod2_lib0.c" in tar_members with tf.extractfile("./metadata.json") as f: metadata = json.load(f) mod2_main_md = metadata["modules"]["mod2"]["memory"]["functions"]["main"] assert mod2_main_md == [ { "constants_size_bytes": 0, "device": 1, "io_size_bytes": 143960, "workspace_size_bytes": 158088, } ] assert metadata["modules"]["mod1"]["model_name"] == "mod1" assert metadata["modules"]["mod2"]["model_name"] == "mod2"
def test_export_model_library_format_workspace(executor, runtime): target = tvm.target.target.micro("host") with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): relay_mod = tvm.parser.fromtext( """ #[version = "0.0.5"] def @main(%p0: Tensor[(1, 56, 56, 128), int16], %p1: Tensor[(3, 3, 128, 1), int16], %p2: Tensor[(1, 1, 1, 128), int32]){ %0 = nn.conv2d(%p0, %p1, padding=[1, 1, 1, 1], groups=128, channels=128, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWOI", out_dtype="int32") /* ty=Tensor[(1, 56, 56, 128), int32] */; %1 = add(%0, %p2) /* ty=Tensor[(1, 56, 56, 128), int32] */; %2 = fixed_point_multiply(%1, multiplier=2080045879, shift=-4) /* ty=Tensor[(1, 56, 56, 128), int32] */; %3 = clip(%2, a_min=0f, a_max=255f) /* ty=Tensor[(1, 56, 56, 128), int32] */; cast(%3, dtype="uint8") /* ty=Tensor[(1, 56, 56, 128), uint8] */ } """ ) factory = tvm.relay.build( relay_mod, target, executor=executor, runtime=runtime, mod_name="qnn_conv2d", ) temp_dir = utils.tempdir() mlf_tar_path = temp_dir.relpath("lib.tar") micro.export_model_library_format(factory, mlf_tar_path) tf = tarfile.open(mlf_tar_path) extract_dir = temp_dir.relpath("extract") os.mkdir(extract_dir) tf.extractall(extract_dir) with open(os.path.join(extract_dir, "metadata.json")) as json_f: metadata = json.load(json_f) module_name = factory.libmod_name assert metadata["version"] == _GENERATED_VERSION assert metadata["modules"][module_name]["model_name"] == "qnn_conv2d" export_datetime = datetime.datetime.strptime( metadata["modules"][module_name]["export_datetime"], "%Y-%m-%d %H:%M:%SZ" ) assert (datetime.datetime.now() - export_datetime) < datetime.timedelta(seconds=60 * 5) assert metadata["modules"][module_name]["target"] == [str(target)] assert metadata["modules"][module_name]["memory"]["functions"]["main"] == [ { "constants_size_bytes": 0, "device": 1, "io_size_bytes": 1207040, "workspace_size_bytes": 2466816, } ] assert metadata["modules"][module_name]["memory"]["functions"]["operator_functions"][0][ "workspace" ] == [{"device": 1, "workspace_size_bytes": 2466816}] assert ( "fused_nn_conv2d_add_fixed_point_multiply_clip_cast" in metadata["modules"][module_name]["memory"]["functions"]["operator_functions"][0][ "function_name" ] )
def test_export_byoc_c_module(): """Test BYOC flow when it produces DSO-exportable modules. NOTE the general BYOC flow is not fully supported by Model Library Format right now. """ x = tvm.relay.var("x", shape=(10, 10)) w0 = tvm.relay.var("w0", shape=(10, 10)) w1 = tvm.relay.var("w1", shape=(10, 10)) w2 = tvm.relay.var("w2", shape=(10, 10)) w3 = tvm.relay.var("w3", shape=(10, 10)) w4 = tvm.relay.var("w4", shape=(10, 10)) w5 = tvm.relay.var("w5", shape=(10, 10)) w6 = tvm.relay.var("w6", shape=(10, 10)) w7 = tvm.relay.var("w7", shape=(10, 10)) # C compiler z0 = tvm.relay.add(x, w0) p0 = tvm.relay.subtract(z0, w1) q0 = tvm.relay.multiply(p0, w2) z1 = tvm.relay.add(x, w3) p1 = tvm.relay.subtract(z1, w4) q1 = tvm.relay.multiply(p1, w5) # Other parts on TVM z2 = tvm.relay.add(x, w6) q2 = tvm.relay.subtract(z2, w7) r = tvm.relay.concatenate((q0, q1, q2), axis=0) f = tvm.relay.Function([x, w0, w1, w2, w3, w4, w5, w6, w7], r) mod = tvm.IRModule() ann = byoc.CcompilerAnnotator() mod["main"] = ann.visit(f) mod = tvm.relay.transform.PartitionGraph("mod_name")(mod) mod = tvm.relay.transform.InferType()(mod) with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): factory = tvm.relay.build(mod, tvm.target.target.micro("host")) temp_dir = utils.tempdir() mlf_tar_path = temp_dir.relpath("lib.tar") from tvm import micro micro.export_model_library_format(factory, mlf_tar_path) with tarfile.open(mlf_tar_path, "r:*") as tf: tar_members = [ti.name for ti in tf.getmembers()] print("tar members", tar_members) assert "./metadata.json" in tar_members with tf.extractfile("./metadata.json") as f: metadata = json.load(f) main_md = metadata["memory"]["functions"]["main"] assert main_md == [{ "constants_size_bytes": 0, "device": 1, "io_size_bytes": 4800, "workspace_size_bytes": 800, }]
def test_multiple_relay_modules_c(): mod = get_conv2d_relay_module() executor = Executor("aot", {"unpacked-api": True, "interface-api": "c"}) runtime = Runtime("crt") target = tvm.target.target.micro("host") with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): factory1 = tvm.relay.build(mod, target, runtime=runtime, executor=executor, mod_name="mod1") factory2 = tvm.relay.build(mod, target, runtime=runtime, executor=executor, mod_name="mod2") temp_dir = utils.tempdir() mlf_tar_path = temp_dir.relpath("lib.tar") micro.export_model_library_format([factory1, factory2], mlf_tar_path) tf = tarfile.open(mlf_tar_path) extract_dir = temp_dir.relpath("extract") os.mkdir(extract_dir) tf.extractall(extract_dir) assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src", "mod1_lib0.c")) assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src", "mod1_lib1.c")) assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src", "mod2_lib0.c")) assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src", "mod2_lib1.c")) assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "include", "tvmgen_mod1.h")) assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "include", "tvmgen_mod2.h")) # check CRT runtime directory assert os.path.exists(os.path.join(extract_dir, "runtime"))
def build_project( temp_dir, zephyr_board, west_cmd, mod, build_config, simd=False, extra_files_tar=None ): project_dir = temp_dir / "project" with tempfile.TemporaryDirectory() as tar_temp_dir: model_tar_path = pathlib.Path(tar_temp_dir) / "model.tar" export_model_library_format(mod, model_tar_path) workspace_size = mlf_extract_workspace_size_bytes(model_tar_path) project_options = { "extra_files_tar": extra_files_tar, "project_type": "aot_demo", "west_cmd": west_cmd, "verbose": bool(build_config.get("debug")), "zephyr_board": zephyr_board, "compile_definitions": [ # TODO(mehrdadh): It fails without offset. f"-DWORKSPACE_SIZE={workspace_size + 128}", ], } if simd: project_options["config_main_stack_size"] = 1536 project = tvm.micro.project.generate_project_from_mlf( str(TEMPLATE_PROJECT_DIR), project_dir, model_tar_path, project_options ) project.build() return project, project_dir
def test_export_operator_model_library_format(): import tvm.micro as micro target = tvm.target.target.micro("host") with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): A = tvm.te.placeholder((2,), dtype="int8") B = tvm.te.placeholder((1,), dtype="int8") C = tvm.te.compute(A.shape, lambda i: A[i] + B[0], name="C") sched = tvm.te.create_schedule(C.op) mod = tvm.build( sched, [A, B, C], tvm.target.Target(target, target), runtime=Runtime("crt", {"system-lib": True}), name="add", ) temp_dir = utils.tempdir() mlf_tar_path = temp_dir.relpath("lib.tar") micro.export_model_library_format(mod, mlf_tar_path) tf = tarfile.open(mlf_tar_path) extract_dir = temp_dir.relpath("extract") os.mkdir(extract_dir) tf.extractall(extract_dir) with open(os.path.join(extract_dir, "metadata.json")) as json_f: metadata = json.load(json_f) assert metadata["version"] == 5 assert metadata["model_name"] == "add" export_datetime = datetime.datetime.strptime( metadata["export_datetime"], "%Y-%m-%d %H:%M:%SZ" ) assert (datetime.datetime.now() - export_datetime) < datetime.timedelta(seconds=60 * 5) assert metadata["target"] == {"1": str(target)} assert metadata["memory"]["add"][0]["dtype"] == "int8" assert metadata["memory"]["add"][0]["shape"] == [2] assert metadata["memory"]["add"][0]["size_bytes"] == 2 assert metadata["memory"]["add"][1]["dtype"] == "int8" assert metadata["memory"]["add"][1]["shape"] == [1] assert metadata["memory"]["add"][1]["size_bytes"] == 1 assert metadata["memory"]["add"][2]["dtype"] == "int8" assert metadata["memory"]["add"][2]["shape"] == [2] assert metadata["memory"]["add"][2]["size_bytes"] == 2 assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src", "lib0.c")) assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "src", "lib1.c")) assert ( len(mod.ir_module_by_target) == 1 ), f"expect 1 ir_model_by_target: {mod.ir_module_by_target!r}" for target, ir_mod in mod.ir_module_by_target.items(): assert int(tvm.runtime.ndarray.device(str(target)).device_type) == 1 with open(os.path.join(extract_dir, "src", "tir-1.txt")) as tir_f: assert tir_f.read() == str(ir_mod)
def test_export_model_library_format_llvm(): with utils.TempDirectory.set_keep_for_debug(True): target = tvm.target.target.micro("host") assert str(target)[:2] == "c " target = tvm.target.Target("llvm " + str(target)[2:]) with tvm.transform.PassContext(opt_level=3): relay_mod = tvm.parser.fromtext( """ #[version = "0.0.5"] def @main(%a : Tensor[(1, 2), uint8], %b : Tensor[(1, 2), float32], %c : Tensor[(1, 2), float32]) { %0 = cast(%a, dtype="float32") + %b * %c; %0 }""" ) factory = tvm.relay.build( relay_mod, target, target_host=target, mod_name="add", params={"c": numpy.array([[2.0, 4.0]], dtype="float32")}, ) temp_dir = utils.tempdir() mlf_tar_path = temp_dir.relpath("lib.tar") import tvm.micro as micro micro.export_model_library_format(factory, mlf_tar_path) tf = tarfile.open(mlf_tar_path) extract_dir = temp_dir.relpath("extract") os.mkdir(extract_dir) tf.extractall(extract_dir) with open(os.path.join(extract_dir, "metadata.json")) as json_f: metadata = json.load(json_f) assert metadata["version"] == 1 assert metadata["model_name"] == "add" export_datetime = datetime.datetime.strptime( metadata["export_datetime"], "%Y-%m-%d %H:%M:%SZ" ) assert (datetime.datetime.now() - export_datetime) < datetime.timedelta(seconds=60 * 5) assert metadata["target"] == {"1": str(target)} assert metadata["memory"] == [ {"storage_id": 0, "size_bytes": 2, "input_binding": "a"}, {"storage_id": 1, "size_bytes": 8, "input_binding": "b"}, {"storage_id": 2, "size_bytes": 8, "input_binding": "p0"}, {"storage_id": 3, "size_bytes": 8}, ] assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "lib", "lib0.o")) validate_graph_json(extract_dir, factory) with open(os.path.join(extract_dir, "relay.txt")) as relay_f: assert relay_f.read() == str(relay_mod) with open(os.path.join(extract_dir, "parameters", "add.params"), "rb") as params_f: params = tvm.relay.load_param_dict(params_f.read()) assert "p0" in params
def check_network(build_dir, target_name, model_path, image_path): model_name = "network" model, shape_dict, dtype_dict = get_tflite_model(model_path) # # Generate random input data # image_data = [] for i in range(NUM_ITERATIONS): assert len(shape_dict) == 1, "Only single input models are supported." image_shape = list(shape_dict.values())[0] in_data = np.random.randint(0, 255, size=image_shape).astype("uint8") # Write raw data for using with the TVM implementation filename = os.path.join(image_path, "{:02d}.raw".format(i)) dump_image(filename, in_data) image_data.append(in_data) mod, params = relay.frontend.from_tflite(model, shape_dict, dtype_dict) # # Build a TVM C module for the ARM CPU (without compiling the kernels # library to the object code form): # with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): rt_module = relay.build(mod, target="c -device=arm_cpu", params=params) # # Export model library format # target_dir = os.path.join(build_dir, target_name + "_gen") if os.path.exists(target_dir): print(f'Removing existing "{target_dir}" directory') try: shutil.rmtree(target_dir) except OSError as err: raise ValueError(f"emit_code.Error: {target_dir} : {err.strerror}") mlf_tar_path = os.path.join(build_dir, target_name + "_lib.tar") import tvm.micro as micro micro.export_model_library_format(rt_module, mlf_tar_path) emitter = stm32.CodeEmitter() quantization = extract_tflite_quantization(model) emitter.parse_library_format(mlf_tar_path, quantization) emitter.emit_code(target_dir, model_name) # # Results # tf_results = run_tflite_model(model_path, image_data) tvm_results = run_tvm_model(build_dir, model_name, target_dir, image_path) check_result(tf_results, tvm_results)
def test_multiple_relay_modules_aot_graph(): mod = get_conv2d_relay_module() executor1 = Executor("graph") executor2 = Executor("aot", {"unpacked-api": True, "interface-api": "c"}) runtime = Runtime("crt") target = tvm.target.target.micro("host") with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): factory1 = tvm.relay.build(mod, target, runtime=runtime, executor=executor1, mod_name="mod1") factory2 = tvm.relay.build(mod, target, runtime=runtime, executor=executor2, mod_name="mod2") temp_dir = utils.tempdir() mlf_tar_path = temp_dir.relpath("lib.tar") micro.export_model_library_format([factory1, factory2], mlf_tar_path) tf = tarfile.open(mlf_tar_path) extract_dir = temp_dir.relpath("extract") os.mkdir(extract_dir) tf.extractall(extract_dir) assert os.path.exists( os.path.join(extract_dir, "codegen", "host", "src", "mod1_lib0.c")) assert os.path.exists( os.path.join(extract_dir, "codegen", "host", "src", "mod1_lib1.c")) assert os.path.exists( os.path.join(extract_dir, "codegen", "host", "src", "mod1_lib2.c")) assert os.path.exists( os.path.join(extract_dir, "codegen", "host", "src", "mod2_lib0.c")) assert os.path.exists( os.path.join(extract_dir, "codegen", "host", "src", "mod2_lib1.c")) assert os.path.exists( os.path.join(extract_dir, "codegen", "host", "include", "tvmgen_mod2.h")) with open(os.path.join(extract_dir, "metadata.json")) as f: metadata = json.load(f) assert metadata["modules"]["mod1"]["executors"] == ["graph"] assert metadata["modules"]["mod2"]["executors"] == ["aot"] assert metadata["version"] == _GENERATED_VERSION
def compile_and_run(mod, input_list, output_list, params=None): """ This method verifies the generated source """ target = "c -runtime=c --link-params --executor=aot" with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): lib = tvm.relay.build(mod, target, target_host=target, params=params) tmp_path = utils.tempdir() tmp_dir = tmp_path.temp_dir base_path = os.path.join(tmp_dir, "test") build_path = os.path.join(base_path, "build") os.makedirs(build_path, exist_ok=True) tar_file = os.path.join(base_path, "test.tar") export_model_library_format(lib, tar_file) t = tarfile.open(tar_file) t.extractall(base_path) for i in range(len(input_list)): create_header_file((f"input_data{i}"), input_list[i], build_path) for i in range(len(output_list)): create_header_file( (f"output_data{i}"), np.zeros(output_list[i].shape, output_list[i].dtype), build_path, ) create_header_file((f"expected_output_data{i}"), output_list[i], build_path) create_main("test.c", input_list, output_list, build_path) # Verify that compiles fine file_dir = os.path.dirname(os.path.abspath(__file__)) makefile = os.path.join(file_dir, "aot_test.mk") make_cmd = f"make -f {makefile} build_dir=" + build_path + f" TVM_ROOT={file_dir}/../../../.." compile_log_path = os.path.join(build_path, "test_compile.log") ret = subprocess_with_stdout_and_log(make_cmd, ".", compile_log_path, False) assert ret == 0 # Verify that runs fine run_log_path = os.path.join(build_path, "test_run.log") ret = subprocess_with_stdout_and_log("./aot_test_runner", build_path, run_log_path, False) assert ret == 0
def export_package( self, executor_factory: Union[GraphExecutorFactoryModule, Executable], package_path: Optional[str] = None, cross: Optional[Union[str, Callable]] = None, cross_options: Optional[str] = None, output_format: str = "so", ): """Save this TVMCModel to file. Parameters ---------- executor_factory : GraphExecutorFactoryModule The factory containing the compiled artifacts needed to run this model. package_path : str, None Where the model should be saved. Note that it will be packaged as a .tar file. If not provided, the package will be saved to a generically named file in tmp. cross : str or callable object, optional Function that performs the actual compilation. cross_options : str, optional Command line options to be passed to the cross compiler. output_format : str How to save the modules function library. Must be one of "so" and "tar" to save using the classic format or "mlf" to save using the Model Library Format. Returns ------- package_path : str The path that the package was saved to. """ if output_format not in ["so", "tar", "mlf"]: raise TVMCException( "Only 'so', 'tar', and 'mlf' output formats are supported.") if output_format == "mlf" and cross: raise TVMCException( "Specifying the MLF output and a cross compiler is not supported." ) if isinstance(executor_factory, Executable): package_path = self.export_vm_format(executor_factory, package_path, output_format) elif output_format in ["so", "tar"]: package_path = self.export_classic_format(executor_factory, package_path, cross, cross_options, output_format) elif output_format == "mlf": if export_model_library_format: package_path = export_model_library_format( executor_factory, package_path) else: raise Exception( "micro tvm is not enabled. Set USE_MICRO to ON in config.cmake" ) return package_path
def test_export_multiple_operator_model_library_format(): target = tvm.target.target.micro("host") with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): A = tvm.te.placeholder((2,), dtype="int8") B = tvm.te.placeholder((1,), dtype="int8") C = tvm.te.compute(A.shape, lambda i: A[i] + B[0], name="C") sched = tvm.te.create_schedule(C.op) mod = tvm.build( sched, [A, B, C], tvm.target.Target(target, target), runtime=Runtime("crt", {"system-lib": True}), name="add", ) temp_dir = utils.tempdir() mlf_tar_path = temp_dir.relpath("lib.tar") with pytest.raises(RuntimeError) as exc: micro.export_model_library_format([mod, mod], mlf_tar_path) assert str(exc.exception) == ("Multiple operator is not supported.")
def test_packed_global_variables(): """Check packed global variables in codegen output.""" dtype = "float32" ishape = (1, 32, 14, 14) wshape = (32, 32, 3, 3) interface_api = "packed" use_unpacked_api = False data0 = relay.var("data", shape=ishape, dtype=dtype) weight0 = relay.var("weight", shape=wshape, dtype=dtype) out = relay.nn.conv2d(data0, weight0, kernel_size=(3, 3), padding=(1, 1), groups=1) main_f = relay.Function([data0, weight0], out) mod = tvm.IRModule() mod["main"] = main_f mod = transform.InferType()(mod) i_data = np.random.uniform(0, 1, ishape).astype(dtype) w1_data = np.random.uniform(0, 1, wshape).astype(dtype) inputs = OrderedDict([("data", i_data), ("weight", w1_data)]) output_list = generate_ref_data(mod, inputs) compiled_models_list = compile_models( models=AOTTestModel(module=mod, inputs=inputs, outputs=output_list), interface_api=interface_api, use_unpacked_api=use_unpacked_api, workspace_byte_alignment=8, enable_op_fusion=True, pass_config=AOT_DEFAULT_RUNNER.pass_config, use_runtime_executor=True, target=tvm.target.Target("c"), ) compiled_model = compiled_models_list[0] tmp_path = utils.tempdir() base_path = tmp_path.temp_dir model = compiled_model.model tar_file = os.path.join(base_path, f"{model.name}.tar") export_model_library_format(compiled_model.executor_factory, tar_file) t = tarfile.open(tar_file) t.extractall(base_path) file_list = [] for path in (pathlib.Path(base_path) / "codegen" / "host" / "src").iterdir(): if path.is_file(): file_list.append(path) assert len(file_list) > 0 for path in file_list: with open(path, "r") as lib_f: lib1 = lib_f.readlines() tvmgen_names = [] tvmgen_funcs = [] for line in lib1: for item in line.split(" "): # Find all names starting with tvmgen_default if item.startswith("tvmgen_default"): # Collect any name starting with tvmgen_default tvmgen_names.append(item) # Collect all functions starting with tvmgen_default tvmgen_funcs += re.findall(r"(?<=).*(?=\()", item) # Check if any function name has a packed variable name in all items that start with tvmgen_default for func in tvmgen_funcs: assert f"{func}_packed" not in tvmgen_names
def run_and_check( models: List[AOTCompiledTestModel], runner: AOTTestRunner, interface_api: str, debug_calculated_workspaces=False, workspace_byte_alignment=8, data_linkage: AOTDataLinkage = None, test_dir: str = None, verbose: bool = False, ): """ This method uses the original test data and compiled runtime.Modules to run in the test runner to verify the results. """ base_path = test_dir if test_dir is None: tmp_path = utils.tempdir() tmp_dir = tmp_path.temp_dir base_path = os.path.join(tmp_dir, "test") cflags = f"-DTVM_RUNTIME_ALLOC_ALIGNMENT_BYTES={workspace_byte_alignment} " # The calculated workspaces will not account for stack allocator tags used for debugging if debug_calculated_workspaces: cflags += "-DTVM_CRT_STACK_ALLOCATOR_ENABLE_LIFO_CHECK " base_path = os.path.abspath(base_path) build_path = os.path.join(base_path, "build") os.makedirs(build_path, exist_ok=True) include_path = os.path.join(base_path, "include") os.mkdir(include_path) crt_root = tvm.micro.get_standalone_crt_dir() shutil.copy2( os.path.join(crt_root, "template", "crt_config-template.h"), os.path.join(include_path, "crt_config.h"), ) workspace_bytes = 0 for compiled_model in models: model = compiled_model.model tar_file = os.path.join(base_path, f"{model.name}.tar") export_model_library_format(compiled_model.executor_factory, tar_file) t = tarfile.open(tar_file) t.extractall(base_path) workspace_bytes = model.extra_memory_in_bytes use_usmp = runner.pass_config.get("tir.usmp.enable", False) if interface_api == "packed" and not use_usmp: workspace_bytes += mlf_extract_workspace_size_bytes(tar_file) for key in model.inputs: sanitized_tensor_name = re.sub(r"\W", "_", key) create_header_file( f'{mangle_name(model.name, "input_data")}_{sanitized_tensor_name}', model.inputs[key], include_path, data_linkage, ) for key in model.outputs: sanitized_tensor_name = re.sub(r"\W", "_", key) create_header_file( f'{mangle_name(model.name, "output_data")}_{sanitized_tensor_name}', np.zeros(model.outputs[key].shape, model.outputs[key].dtype), include_path, data_linkage, ) create_header_file( f'{mangle_name(model.name, "expected_output_data")}_{sanitized_tensor_name}', model.outputs[key], include_path, data_linkage, ) create_main( "test.c", models, build_path, runner.includes, runner.prologue, runner.epilogue, data_linkage, interface_api, workspace_bytes, ) # Verify that compiles fine file_dir = os.path.dirname(os.path.abspath(__file__)) codegen_path = os.path.join(base_path, "codegen") makefile = os.path.join(file_dir, f"{runner.makefile}.mk") fvp_dir = "/opt/arm/FVP_Corstone_SSE-300/models/Linux64_GCC-6.4/" # TODO(@grant-arm): Remove once ci_cpu docker image has been updated to FVP_Corstone_SSE if not os.path.isdir(fvp_dir): fvp_dir = "/opt/arm/FVP_Corstone_SSE-300_Ethos-U55/models/Linux64_GCC-6.4/" custom_params = " ".join( [f" {param}='{value}'" for param, value in runner.parameters.items()]) make_command = ( f"make -f {makefile} build_dir={build_path}" + f" CFLAGS='{cflags}'" + f" TVM_ROOT={file_dir}/../../../.." + f" AOT_TEST_ROOT={file_dir}" + f" CODEGEN_ROOT={codegen_path}" + f" STANDALONE_CRT_DIR={tvm.micro.get_standalone_crt_dir()}" + f" FVP_DIR={fvp_dir}" + custom_params) compile_log_path = os.path.join(build_path, "test_compile.log") compile_command = f"{make_command} aot_test_runner" if verbose: print("Compile command:\n", compile_command) ret = subprocess_log_output(compile_command, ".", compile_log_path) assert ret == 0 # Verify that runs fine run_log_path = os.path.join(build_path, "test_run.log") run_command = f"{make_command} run" if verbose: print("Run command:\n", run_command) ret = subprocess_log_output(run_command, build_path, run_log_path) assert ret == 0 with open(run_log_path) as run_log: assert AOT_SUCCESS_TOKEN in run_log.read()
def test_export_model_library_format_llvm(): with utils.TempDirectory.set_keep_for_debug(True): target = tvm.target.target.micro("host") assert str(target)[:2] == "c " target = tvm.target.Target("llvm " + str(target)[2:]) with tvm.transform.PassContext(opt_level=3): relay_mod = tvm.parser.fromtext( """ #[version = "0.0.5"] def @main(%a : Tensor[(1, 2), uint8], %b : Tensor[(1, 2), float32], %c : Tensor[(1, 2), float32]) { %0 = cast(%a, dtype="float32") + %b * %c; %0 }""" ) factory = tvm.relay.build( relay_mod, target, runtime=Runtime("crt", {"system-lib": True}), mod_name="add", params={"c": np.array([[2.0, 4.0]], dtype="float32")}, ) temp_dir = utils.tempdir() mlf_tar_path = temp_dir.relpath("lib.tar") micro.export_model_library_format(factory, mlf_tar_path) tf = tarfile.open(mlf_tar_path) extract_dir = temp_dir.relpath("extract") os.mkdir(extract_dir) tf.extractall(extract_dir) with open(os.path.join(extract_dir, "metadata.json")) as json_f: metadata = json.load(json_f) module_name = factory.libmod_name assert metadata["version"] == _GENERATED_VERSION assert metadata["modules"][module_name]["model_name"] == "add" export_datetime = datetime.datetime.strptime( metadata["modules"][module_name]["export_datetime"], "%Y-%m-%d %H:%M:%SZ" ) assert (datetime.datetime.now() - export_datetime) < datetime.timedelta(seconds=60 * 5) assert metadata["modules"][module_name]["target"] == [str(target)] assert metadata["modules"][module_name]["memory"]["sids"] == [ {"storage_id": 0, "size_bytes": 2, "input_binding": "a"}, {"storage_id": 1, "size_bytes": 8, "input_binding": "b"}, {"storage_id": 2, "size_bytes": 8, "input_binding": "p0"}, {"storage_id": 3, "size_bytes": 8}, ] assert metadata["modules"][module_name]["memory"]["functions"]["main"] == [ { "constants_size_bytes": 8, "device": 1, "io_size_bytes": 18, "workspace_size_bytes": 0, } ] assert metadata["modules"][module_name]["memory"]["functions"]["operator_functions"][0][ "workspace" ] == [{"device": 1, "workspace_size_bytes": 0}] assert ( "fused_cast_multiply_add" in metadata["modules"][module_name]["memory"]["functions"]["operator_functions"][0][ "function_name" ] ) assert os.path.exists(os.path.join(extract_dir, "codegen", "host", "lib", "add_lib0.o")) validate_graph_json(extract_dir, factory) with open(os.path.join(extract_dir, "src", f"{module_name}.relay")) as relay_f: assert relay_f.read() == str(relay_mod) with open(os.path.join(extract_dir, "parameters", "add.params"), "rb") as params_f: params = tvm.relay.load_param_dict(params_f.read()) assert "p0" in params
def run_and_check_body(base_path): cflags = f"-DTVM_RUNTIME_ALLOC_ALIGNMENT_BYTES={workspace_byte_alignment} " # The calculated workspaces will not account for stack allocator tags used for debugging if debug_calculated_workspaces: cflags += "-DTVM_CRT_STACK_ALLOCATOR_ENABLE_LIFO_CHECK " base_path = os.path.abspath(base_path) build_path = os.path.join(base_path, "build") os.makedirs(build_path, exist_ok=True) include_path = os.path.join(base_path, "include") os.mkdir(include_path) crt_root = tvm.micro.get_standalone_crt_dir() shutil.copy2( os.path.join(crt_root, "template", "crt_config-template.h"), os.path.join(include_path, "crt_config.h"), ) workspace_bytes = 0 for compiled_model in models: model = compiled_model.model tar_file = os.path.join(base_path, f"{model.name}.tar") export_model_library_format(compiled_model.executor_factory, tar_file) t = tarfile.open(tar_file) t.extractall(base_path) # Interface C APIs does not need compiler generated # workspace to generate the test application, because # workspace size is codegen'd as a macro to # tvmgen_<model_name>.h. if interface_api != "c": workspace_bytes += mlf_extract_workspace_size_bytes(tar_file) workspace_bytes += model.extra_memory_in_bytes for key in model.inputs: sanitized_tensor_name = re.sub(r"\W", "_", key) create_header_file( f'{mangle_name(model.name, "input_data")}_{sanitized_tensor_name}', model.inputs[key], include_path, data_linkage, ) for key in model.outputs: sanitized_tensor_name = re.sub(r"\W", "_", key) create_header_file( f'{mangle_name(model.name, "output_data")}_{sanitized_tensor_name}', np.zeros(model.outputs[key].shape, model.outputs[key].dtype), include_path, data_linkage, ) create_header_file( f'{mangle_name(model.name, "expected_output_data")}_{sanitized_tensor_name}', model.outputs[key], include_path, data_linkage, ) use_usmp = runner.pass_config.get("tir.usmp.enable", False) # We only need the stack allocator if USMP is not used use_stack_allocator = not use_usmp create_main( "test.c", models, build_path, runner.includes, runner.prologue, runner.epilogue, data_linkage, interface_api, workspace_bytes, use_stack_allocator, ) # Verify that compiles fine file_dir = os.path.dirname(os.path.abspath(__file__)) codegen_path = os.path.join(base_path, "codegen") makefile = os.path.join(file_dir, f"{runner.makefile}.mk") fvp_dir = "/opt/arm/FVP_Corstone_SSE-300/models/Linux64_GCC-6.4/" # TODO(@grant-arm): Remove once ci_cpu docker image has been updated to FVP_Corstone_SSE if not os.path.isdir(fvp_dir): fvp_dir = "/opt/arm/FVP_Corstone_SSE-300_Ethos-U55/models/Linux64_GCC-6.4/" custom_params = " ".join([ f" {param}='{value}'" for param, value in runner.parameters.items() ]) make_command = ( f"make -f {makefile} build_dir={build_path}" + f" CFLAGS='{cflags}'" + f" TVM_ROOT={file_dir}/../../../.." + f" AOT_TEST_ROOT={file_dir}" + f" CODEGEN_ROOT={codegen_path}" + f" STANDALONE_CRT_DIR={tvm.micro.get_standalone_crt_dir()}" + f" FVP_DIR={fvp_dir}" + custom_params) compile_log_path = os.path.join(build_path, "test_compile.log") compile_command = f"{make_command} aot_test_runner" if verbose: print("Compile command:\n", compile_command) subprocess_check_log_output(compile_command, ".", compile_log_path) # Verify that runs fine run_log_path = os.path.join(build_path, "test_run.log") run_command = f"{make_command} run" if verbose: print("Run command:\n", run_command) # TODO(lhutton1) This is a quick and dirty work around to help temporarily reduce # the flakyness of the tests. Will remove once #10300 and #10314 are resolved. try: subprocess_check_log_output(run_command, build_path, run_log_path) except RuntimeError as err: print("Failed to run the module, having a second attempt...", file=sys.stderr) print(err, file=sys.stderr) subprocess_check_log_output(run_command, build_path, run_log_path) with open(run_log_path) as run_log: assert AOT_SUCCESS_TOKEN in run_log.read()
def run_and_check( models: List[AOTCompiledTestModel], runner: AOTTestRunner, interface_api: str, debug_calculated_workspaces=False, workspace_byte_alignment=8, data_linkage: AOTDataLinkage = None, ): """ This method uses the original test data and compiled runtime.Modules to run in the test runner to verify the results. """ tmp_path = utils.tempdir() tmp_dir = tmp_path.temp_dir cflags = f"-DTVM_RUNTIME_ALLOC_ALIGNMENT_BYTES={workspace_byte_alignment} " # The calculated workspaces will not account for stack allocator tags used for debugging if debug_calculated_workspaces: cflags += "-DTVM_CRT_STACK_ALLOCATOR_ENABLE_LIFO_CHECK " base_path = os.path.join(tmp_dir, "test") build_path = os.path.join(base_path, "build") os.makedirs(build_path, exist_ok=True) include_path = os.path.join(base_path, "include") os.mkdir(include_path) crt_root = tvm.micro.get_standalone_crt_dir() shutil.copy2( os.path.join(crt_root, "template", "crt_config-template.h"), os.path.join(include_path, "crt_config.h"), ) workspace_bytes = 0 for compiled_model in models: model = compiled_model.model tar_file = os.path.join(base_path, f"{model.name}.tar") export_model_library_format(compiled_model.executor_factory, tar_file) t = tarfile.open(tar_file) t.extractall(base_path) workspace_bytes += model.extra_memory_in_bytes workspace_bytes += extract_main_workspace_size_bytes(base_path) for key in model.inputs: sanitized_tensor_name = re.sub(r"\W", "_", key) create_header_file( f'{mangle_name(model.name, "input_data")}_{sanitized_tensor_name}', model.inputs[key], include_path, data_linkage, ) for i in range(len(model.outputs)): create_header_file( (f'{mangle_name(model.name,"output_data")}{i}'), np.zeros(model.outputs[i].shape, model.outputs[i].dtype), include_path, data_linkage, ) create_header_file( (f'{mangle_name(model.name, "expected_output_data")}{i}'), model.outputs[i], include_path, data_linkage, ) create_main( "test.c", [compiled_model.model for compiled_model in models], build_path, runner.includes, runner.prologue, data_linkage, interface_api, workspace_bytes, ) # Verify that compiles fine file_dir = os.path.dirname(os.path.abspath(__file__)) codegen_path = os.path.join(base_path, "codegen") makefile = os.path.join(file_dir, f"{runner.makefile}.mk") custom_params = " ".join([f" {param}='{value}'" for param, value in runner.parameters.items()]) make_command = ( f"make -f {makefile} build_dir={build_path}" + f" CFLAGS='{cflags}'" + f" TVM_ROOT={file_dir}/../../../.." + f" AOT_TEST_ROOT={file_dir}" + f" CODEGEN_ROOT={codegen_path}" + f" STANDALONE_CRT_DIR={tvm.micro.get_standalone_crt_dir()}" + custom_params ) compile_log_path = os.path.join(build_path, "test_compile.log") compile_command = f"{make_command} aot_test_runner" ret = subprocess_log_output(compile_command, ".", compile_log_path) assert ret == 0 # Verify that runs fine run_log_path = os.path.join(build_path, "test_run.log") run_command = f"{make_command} run" ret = subprocess_log_output(run_command, build_path, run_log_path) assert ret == 0 with open(run_log_path) as run_log: assert AOT_SUCCESS_TOKEN in run_log.read()
def compile_and_run(mod, input_list, output_list, use_calculated_workspaces, params=None, workspace_byte_alignment=8): """ This method verifies the generated source """ target = f"c -runtime=c --link-params --executor=aot --workspace-byte-alignment={workspace_byte_alignment}" cflags = f"-DTVM_RUNTIME_ALLOC_ALIGNMENT_BYTES={workspace_byte_alignment} " # The calculated workspaces will not account for stack allocator tags used for debugging if not use_calculated_workspaces: cflags += "-DTVM_CRT_STACK_ALLOCATOR_ENABLE_LIFO_CHECK " with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): lib = tvm.relay.build(mod, target, target_host=target, params=params) tmp_path = utils.tempdir() tmp_dir = tmp_path.temp_dir base_path = os.path.join(tmp_dir, "test") build_path = os.path.join(base_path, "build") os.makedirs(build_path, exist_ok=True) tar_file = os.path.join(base_path, "test.tar") export_model_library_format(lib, tar_file) t = tarfile.open(tar_file) t.extractall(base_path) if use_calculated_workspaces: workspace_bytes = extract_main_workspace_sizebytes(base_path) else: workspace_bytes = 16384 * 1024 for i in range(len(input_list)): create_header_file((f"input_data{i}"), input_list[i], build_path) for i in range(len(output_list)): create_header_file( (f"output_data{i}"), np.zeros(output_list[i].shape, output_list[i].dtype), build_path, ) create_header_file((f"expected_output_data{i}"), output_list[i], build_path) create_main("test.c", input_list, output_list, build_path, workspace_bytes) # Verify that compiles fine file_dir = os.path.dirname(os.path.abspath(__file__)) makefile = os.path.join(file_dir, "aot_test.mk") make_cmd = (f"make CFLAGS='{cflags}' -f {makefile} build_dir=" + build_path + f" TVM_ROOT={file_dir}/../../../..") compile_log_path = os.path.join(build_path, "test_compile.log") ret = subprocess_with_stdout_and_log(make_cmd, ".", compile_log_path, False) assert ret == 0 # Verify that runs fine run_log_path = os.path.join(build_path, "test_run.log") ret = subprocess_with_stdout_and_log("./aot_test_runner", build_path, run_log_path, False) assert ret == 0
def compile_and_run( models: Union[List[AOTTestModel], AOTTestModel], runner: AOTTestRunner, interface_api, use_unpacked_api, debug_calculated_workspaces=False, workspace_byte_alignment=8, enable_op_fusion=True, ): """ This method verifies the generated source """ base_target = "c -runtime=c --link-params --executor=aot" extra_target = f"--workspace-byte-alignment={workspace_byte_alignment} --interface-api={interface_api} --unpacked-api={int(use_unpacked_api)}" target = f"{base_target} {extra_target}" cflags = f"-DTVM_RUNTIME_ALLOC_ALIGNMENT_BYTES={workspace_byte_alignment} " if not isinstance(models, list): models = [models] # The calculated workspaces will not account for stack allocator tags used for debugging if debug_calculated_workspaces: cflags += "-DTVM_CRT_STACK_ALLOCATOR_ENABLE_LIFO_CHECK " config = {"tir.disable_vectorize": True} if not enable_op_fusion: config["relay.FuseOps.max_depth"] = 1 tmp_path = utils.tempdir() tmp_dir = tmp_path.temp_dir base_path = os.path.join(tmp_dir, "test") build_path = os.path.join(base_path, "build") os.makedirs(build_path, exist_ok=True) include_path = os.path.join(base_path, "include") os.mkdir(include_path) crt_root = tvm.micro.get_standalone_crt_dir() shutil.copy2( os.path.join(crt_root, "template", "crt_config-template.h"), os.path.join(include_path, "crt_config.h"), ) workspace_bytes = 0 for model in models: with tvm.transform.PassContext(opt_level=3, config=config): lib = tvm.relay.build( model.module, target, target_host=target, params=model.params, mod_name=model.name, ) tar_file = os.path.join(base_path, f"{model.name}.tar") export_model_library_format(lib, tar_file) t = tarfile.open(tar_file) t.extractall(base_path) workspace_bytes += extract_main_workspace_size_bytes(base_path) for key in model.inputs: sanitized_tensor_name = re.sub(r"\W", "_", key) create_header_file( f'{mangle_name(model.name, "input_data")}_{sanitized_tensor_name}', model.inputs[key], include_path, ) for i in range(len(model.outputs)): create_header_file( (f'{mangle_name(model.name,"output_data")}{i}'), np.zeros(model.outputs[i].shape, model.outputs[i].dtype), include_path, ) create_header_file( (f'{mangle_name(model.name, "expected_output_data")}{i}'), model.outputs[i], include_path, ) create_main( "test.c", models, build_path, runner.includes, runner.prologue, interface_api, workspace_bytes, ) # Verify that compiles fine file_dir = os.path.dirname(os.path.abspath(__file__)) codegen_path = os.path.join(base_path, "codegen") makefile = os.path.join(file_dir, f"{runner.makefile}.mk") custom_params = " ".join( [f" {param}='{value}'" for param, value in runner.parameters.items()]) make_command = ( f"make -f {makefile} build_dir={build_path}" + f" CFLAGS='{cflags}'" + f" TVM_ROOT={file_dir}/../../../.." + f" AOT_TEST_ROOT={file_dir}" + f" CODEGEN_ROOT={codegen_path}" + f" STANDALONE_CRT_DIR={tvm.micro.get_standalone_crt_dir()}" + custom_params) compile_log_path = os.path.join(build_path, "test_compile.log") compile_command = f"{make_command} aot_test_runner" ret = subprocess_log_output(compile_command, ".", compile_log_path) assert ret == 0 # Verify that runs fine run_log_path = os.path.join(build_path, "test_run.log") run_command = f"{make_command} run" ret = subprocess_log_output(run_command, build_path, run_log_path) assert ret == 0 with open(run_log_path) as run_log: assert AOT_SUCCESS_TOKEN in run_log.read()
def compile_and_run( mod, inputs, output_list, interface_api, use_unpacked_api, use_calculated_workspaces, params=None, workspace_byte_alignment=8, mod_name="default", enable_op_fusion=True, ): """ This method verifies the generated source """ base_target = "c -runtime=c --link-params --executor=aot" extra_target = f"--workspace-byte-alignment={workspace_byte_alignment} --interface-api={interface_api} --unpacked-api={int(use_unpacked_api)}" target = f"{base_target} {extra_target}" cflags = f"-DTVM_RUNTIME_ALLOC_ALIGNMENT_BYTES={workspace_byte_alignment} " # The calculated workspaces will not account for stack allocator tags used for debugging if not use_calculated_workspaces: cflags += "-DTVM_CRT_STACK_ALLOCATOR_ENABLE_LIFO_CHECK " config = {"tir.disable_vectorize": True} if not enable_op_fusion: config["relay.FuseOps.max_depth"] = 1 with tvm.transform.PassContext(opt_level=3, config=config): lib = tvm.relay.build(mod, target, target_host=target, params=params, mod_name=mod_name) tmp_path = utils.tempdir() tmp_dir = tmp_path.temp_dir base_path = os.path.join(tmp_dir, "test") build_path = os.path.join(base_path, "build") os.makedirs(build_path, exist_ok=True) tar_file = os.path.join(base_path, "test.tar") export_model_library_format(lib, tar_file) t = tarfile.open(tar_file) t.extractall(base_path) if use_calculated_workspaces: workspace_bytes = extract_main_workspace_sizebytes(base_path) else: workspace_bytes = 16384 * 1024 include_path = os.path.join(base_path, "include") os.mkdir(include_path) crt_root = tvm.micro.get_standalone_crt_dir() shutil.copy2( os.path.join(crt_root, "template", "crt_config-template.h"), os.path.join(include_path, "crt_config.h"), ) for key in inputs: create_header_file( f'{mangle_name(mod_name, "input_data")}_{key}', inputs[key], os.path.join(base_path, "include"), ) for i in range(len(output_list)): create_header_file( f'{mangle_name(mod_name,"output_data")}{i}', np.zeros(output_list[i].shape, output_list[i].dtype), os.path.join(base_path, "include"), ) create_header_file( f'{mangle_name(mod_name, "expected_output_data")}{i}', output_list[i], os.path.join(base_path, "include"), ) create_main( "test.c", {mod_name: inputs}, {mod_name: output_list}, build_path, interface_api, workspace_bytes, ) # Verify that compiles fine file_dir = os.path.dirname(os.path.abspath(__file__)) codegen_path = os.path.join(base_path, "codegen") makefile = os.path.join(file_dir, "aot_test.mk") make_cmd = (f"make CFLAGS='{cflags}' -f {makefile} build_dir=" + build_path + f" TVM_ROOT={file_dir}/../../../.." + f" CODEGEN_ROOT={codegen_path}" + f" STANDALONE_CRT_DIR={tvm.micro.get_standalone_crt_dir()}") compile_log_path = os.path.join(build_path, "test_compile.log") ret = subprocess_log_output(make_cmd, ".", compile_log_path) assert ret == 0 # Verify that runs fine run_log_path = os.path.join(build_path, "test_run.log") ret = subprocess_log_output("./aot_test_runner", build_path, run_log_path) assert ret == 0
def compile_and_run_multiple_models( mod_map, input_list_map, output_list_map, interface_api, use_unpacked_api, use_calculated_workspaces, param_map, workspace_byte_alignment=8, ): """ This method verifies the generated source """ base_target = "c -runtime=c --link-params --executor=aot" extra_target = f"--workspace-byte-alignment={workspace_byte_alignment} --interface-api={interface_api} --unpacked-api={int(use_unpacked_api)}" target = f"{base_target} {extra_target}" tmp_path = utils.tempdir() tmp_dir = tmp_path.temp_dir base_path = os.path.join(tmp_dir, "test") build_path = os.path.join(base_path, "build") os.makedirs(build_path, exist_ok=True) include_path = os.path.join(base_path, "include") os.mkdir(include_path) crt_root = tvm.micro.get_standalone_crt_dir() shutil.copy2( os.path.join(crt_root, "template", "crt_config-template.h"), os.path.join(include_path, "crt_config.h"), ) for mod_name, mod in mod_map.items(): with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): lib = tvm.relay.build(mod, target, target_host=target, params=param_map[mod_name], mod_name=mod_name) tar_file = os.path.join(base_path, "test.tar") export_model_library_format(lib, tar_file) t = tarfile.open(tar_file) t.extractall(base_path) input_list = input_list_map[mod_name] output_list = output_list_map[mod_name] for key in input_list: create_header_file((f'{mangle_name(mod_name,"input_data")}_{key}'), input_list[key], build_path) for i in range(len(output_list_map[mod_name])): create_header_file( (f'{mangle_name(mod_name,"output_data")}{i}'), np.zeros(output_list[i].shape, output_list[i].dtype), build_path, ) create_header_file( (f'{mangle_name(mod_name,"expected_output_data")}{i}'), output_list[i], build_path) create_main( "test.c", input_list_map, output_list_map, build_path, interface_api, workspace_bytes=16384 * 1024, ) # Verify that compiles fine file_dir = os.path.dirname(os.path.abspath(__file__)) codegen_path = os.path.join(base_path, "codegen") makefile = os.path.join(file_dir, "aot_test.mk") make_cmd = (f"make -f {makefile} build_dir=" + build_path + f" TVM_ROOT={file_dir}/../../../.." + f" CODEGEN_ROOT={codegen_path}" + f" STANDALONE_CRT_DIR={tvm.micro.get_standalone_crt_dir()}") compile_log_path = os.path.join(build_path, "test_compile.log") ret = subprocess_log_output(make_cmd, ".", compile_log_path) assert ret == 0 # Verify that runs fine run_log_path = os.path.join(build_path, "test_run.log") ret = subprocess_log_output("./aot_test_runner", build_path, run_log_path) assert ret == 0
def test_export_model_library_format_c(target): executor, _target = target with utils.TempDirectory.set_keep_for_debug(True): with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): relay_mod = tvm.parser.fromtext(""" #[version = "0.0.5"] def @main(%a : Tensor[(1, 2), uint8], %b : Tensor[(1, 2), float32], %c : Tensor[(1, 2), float32]) { %0 = cast(%a, dtype="float32") + %b * %c; %0 }""") factory = tvm.relay.build( relay_mod, _target, target_host=_target, mod_name="add", params={"c": numpy.array([[2.0, 4.0]], dtype="float32")}, ) temp_dir = utils.tempdir() mlf_tar_path = temp_dir.relpath("lib.tar") import tvm.micro as micro micro.export_model_library_format(factory, mlf_tar_path) tf = tarfile.open(mlf_tar_path) extract_dir = temp_dir.relpath("extract") os.mkdir(extract_dir) tf.extractall(extract_dir) with open(os.path.join(extract_dir, "metadata.json")) as json_f: metadata = json.load(json_f) assert metadata["version"] == 3 assert metadata["model_name"] == "add" export_datetime = datetime.datetime.strptime( metadata["export_datetime"], "%Y-%m-%d %H:%M:%SZ") assert (datetime.datetime.now() - export_datetime) < datetime.timedelta(seconds=60 * 5) assert metadata["target"] == {"1": str(_target)} if executor == "graph": assert metadata["memory"]["sids"] == [ { "storage_id": 0, "size_bytes": 2, "input_binding": "a" }, { "storage_id": 1, "size_bytes": 8, "input_binding": "b" }, { "storage_id": 2, "size_bytes": 8, "input_binding": "p0" }, { "storage_id": 3, "size_bytes": 8 }, ] assert metadata["memory"]["functions"]["main"] == [{ "constants_size_bytes": 8, "device": 1, "io_size_bytes": 18, "workspace_size_bytes": 0, }] assert metadata["memory"]["functions"]["operator_functions"][0][ "workspace"] == [{ "device": 1, "workspace_size_bytes": 0 }] assert ("fused_cast_multiply_add" in metadata["memory"] ["functions"]["operator_functions"][0]["function_name"]) assert os.path.exists( os.path.join(extract_dir, "codegen", "host", "src", "add_lib0.c")) assert os.path.exists( os.path.join(extract_dir, "codegen", "host", "src", "add_lib1.c")) if executor == "graph": validate_graph_json(extract_dir, factory) with open(os.path.join(extract_dir, "relay.txt")) as relay_f: assert relay_f.read() == str(relay_mod) with open(os.path.join(extract_dir, "parameters", "add.params"), "rb") as params_f: params = tvm.relay.load_param_dict(params_f.read()) assert "p0" in params