[CRT][microTVM] Enable USMP by default for AoTExecutor + CRT runtime

apps/microtvm/zephyr/template_project/src/mlperftiny/platform.cc

@@ -38,10 +38,6 @@
 
 #include "crt_config.h"
 
-// TVM_WORKSPACE_SIZE_BYTES is defined in python
-static uint8_t g_aot_memory[TVM_WORKSPACE_SIZE_BYTES];
-tvm_workspace_t app_workspace;
-
 size_t TVMPlatformFormatMessage(char* out_buf, size_t out_buf_size_bytes, const char* fmt,
                                 va_list args) {
   return vsnprintk(out_buf, out_buf_size_bytes, fmt, args);
@@ -53,16 +49,3 @@ void TVMPlatformAbort(tvm_crt_error_t error) {
   for (;;)
     ;
 }
-
-tvm_crt_error_t TVMPlatformMemoryAllocate(size_t num_bytes, DLDevice dev, void** out_ptr) {
-  return StackMemoryManager_Allocate(&app_workspace, num_bytes, out_ptr);
-}
-
-tvm_crt_error_t TVMPlatformMemoryFree(void* ptr, DLDevice dev) {
-  return StackMemoryManager_Free(&app_workspace, ptr);
-}
-
-tvm_crt_error_t TVMPlatformInitialize() {
-  StackMemoryManager_Init(&app_workspace, g_aot_memory, sizeof(g_aot_memory));
-  return kTvmErrorNoError;
-}

apps/microtvm/zephyr/template_project/src/mlperftiny/submitter_implemented.cc

@@ -259,7 +259,7 @@ void th_infer() { Infer(g_input_data); }
 
 /// \brief optional API.
 // Modified from source
-void th_final_initialize(void) { TVMPlatformInitialize(); }
+void th_final_initialize(void) {}
 
 void th_pre() {}
 void th_post() {}

src/relay/backend/aot_executor_codegen.cc

@@ -1209,7 +1209,15 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     // Parallel for loops are not supported in AoT codegen.
     lowered_mod = tir::transform::ConvertForLoopsToSerial()(lowered_mod);
 
-    bool enable_usmp = pass_ctx->GetConfig<Bool>(kUSMPEnableOption, Bool(false)).value();
+    // Check USMP option
+    bool enable_usmp = false;
+    if (runtime_config->name == kTvmRuntimeCrt) {
+      enable_usmp = true;
+    }
+    if (pass_ctx->GetConfig<Bool>(kUSMPEnableOption) != nullptr) {
+      enable_usmp = pass_ctx->GetConfig<Bool>(kUSMPEnableOption, Bool(false)).value();
+    }
+
     if (enable_usmp) {
       lowered_mod = PlanMemoryWithUSMP(lowered_mod);
     } else {

tests/micro/arduino/test_arduino_rpc_server.py

@@ -23,7 +23,6 @@
 """
 
 import pathlib
-import sys
 import numpy as np
 import onnx
 import pytest

tests/micro/arduino/test_arduino_workflow.py

@@ -18,7 +18,6 @@
 import pathlib
 import re
 import shutil
-import sys
 import pytest
 
 import tvm.testing

tests/micro/zephyr/test_ms_tuning.py

@@ -24,7 +24,7 @@
 from tvm import relay
 import tvm.micro.testing
 from tvm.relay.backend import Executor
-from tvm.contrib import graph_executor, utils
+from tvm.contrib import graph_executor
 from tvm import meta_schedule as ms
 from tvm.contrib.micro.meta_schedule.local_builder_micro import get_local_builder_micro
 from tvm.contrib.micro.meta_schedule.rpc_runner_micro import get_rpc_runner_micro

tests/micro/zephyr/test_zephyr_aot_exec.py

@@ -22,7 +22,6 @@
 import tvm.micro.testing
 import tvm.relay as relay
 from tvm.relay.backend import Executor, Runtime
-from tvm.contrib import utils
 
 from . import utils
 

tests/micro/zephyr/test_zephyr_aot_exec_standalone.py

@@ -14,9 +14,6 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-import os
-import pathlib
-
 import pytest
 import numpy as np
 
@@ -62,9 +59,7 @@ def test_tflite(workspace_dir, board, microtvm_debug, serial_number):
         "aot", {"unpacked-api": True, "interface-api": "c", "workspace-byte-alignment": 4}
     )
     runtime = Runtime("crt")
-    with tvm.transform.PassContext(
-        opt_level=3, config={"tir.disable_vectorize": True, "tir.usmp.enable": True}
-    ):
+    with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
         lowered = relay.build(relay_mod, target, params=params, runtime=runtime, executor=executor)
 
     sample_url = "https://github.com/tlc-pack/web-data/raw/main/testdata/microTVM/data/keyword_spotting_int8_6.pyc.npy"

tests/python/contrib/test_cmsisnn/test_conv2d.py

@@ -210,6 +210,7 @@ def test_conv2d_number_primfunc_args(
         AOTTestModel(module=cmsisnn_mod, inputs=inputs, outputs=output_list, params=params),
         interface_api,
         use_unpacked_api,
+        pass_config={"tir.usmp.enable": False},
     )
 
     # validate number of TIR primfunc args

tests/python/relay/aot/test_crt_aot.py

@@ -802,6 +802,7 @@ def @main(%data: Tensor[(1, 4, 4, 4), float32], %weight: Tensor[(4, 4, 3, 3), fl
         models=AOTTestModel(module=relay_mod, inputs=None, outputs=None),
         interface_api="c",
         use_unpacked_api=True,
+        pass_config={"tir.usmp.enable": False},
     )
     source = compiled_test_mods[0].executor_factory.lib.imported_modules[0].get_source()
     # There should be three allocates created for three primitive relay function
@@ -827,6 +828,7 @@ def test_constants_alignment(constants_byte_alignment):
         interface_api,
         use_unpacked_api,
         target=tvm.target.Target(target, host=target),
+        pass_config={"tir.usmp.enable": False},
     )
     source = compiled_test_mods[0].executor_factory.lib.imported_modules[0].get_source()
     assert f'__attribute__((section(".rodata.tvm"), aligned({constants_byte_alignment})))' in source
@@ -966,6 +968,7 @@ def test_workspace_calculation(workspace_byte_alignment, main_workspace_size):
         opt_level=3,
         config={
             "tir.disable_vectorize": True,
+            "tir.usmp.enable": False,
         },
     ):
         lib = tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)

tests/python/relay/aot/test_crt_aot_usmp.py

@@ -907,5 +907,40 @@ def test_incompatible_interface_api_errors():
             tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
 
 
+@parametrize_aot_options
+def test_usmp_enabled_by_default_for_crt(interface_api, use_unpacked_api, test_runner):
+    """This test checks whether USMP is enabled by default
+    for cortex-M targets.
+    """
+    dtype = "float32"
+    ishape = (1, 32, 14, 14)
+    wshape = (32, 32, 3, 3)
+
+    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_test_mods = compile_models(
+        models=AOTTestModel(module=mod, inputs=inputs, outputs=output_list),
+        interface_api=interface_api,
+        use_unpacked_api=use_unpacked_api,
+        pass_config=test_runner.pass_config,
+        target=tvm.target.target.micro("host"),
+    )
+
+    for compiled_model in compiled_test_mods:
+        _check_for_no_tvm_backendallocworkspace_calls(compiled_model.executor_factory.lib)
+
+
 if __name__ == "__main__":
     tvm.testing.main()

tests/python/unittest/test_micro_model_library_format.py

@@ -159,7 +159,9 @@ def test_export_model_library_format_c(
 ):
     target = tvm.target.target.micro("host")
     with utils.TempDirectory.set_keep_for_debug(True):
-        with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
+        with tvm.transform.PassContext(
+            opt_level=3, config={"tir.disable_vectorize": True, "tir.usmp.enable": False}
+        ):
             relay_mod = tvm.relay.fromtext(
                 """
             #[version = "0.0.5"]
@@ -338,7 +340,9 @@ def @main(%a : Tensor[(1, 2), uint8], %b : Tensor[(1, 2), float32], %c : Tensor[
 )
 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}):
+    with tvm.transform.PassContext(
+        opt_level=3, config={"tir.disable_vectorize": True, "tir.usmp.enable": False}
+    ):
         relay_mod = tvm.relay.fromtext(
             """
             #[version = "0.0.5"]