From c93d9197e5beabb9518c2b62a331dab893dc5108 Mon Sep 17 00:00:00 2001 From: Ashutosh Parkhi Date: Fri, 23 Dec 2022 12:49:23 +0000 Subject: [PATCH 1/2] [AOT] Added a test for output size post MLF export Follow up: https://github.com/apache/tvm/pull/12789 -Added a separate test to detect output size from MLF codegen -Updated test harness to detect correct size from IO arrays -MLF size was not used in aot.py because its unavailable in case of packed apis --- python/tvm/micro/model_library_format.py | 10 +++- tests/python/relay/aot/test_crt_aot.py | 58 ++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 2 deletions(-) diff --git a/python/tvm/micro/model_library_format.py b/python/tvm/micro/model_library_format.py index 5aa2d154ba57..f16d57983a11 100644 --- a/python/tvm/micro/model_library_format.py +++ b/python/tvm/micro/model_library_format.py @@ -485,6 +485,12 @@ def _export_graph_model_library_format( "functions" ]["main"][0]["outputs"][key] + input_name_to_size_map = { + name: property_map["size"] for name, property_map in inputs_sizes.items() + } + output_name_to_size_map = { + name: property_map["size"] for name, property_map in output_sizes.items() + } generate_c_interface_header( mod.libmod_name, inputs, @@ -494,8 +500,8 @@ def _export_graph_model_library_format( devices, workspace_size, include_path, - inputs_sizes, - output_sizes, + input_name_to_size_map, + output_name_to_size_map, ) is_aot = isinstance(mod, executor_factory.AOTExecutorFactoryModule) diff --git a/tests/python/relay/aot/test_crt_aot.py b/tests/python/relay/aot/test_crt_aot.py index c3426f147e0d..50d44b0ed384 100644 --- a/tests/python/relay/aot/test_crt_aot.py +++ b/tests/python/relay/aot/test_crt_aot.py @@ -225,6 +225,64 @@ def test_packed_global_variables(): assert f"{func}_packed" not in tvmgen_names +def test_io_size_definition(): + """Check network IO size definitions in the codegen output.""" + dtype = "float32" + ishape = (1, 32, 14, 14) + wshape = (32, 32, 3, 3) + interface_api = "c" + use_unpacked_api = True + + 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"), + ) + ref_output_size = output_list["output"].size * np.dtype(dtype).itemsize + 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" / "include").iterdir(): + if path.is_file(): + file_list.append(path) + assert len(file_list) > 0 + + for path in file_list: + with open(path, "r") as header: + contents = header.readlines() + contents = "".join(map(str, contents)) + assert contents.count("_SIZE") == 4 + assert str(ref_output_size) in contents + + @parametrize_aot_options def test_concatenate(interface_api, use_unpacked_api, test_runner): """Tests compilation of concatenate""" From 8115e68b61e8ac1aef79de2496375e821e30bef2 Mon Sep 17 00:00:00 2001 From: Ashutosh Parkhi Date: Tue, 3 Jan 2023 16:30:22 +0000 Subject: [PATCH 2/2] Introduced better checks in the test Change-Id: I07d470e4a8eabc75912675a56c528354c7d40981 --- tests/python/relay/aot/test_crt_aot.py | 29 +++++++++++++------------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/tests/python/relay/aot/test_crt_aot.py b/tests/python/relay/aot/test_crt_aot.py index 50d44b0ed384..b3db410156b3 100644 --- a/tests/python/relay/aot/test_crt_aot.py +++ b/tests/python/relay/aot/test_crt_aot.py @@ -242,9 +242,9 @@ def test_io_size_definition(): mod = transform.InferType()(mod) i_data = np.random.uniform(0, 1, ishape).astype(dtype) - w1_data = np.random.uniform(0, 1, wshape).astype(dtype) + w_data = np.random.uniform(0, 1, wshape).astype(dtype) - inputs = OrderedDict([("data", i_data), ("weight", w1_data)]) + inputs = OrderedDict([("data", i_data), ("weight", w_data)]) output_list = generate_ref_data(mod, inputs) compiled_models_list = compile_models( @@ -257,7 +257,10 @@ def test_io_size_definition(): use_runtime_executor=True, target=tvm.target.Target("c"), ) - ref_output_size = output_list["output"].size * np.dtype(dtype).itemsize + dtype_itemsize = np.dtype(dtype).itemsize + ref_input_size = i_data.size * dtype_itemsize + ref_weight_size = w_data.size * dtype_itemsize + ref_output_size = output_list["output"].size * dtype_itemsize compiled_model = compiled_models_list[0] tmp_path = utils.tempdir() @@ -269,18 +272,14 @@ def test_io_size_definition(): t = tarfile.open(tar_file) t.extractall(base_path) - file_list = [] - for path in (pathlib.Path(base_path) / "codegen" / "host" / "include").iterdir(): - if path.is_file(): - file_list.append(path) - assert len(file_list) > 0 - - for path in file_list: - with open(path, "r") as header: - contents = header.readlines() - contents = "".join(map(str, contents)) - assert contents.count("_SIZE") == 4 - assert str(ref_output_size) in contents + header_path = f"{base_path}/codegen/host/include/tvmgen_{model.name}.h" + with open(header_path, "r") as header: + contents = header.readlines() + contents = "".join(map(str, contents)) + assert contents.count("_SIZE") == 4 + assert f"TVMGEN_DEFAULT_DATA_SIZE {ref_input_size}" in contents + assert f"TVMGEN_DEFAULT_WEIGHT_SIZE {ref_weight_size}" in contents + assert f"TVMGEN_DEFAULT_OUTPUT_SIZE {ref_output_size}" in contents @parametrize_aot_options