From ac179068094186124d9c15603da4bb57e9100889 Mon Sep 17 00:00:00 2001 From: Alan MacDonald Date: Wed, 18 Jan 2023 10:26:57 -0800 Subject: [PATCH 1/5] change host-driven AOT to call the _run() function instead of __tvm_main__() to properly initialize memory pools --- src/runtime/crt/aot_executor/aot_executor.c | 13 +------------ src/target/source/source_module.cc | 4 ++++ 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/src/runtime/crt/aot_executor/aot_executor.c b/src/runtime/crt/aot_executor/aot_executor.c index a40c1d530fa9..ae007037e6cc 100644 --- a/src/runtime/crt/aot_executor/aot_executor.c +++ b/src/runtime/crt/aot_executor/aot_executor.c @@ -83,7 +83,7 @@ int TVMAotExecutor_GetInputIndex(TVMAotExecutor* executor, const char* name) { } int TVMAotExecutor_Run(TVMAotExecutor* executor) { - const char* tvm_main_suffix = "___tvm_main__"; + const char* tvm_main_suffix = "_run"; char tvm_main_name[TVM_CRT_MAX_STRLEN_FUNCTION_NAME]; { @@ -203,17 +203,6 @@ int TVMAotExecutor_Init(TVMAotExecutor* executor, TVMModuleHandle module_handle, TVMNDArray_IncrementReference(array); } - for (i = 0; i < md->num_workspace_pools; ++i) { - LOG_DEBUG("pools allocate[%d]: %s\n", i, md->workspace_pools[i].name); - - status = TVMNDArray_Empty(md->workspace_pools[i].num_shape, md->workspace_pools[i].shape, - md->workspace_pools[i].dtype, executor->device, - &executor->args[arg_idx++]); - if (status != 0) { - return status; - } - } - CHECK_EQ(0, md->num_constant_pools, "Constant pools not supported"); return status; } diff --git a/src/target/source/source_module.cc b/src/target/source/source_module.cc index ccc15fc1ee49..319727ce1956 100644 --- a/src/target/source/source_module.cc +++ b/src/target/source/source_module.cc @@ -944,6 +944,10 @@ runtime::Module CreateCSourceCrtMetadataModule(const Array& mod } } + // add the run function (typically "tvmgen_default_run") to function registry + std::string run_func = runtime::get_name_mangled(metadata->mod_name, "run"); + func_names.push_back(run_func); + auto n = make_object(func_names, "c", target, runtime, metadata); auto csrc_metadata_module = runtime::Module(n); for (const auto& mod : final_modules) { From 3a515a741578d308edc2cae442a4fd9c20b5bafc Mon Sep 17 00:00:00 2001 From: Alan MacDonald Date: Wed, 18 Jan 2023 14:48:40 -0800 Subject: [PATCH 2/5] only add run function when AOT executor is specified --- src/target/source/source_module.cc | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/target/source/source_module.cc b/src/target/source/source_module.cc index 319727ce1956..8b3ec8e50612 100644 --- a/src/target/source/source_module.cc +++ b/src/target/source/source_module.cc @@ -944,9 +944,14 @@ runtime::Module CreateCSourceCrtMetadataModule(const Array& mod } } - // add the run function (typically "tvmgen_default_run") to function registry - std::string run_func = runtime::get_name_mangled(metadata->mod_name, "run"); - func_names.push_back(run_func); + if (metadata.defined()) { + if (metadata->executor == "aot") { + // add the run function (typically "tvmgen_default_run") to function registry + // when using AOT executor + std::string run_func = runtime::get_name_mangled(metadata->mod_name, "run"); + func_names.push_back(run_func); + } + } auto n = make_object(func_names, "c", target, runtime, metadata); auto csrc_metadata_module = runtime::Module(n); From 57fd0d39d287bf5810002d1c35a80220c6545175 Mon Sep 17 00:00:00 2001 From: Alan MacDonald Date: Wed, 18 Jan 2023 16:38:37 -0800 Subject: [PATCH 3/5] only generate get_c_metadata function when executor is AOT --- src/target/source/source_module.cc | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/src/target/source/source_module.cc b/src/target/source/source_module.cc index 8b3ec8e50612..ee5a7cd33de9 100644 --- a/src/target/source/source_module.cc +++ b/src/target/source/source_module.cc @@ -929,11 +929,21 @@ runtime::Module CreateCSourceCrtMetadataModule(const Array& mod relay::backend::ExecutorCodegenMetadata metadata, runtime::metadata::Metadata aot_metadata) { Array final_modules(modules); - if (aot_metadata.defined()) { - final_modules.push_back(CreateAotMetadataModule(aot_metadata, true)); + Array func_names; + + if (metadata.defined()) { + if (metadata->executor == "aot") { + if (aot_metadata.defined()) { + final_modules.push_back(CreateAotMetadataModule(aot_metadata, true)); + } + + // add the run function (typically "tvmgen_default_run") to function registry + // when using AOT executor + std::string run_func = runtime::get_name_mangled(metadata->mod_name, "run"); + func_names.push_back(run_func); + } } - Array func_names; for (runtime::Module mod : final_modules) { auto pf_funcs = mod.GetFunction("get_func_names"); if (pf_funcs != nullptr) { @@ -944,15 +954,6 @@ runtime::Module CreateCSourceCrtMetadataModule(const Array& mod } } - if (metadata.defined()) { - if (metadata->executor == "aot") { - // add the run function (typically "tvmgen_default_run") to function registry - // when using AOT executor - std::string run_func = runtime::get_name_mangled(metadata->mod_name, "run"); - func_names.push_back(run_func); - } - } - auto n = make_object(func_names, "c", target, runtime, metadata); auto csrc_metadata_module = runtime::Module(n); for (const auto& mod : final_modules) { From 40562d23467d2a4f826402ec60bdc7626bd8d9d4 Mon Sep 17 00:00:00 2001 From: Alan MacDonald Date: Thu, 19 Jan 2023 17:07:30 -0800 Subject: [PATCH 4/5] remove check for mod1_lib2.c since the file with the get_c_metadata is no longer generated --- tests/python/unittest/test_micro_model_library_format.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/python/unittest/test_micro_model_library_format.py b/tests/python/unittest/test_micro_model_library_format.py index e664c2ebb858..39919f337197 100644 --- a/tests/python/unittest/test_micro_model_library_format.py +++ b/tests/python/unittest/test_micro_model_library_format.py @@ -618,7 +618,6 @@ def test_multiple_relay_modules_aot_graph(): 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")) From 35cc3f0449fb17b16e67fb7c0550a9ac5922a692 Mon Sep 17 00:00:00 2001 From: Alan MacDonald Date: Mon, 23 Jan 2023 15:01:45 -0800 Subject: [PATCH 5/5] update USMP const pool test to always enable USMP and expect to succeed --- tests/python/unittest/test_crt.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/tests/python/unittest/test_crt.py b/tests/python/unittest/test_crt.py index 83fa91af06c9..3309aad0a5db 100644 --- a/tests/python/unittest/test_crt.py +++ b/tests/python/unittest/test_crt.py @@ -229,15 +229,10 @@ def do_test(): do_test() -enable_usmp, expect_exception = tvm.testing.parameters((True, True), (False, False)) - - @tvm.testing.requires_micro -def test_aot_executor_usmp_const_pool(enable_usmp, expect_exception): - """Test the AOT executor with microTVM using usmp. - Test should fail if const pool is supplied to executor - as these are currently not supported - """ +def test_aot_executor_usmp_const_pool(): + """Test the AOT executor with microTVM using USMP to generate a constant data pool.""" + ws_root = pathlib.Path(os.path.dirname(__file__) + "/micro-workspace-usmp") if ws_root.exists(): shutil.rmtree(ws_root) @@ -260,7 +255,7 @@ def @main(%a : Tensor[(1, 2), uint8], %b : Tensor[(1, 2), uint8], %c : Tensor[(1 C_np = np.array([[8, 9]], dtype="uint8").astype(type_dict["c"]) params = {"c": C_np} with tvm.transform.PassContext( - opt_level=3, config={"tir.disable_vectorize": True, "tir.usmp.enable": enable_usmp} + opt_level=3, config={"tir.disable_vectorize": True, "tir.usmp.enable": True} ): factory = tvm.relay.build( relay_mod, @@ -278,10 +273,7 @@ def do_test(): ) ) except tvm._ffi.base.TVMError as e: - if expect_exception: - return - else: - raise e + raise e assert aot_executor.get_input_index("a") == 0 assert aot_executor.get_input_index("b") == 1