[microTVM][ARM][Zephyr] Add CMSIS dependencies in Zephyr project build

apps/microtvm/zephyr/template_project/microtvm_api_server.py

@@ -27,15 +27,14 @@
 import pathlib
 import queue
 import re
-import select
 import shlex
 import shutil
 import subprocess
 import sys
 import tarfile
 import tempfile
 import threading
-import time
+from typing import Union
 import usb
 
 import serial
@@ -323,6 +322,12 @@ def _get_nrf_device_args(options):
         type="str",
         help="Extra definitions added project compile.",
     ),
+    server.ProjectOption(
+        "cmsis_path",
+        optional=["generate_project"],
+        type="str",
+        help="Path to the CMSIS directory.",
+    ),
 ]
 
 
@@ -333,6 +338,13 @@ def get_zephyr_base(options: dict):
     return zephyr_base
 
 
+def get_cmsis_path(options: dict) -> pathlib.Path:
+    """Returns CMSIS dependency path"""
+    cmsis_path = options.get("cmsis_path")
+    assert cmsis_path, "'cmsis_path' option not passed!"
+    return pathlib.Path(cmsis_path)
+
+
 class Handler(server.ProjectAPIHandler):
     def __init__(self):
         super(Handler, self).__init__()
@@ -424,6 +436,17 @@ def _get_platform_version(self, zephyr_base: str) -> float:
 
         return float(f"{version_major}.{version_minor}")
 
+    def _cmsis_required(self, project_path: Union[str, pathlib.Path]) -> bool:
+        """Check if CMSIS dependency is required."""
+        project_path = pathlib.Path(project_path)
+        for path in (project_path / "codegen" / "host" / "src").iterdir():
+            if path.is_file():
+                with open(path, "r") as lib_f:
+                    lib_content = lib_f.read()
+                if "<arm_nnsupportfunctions.h>" in lib_content and "<arm_math.h>" in lib_content:
+                    return True
+        return False
+
     def generate_project(self, model_library_format_path, standalone_crt_dir, project_dir, options):
         # Check Zephyr version
         version = self._get_platform_version(get_zephyr_base(options))
