You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by mo...@apache.org on 2023/01/09 17:48:10 UTC
[tvm] branch main updated: [AOT] Added a test for detecting output size post MLF export (#13655)
This is an automated email from the ASF dual-hosted git repository.
mousius pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 614e16d73a [AOT] Added a test for detecting output size post MLF export (#13655)
614e16d73a is described below
commit 614e16d73aef703ffd7edad774738d07a859020a
Author: Ashutosh Parkhi <86...@users.noreply.github.com>
AuthorDate: Mon Jan 9 17:48:03 2023 +0000
[AOT] Added a test for detecting output size post MLF export (#13655)
Follow up: https://github.com/apache/tvm/pull/12789
-Added a separate test to detect output size from MLF codegen
---
python/tvm/micro/model_library_format.py | 10 ++++--
tests/python/relay/aot/test_crt_aot.py | 57 ++++++++++++++++++++++++++++++++
2 files changed, 65 insertions(+), 2 deletions(-)
diff --git a/python/tvm/micro/model_library_format.py b/python/tvm/micro/model_library_format.py
index 5aa2d154ba..f16d57983a 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 c3426f147e..b3db410156 100644
--- a/tests/python/relay/aot/test_crt_aot.py
+++ b/tests/python/relay/aot/test_crt_aot.py
@@ -225,6 +225,63 @@ 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)
+ w_data = np.random.uniform(0, 1, wshape).astype(dtype)
+
+ inputs = OrderedDict([("data", i_data), ("weight", w_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"),
+ )
+ 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()
+ 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)
+
+ 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
def test_concatenate(interface_api, use_unpacked_api, test_runner):
"""Tests compilation of concatenate"""