@@ -470,8 +493,8 @@ def generate_project(self, model_library_format_path, standalone_crt_dir, projec
                 shutil.copy2(src_path, dst_path)
 
         # Populate Makefile.
-        with open(API_SERVER_DIR / "CMakeLists.txt.template", "r") as cmake_template_f:
-            with open(project_dir / "CMakeLists.txt", "w") as cmake_f:
+        with open(project_dir / "CMakeLists.txt", "w") as cmake_f:
+            with open(API_SERVER_DIR / "CMakeLists.txt.template", "r") as cmake_template_f:
                 for line in cmake_template_f:
                     if self.API_SERVER_CRT_LIBS_TOKEN in line:
                         crt_libs = self.CRT_LIBS_BY_PROJECT_TYPE[options["project_type"]]
@@ -484,6 +507,20 @@ def generate_project(self, model_library_format_path, standalone_crt_dir, projec
                     for item in flags:
                         cmake_f.write(f"target_compile_definitions(app PUBLIC {item})\n")
 
+            # Include CMSIS libraries if required.
+            if self._cmsis_required(extract_path):
+                cmsis_path = get_cmsis_path(options)
+                cmake_f.write("\n")
+                cmake_f.write(
+                    f'target_include_directories(tvm_model PRIVATE {str(cmsis_path / "CMSIS" / "DSP" / "Include")})\n'
+                )
+                cmake_f.write(
+                    f'target_include_directories(tvm_model PRIVATE {str(cmsis_path / "CMSIS" / "DSP" / "Include" / "dsp")})\n'
+                )
+                cmake_f.write(
+                    f'target_include_directories(tvm_model PRIVATE {str(cmsis_path / "CMSIS" / "NN" / "Include")})\n'
+                )
+
         self._create_prj_conf(project_dir, options)
 
         # Populate crt-config.h

python/tvm/contrib/utils.py

@@ -93,11 +93,15 @@ def set_keep_for_debug(cls, set_to=True):
         finally:
             cls._KEEP_FOR_DEBUG = old_keep_for_debug
 
-    def __init__(self, custom_path=None):
+    def __init__(self, custom_path=None, keep_for_debug=None):
         if self.TEMPDIRS is None:
             raise DirectoryCreatedPastAtExit()
 
-        self._created_with_keep_for_debug = self._KEEP_FOR_DEBUG
+        if keep_for_debug is not None:
+            self._created_with_keep_for_debug = keep_for_debug
+        else:
+            self._created_with_keep_for_debug = self._KEEP_FOR_DEBUG
+
         if custom_path:
             os.mkdir(custom_path)
             self.temp_dir = custom_path
@@ -169,20 +173,22 @@ def listdir(self):
 atexit.register(TempDirectory.remove_tempdirs)
 
 
-def tempdir(custom_path=None):
+def tempdir(custom_path=None, keep_for_debug=None):
     """Create temp dir which deletes the contents when exit.
 
     Parameters
     ----------
     custom_path : str, optional
         Manually specify the exact temp dir path
 
+    keep_for_debug : bool
+        Keep temp directory for debugging purposes
     Returns
     -------
     temp : TempDirectory
         The temp directory object
     """
-    return TempDirectory(custom_path)
+    return TempDirectory(custom_path=custom_path, keep_for_debug=keep_for_debug)
 
 
 class FileLock(object):

tests/micro/zephyr/conftest.py

@@ -59,7 +59,7 @@ def tvm_debug(request):
 
 
 @pytest.fixture
-def temp_dir(board):
+def temp_dir(board, tvm_debug):
     parent_dir = pathlib.Path(os.path.dirname(__file__))
     filename = os.path.splitext(os.path.basename(__file__))[0]
     board_workspace = (
@@ -76,4 +76,21 @@ def temp_dir(board):
     if not os.path.exists(board_workspace.parent):
         os.makedirs(board_workspace.parent)
 
-    return tempdir(board_workspace)
+    keep_for_debug = tvm_debug if tvm_debug else None
+    test_temp_dir = tempdir(custom_path=board_workspace, keep_for_debug=keep_for_debug)
+    return test_temp_dir
+
+
+@pytest.fixture(autouse=True)
+def skip_by_board(request, board):
+    """Skip test if board is in the list."""
+    if request.node.get_closest_marker("skip_boards"):
+        if board in request.node.get_closest_marker("skip_boards").args[0]:
+            pytest.skip("skipped on this board: {}".format(board))
+
+
+def pytest_configure(config):
+    config.addinivalue_line(
+        "markers",
+        "skip_by_board(board): skip test for the given board",
+    )

tests/micro/zephyr/test_zephyr.py

@@ -22,6 +22,7 @@
 
 import pytest
 import numpy as np
+
 import onnx
 from PIL import Image
 
@@ -32,6 +33,7 @@
 from tvm.relay.testing import byoc
 from tvm.contrib import utils
 from tvm.micro.testing.utils import check_tune_log
+from tvm.target import arm_isa
 
 import test_utils
 
@@ -87,6 +89,7 @@ def _make_add_sess(temp_dir, model, zephyr_board, west_cmd, build_config, dtype=
 
 # The same test code can be executed on both the QEMU simulation and on real hardware.
 @tvm.testing.requires_micro
+@pytest.mark.skip_boards(["mps2_an521"])
 def test_add_uint(temp_dir, board, west_cmd, tvm_debug):
     """Test compiling the on-device runtime."""
 
@@ -112,6 +115,7 @@ def test_basic_add(sess):
 
 # The same test code can be executed on both the QEMU simulation and on real hardware.
 @tvm.testing.requires_micro
+@pytest.mark.skip_boards(["mps2_an521"])
 def test_add_float(temp_dir, board, west_cmd, tvm_debug):
     """Test compiling the on-device runtime."""
     model = test_utils.ZEPHYR_BOARDS[board]
@@ -138,6 +142,7 @@ def test_basic_add(sess):
 
 
 @tvm.testing.requires_micro
+@pytest.mark.skip_boards(["mps2_an521"])
 def test_platform_timer(temp_dir, board, west_cmd, tvm_debug):
     """Test compiling the on-device runtime."""
 
@@ -167,6 +172,7 @@ def test_basic_add(sess):
 
 
 @tvm.testing.requires_micro
+@pytest.mark.skip_boards(["mps2_an521"])
 def test_relay(temp_dir, board, west_cmd, tvm_debug):
     """Testing a simple relay graph"""
     model = test_utils.ZEPHYR_BOARDS[board]
@@ -199,6 +205,7 @@ def test_relay(temp_dir, board, west_cmd, tvm_debug):
 
 
 @tvm.testing.requires_micro
+@pytest.mark.skip_boards(["mps2_an521"])
 def test_onnx(temp_dir, board, west_cmd, tvm_debug):
     """Testing a simple ONNX model."""
     model = test_utils.ZEPHYR_BOARDS[board]
@@ -279,6 +286,7 @@ def check_result(
 
 
 @tvm.testing.requires_micro
+@pytest.mark.skip_boards(["mps2_an521"])
 def test_byoc_microtvm(temp_dir, board, west_cmd, tvm_debug):
     """This is a simple test case to check BYOC capabilities of microTVM"""
     model = test_utils.ZEPHYR_BOARDS[board]
@@ -359,6 +367,7 @@ def _make_add_sess_with_shape(temp_dir, model, zephyr_board, west_cmd, shape, bu
     ],
 )
 @tvm.testing.requires_micro
+@pytest.mark.skip_boards(["mps2_an521"])
 def test_rpc_large_array(temp_dir, board, west_cmd, tvm_debug, shape):
     """Test large RPC array transfer."""
     model = test_utils.ZEPHYR_BOARDS[board]
@@ -504,5 +513,66 @@ def test_autotune_conv2d(temp_dir, board, west_cmd, tvm_debug):
     tvm.testing.assert_allclose(output, expected_output, rtol=1e-4, atol=1e-5)
 
 
+@tvm.testing.requires_micro
+def test_schedule_build_with_cmsis_dependency(temp_dir, board, west_cmd, tvm_debug):
+    """Test Relay schedule with CMSIS dependency. This test shows if microTVM Auto tuning
+    with Zephyr breaks if CMSIS dependency was required for a schedule.
+    """
+    model = test_utils.ZEPHYR_BOARDS[board]
+    build_config = {"debug": tvm_debug}
+    target = tvm.target.target.micro(model, options=["-keys=arm_cpu,cpu"])
+
+    isa = arm_isa.IsaAnalyzer(target)
+    if not isa.has_dsp_support:
+        pytest.skip(f"ISA does not support DSP. target: {target}")
+
+    # Create a Relay conv2d
+    data_shape = (1, 16, 16, 3)
+    weight_shape = (5, 5, 8, 3)
+    data = relay.var("data", relay.TensorType(data_shape, "int8"))
+    weight = relay.var("weight", relay.TensorType(weight_shape, "int8"))
+    y = relay.nn.conv2d(
+        data,
+        weight,
+        padding=(2, 2),
+        kernel_size=(5, 5),
+        data_layout="NHWC",
+        kernel_layout="HWOI",
+        out_dtype="int32",
+    )
+    func = relay.Function([data, weight], y)
+    ir_mod = tvm.IRModule.from_expr(func)
+
+    runtime = Runtime("crt", {"system-lib": True})
+
+    with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
+        mod = tvm.relay.build(ir_mod, target=target, runtime=runtime)
+
+    project_options = {
+        "project_type": "host_driven",
+        "west_cmd": west_cmd,
+        "verbose": bool(build_config.get("debug")),
+        "zephyr_board": board,
+        "cmsis_path": os.getenv("CMSIS_PATH"),
+    }
+
+    project_dir = temp_dir / "project"
+    project = tvm.micro.generate_project(
+        str(test_utils.TEMPLATE_PROJECT_DIR),
+        mod,
+        project_dir,
+        project_options,
+    )
+    project.build()
+
+    with open(project_dir / "CMakeLists.txt", "r") as cmake_f:
+        cmake_content = cmake_f.read()
+
+    assert "CMSIS/DSP/Include" in cmake_content
+    assert "CMSIS/DSP/Include/dsp" in cmake_content
+    assert "CMSIS/DSP/Include" in cmake_content
+    assert "CMSIS/NN/Include" in cmake_content
+
+
 if __name__ == "__main__":
     tvm.testing.main()

tests/micro/zephyr/test_zephyr_aot.py

@@ -38,6 +38,7 @@
 
 
 @tvm.testing.requires_micro
+@pytest.mark.skip_boards(["mps2_an521"])
 def test_tflite(temp_dir, board, west_cmd, tvm_debug):
     """Testing a TFLite model."""
     model = test_utils.ZEPHYR_BOARDS[board]
@@ -93,6 +94,7 @@ def test_tflite(temp_dir, board, west_cmd, tvm_debug):
 
 
 @tvm.testing.requires_micro
+@pytest.mark.skip_boards(["mps2_an521"])
 def test_qemu_make_fail(temp_dir, board, west_cmd, tvm_debug):
     """Testing QEMU make fail."""
     if board not in ["qemu_x86", "mps2_an521", "mps3_an547"]:

tests/micro/zephyr/test_zephyr_armv7m.py

@@ -103,6 +103,7 @@ def _apply_desired_layout_no_simd(relay_mod):
 
 
 @tvm.testing.requires_micro
+@pytest.mark.skip_boards(["mps2_an521"])
 def test_armv7m_intrinsic(temp_dir, board, west_cmd, tvm_debug):
     """Testing a ARM v7m SIMD extension."""
 

tests/scripts/task_python_microtvm.sh

@@ -28,6 +28,7 @@ make cython3
 run_pytest ctypes python-microtvm-zephyr-qemu_x86 tests/micro/zephyr --zephyr-board=qemu_x86
 run_pytest ctypes python-microtvm-zephyr-qemu_riscv32 tests/micro/zephyr --zephyr-board=qemu_riscv32
 run_pytest ctypes python-microtvm-zephyr-qemu_riscv64 tests/micro/zephyr --zephyr-board=qemu_riscv64
+run_pytest ctypes python-microtvm-zephyr-mps2_an521 tests/micro/zephyr --zephyr-board=mps2_an521
 
 # Arduino
 run_pytest ctypes python-microtvm-arduino apps/microtvm/arduino/template_project/tests