[microTVM] Fix AOT/ARMv7m tests on physical devices. (#9364)

* init

* test on hardware

* moved to testing.py

* add option to set workspace size

* changed model

* fix memory issue

* fix tests

* conv2d test

* fix simd test

* format

* fix exception

* refactor workspace helper function

* address comments

* revert to fix the bug

* address comment

apps/microtvm/zephyr/template_project/microtvm_api_server.py

@@ -275,6 +275,10 @@ def _get_nrf_device_args(options):
         choices=(True, False),
         help="Treat warnings as errors and raise an Exception.",
     ),
+    server.ProjectOption(
+        "compile_definitions",
+        help="Extra definitions added project compile.",
+    ),
 ]
 
 
@@ -419,6 +423,11 @@ def generate_project(self, model_library_format_path, standalone_crt_dir, projec
 
                     cmake_f.write(line)
 
+                if options.get("compile_definitions"):
+                    flags = options.get("compile_definitions")
+                    for item in flags:
+                        cmake_f.write(f"target_compile_definitions(app PUBLIC {item})\n")
+
         self._create_prj_conf(project_dir, options)
 
         # Populate crt-config.h

apps/microtvm/zephyr/template_project/src/aot_demo/main.c

@@ -38,14 +38,21 @@
 #include "posix_board_if.h"
 #endif
 
-#define WORKSPACE_SIZE (270 * 1024)
+// WORKSPACE_SIZE defined in Project API Makefile
 
 static uint8_t g_aot_memory[WORKSPACE_SIZE];
 tvm_workspace_t app_workspace;
 
-// Wakeup sequence used to wake up QEMU on the host.
-const unsigned char g_wakeup_sequence[] = "#wakeup\n";
-const char g_start_cmd[] = "start\n";
+// Transport Commands.
+// Commands on host end with `\n`
+// Commands on microTVM device end with `%`
+const unsigned char CMD_WAKEUP[] = "wakeup\n";
+const unsigned char CMD_READY[] = "ready\n";
+const unsigned char CMD_INIT[] = "init";
+const unsigned char CMD_INFER[] = "infer";
+
+#define CMD_SIZE 80u
+#define CMD_TERMINATOR '%'
 
 size_t TVMPlatformFormatMessage(char* out_buf, size_t out_buf_size_bytes, const char* fmt,
                                 va_list args) {
@@ -163,35 +170,10 @@ int TVMBackendFreeWorkspace(int device_type, int device_id, void* ptr) {
 }
 
 static uint8_t main_rx_buf[128];
-static uint8_t cmd_buf[128];
+static uint8_t g_cmd_buf[128];
 static size_t g_cmd_buf_ind;
 
-void main(void) {
-  g_cmd_buf_ind = 0;
-  memset((char*)cmd_buf, 0, sizeof(cmd_buf));
-  TVMPlatformUARTInit();
-  k_timer_init(&g_microtvm_timer, NULL, NULL);
-  // Wake up host side.
-  TVMPlatformWriteSerial(g_wakeup_sequence, sizeof(g_wakeup_sequence));
-
-  // Wait for start command
-  while (true) {
-    int bytes_read = TVMPlatformUartRxRead(main_rx_buf, sizeof(main_rx_buf));
-    if (bytes_read > 0) {
-      memcpy((char*)cmd_buf + g_cmd_buf_ind, main_rx_buf, bytes_read);
-      g_cmd_buf_ind += bytes_read;
-    }
-    if (g_cmd_buf_ind >= 6) {
-      if (!strcmp((char*)(cmd_buf), g_start_cmd)) {
-        break;
-      } else {
-        memset((char*)cmd_buf, 0, sizeof(cmd_buf));
-        g_cmd_buf_ind = 0;
-      }
-    }
-  }
-  TVMLogf("Zephyr AOT Runtime\n");
-
+void TVMInfer() {
   struct tvmgen_default_inputs inputs = {
       .input_1 = input_data,
   };
@@ -219,7 +201,47 @@ void main(void) {
       max_val = output_data[i];
     }
   }
-  TVMLogf("#result:%d:%d\n", max_ind, (uint32_t)(elapsed_time * 1000));
+  TVMLogf("result:%d:%d\n", max_ind, (uint32_t)(elapsed_time * 1000));
+}
+
+// Execute functions based on received command
+void command_ready(char* command) {
+  if (strncmp(command, CMD_INIT, CMD_SIZE) == 0) {
+    TVMPlatformWriteSerial(CMD_WAKEUP, sizeof(CMD_WAKEUP));
+  } else if (strncmp(command, CMD_INFER, CMD_SIZE) == 0) {
+    TVMInfer();
+  } else {
+    TVMPlatformWriteSerial(CMD_READY, sizeof(CMD_READY));
+  }
+}
+
+// Append received characters to buffer and check for termination character.
+void serial_callback(char* message, int len_bytes) {
+  for (int i = 0; i < len_bytes; i++) {
+    if (message[i] == CMD_TERMINATOR) {
+      g_cmd_buf[g_cmd_buf_ind] = (char)0;
+      command_ready(g_cmd_buf);
+      g_cmd_buf_ind = 0;
+    } else {
+      g_cmd_buf[g_cmd_buf_ind] = message[i];
+      g_cmd_buf_ind += 1;
+    }
+  }
+}
+
+void main(void) {
+  g_cmd_buf_ind = 0;
+  memset((char*)g_cmd_buf, 0, sizeof(g_cmd_buf));
+  TVMPlatformUARTInit();
+  k_timer_init(&g_microtvm_timer, NULL, NULL);
+
+  while (true) {
+    int bytes_read = TVMPlatformUartRxRead(main_rx_buf, sizeof(main_rx_buf));
+    if (bytes_read > 0) {
+      serial_callback(main_rx_buf, bytes_read);
+    }
+  }
+
 #ifdef CONFIG_ARCH_POSIX
   posix_exit(0);
 #endif

python/tvm/micro/project.py

@@ -184,7 +184,7 @@ def generate_project_from_mlf(
 
     mlf_path : pathlib.Path or str
         Path to the Model Library Format archive that will be used when creating
-        the new project.
+        the new project. The archive file will be copied to project_dir.
 
     options : dict
         Project API options given to the microTVM API server for the specified platform.

python/tvm/micro/testing.py

@@ -19,8 +19,16 @@
 
 import pathlib
 import json
+import logging
+import tarfile
+import time
 from typing import Union
 
+from tvm.micro.project_api.server import IoTimeoutError
+
+# Timeout in seconds for AOT transport.
+TIMEOUT_SEC = 10
+
 
 def check_tune_log(log_path: Union[pathlib.Path, str]):
     """Read the tuning log and check each result."""
@@ -31,3 +39,47 @@ def check_tune_log(log_path: Union[pathlib.Path, str]):
         if len(line) > 0:
             tune_result = json.loads(line)
             assert tune_result["result"][0][0] < 1000000000.0
+
+
+def aot_transport_init_wait(transport):
+    """Send init message to microTVM device until it receives wakeup sequence."""
+    while True:
+        try:
+            aot_transport_find_message(transport, "wakeup", timeout_sec=TIMEOUT_SEC)
+            break
+        except IoTimeoutError:
+            transport.write(b"init%", timeout_sec=TIMEOUT_SEC)
+
+
+def aot_transport_find_message(transport, expression: str, timeout_sec: int) -> str:
+    """Read transport message until it finds the expression."""
+    timeout = timeout_sec
+    start_time = time.monotonic()
+    while True:
+        data = _read_line(transport, timeout)
+        logging.debug("new line: %s", data)
+        if expression in data:
+            return data
+        timeout = max(0, timeout_sec - (time.monotonic() - start_time))
+
+
+def _read_line(transport, timeout_sec: int) -> str:
+    data = bytearray()
+    while True:
+        new_data = transport.read(1, timeout_sec=timeout_sec)
+        logging.debug("read data: %s", new_data)
+        for item in new_data:
+            data.append(item)
+            if str(chr(item)) == "\n":
+                return data.decode(encoding="utf-8")
+
+
+def mlf_extract_workspace_size_bytes(mlf_tar_path: Union[pathlib.Path, str]) -> int:
+    """Extract an MLF archive file and read workspace size from metadata file."""
+
+    with tarfile.open(mlf_tar_path, "r:*") as tar_file:
+        tar_members = [ti.name for ti in tar_file.getmembers()]
+        assert "./metadata.json" in tar_members
+        with tar_file.extractfile("./metadata.json") as f:
+            metadata = json.load(f)
+            return metadata["memory"]["functions"]["main"][0]["workspace_size_bytes"]

tests/micro/zephyr/test_utils.py

@@ -18,8 +18,9 @@
 import os
 import json
 import pathlib
-import logging
 import tarfile
+import tempfile
+from typing import Union
 
 import numpy as np
 
@@ -29,7 +30,8 @@
 import requests
 
 import tvm.micro
-
+from tvm.micro import export_model_library_format
+from tvm.micro.testing import mlf_extract_workspace_size_bytes
 
 TEMPLATE_PROJECT_DIR = (
     pathlib.Path(__file__).parent
@@ -77,19 +79,29 @@ def has_fpu(board: str):
 
 def build_project(temp_dir, zephyr_board, west_cmd, mod, build_config, extra_files_tar=None):
     project_dir = temp_dir / "project"
-    project = tvm.micro.generate_project(
-        str(TEMPLATE_PROJECT_DIR),
-        mod,
-        project_dir,
-        {
-            "extra_files_tar": extra_files_tar,
-            "project_type": "aot_demo",
-            "west_cmd": west_cmd,
-            "verbose": bool(build_config.get("debug")),
-            "zephyr_board": zephyr_board,
-        },
-    )
-    project.build()
+
+    with tempfile.TemporaryDirectory() as tar_temp_dir:
+        model_tar_path = pathlib.Path(tar_temp_dir) / "model.tar"
+        export_model_library_format(mod, model_tar_path)
+
+        workspace_size = mlf_extract_workspace_size_bytes(model_tar_path)
+        project = tvm.micro.project.generate_project_from_mlf(
+            str(TEMPLATE_PROJECT_DIR),
+            project_dir,
+            model_tar_path,
+            {
+                "extra_files_tar": extra_files_tar,
+                "project_type": "aot_demo",
+                "west_cmd": west_cmd,
+                "verbose": bool(build_config.get("debug")),
+                "zephyr_board": zephyr_board,
+                "compile_definitions": [
+                    # TODO(mehrdadh): It fails without offset.
+                    f"-DWORKSPACE_SIZE={workspace_size + 128}",
+                ],
+            },
+        )
+        project.build()
     return project, project_dir
 
 
@@ -129,31 +141,6 @@ def create_header_file(tensor_name, npy_data, output_path, tar_file):
     tar_file.addfile(ti, io.BytesIO(header_file_bytes))
 
 
-def _read_line(fd, timeout_sec: int):
-    data = ""
-    new_line = False
-    while True:
-        if new_line:
-            break
-        new_data = fd.read(1, timeout_sec=timeout_sec)
-        logging.debug(f"read data: {new_data}")
-        for item in new_data:
-            new_c = chr(item)
-            data = data + new_c
-            if new_c == "\n":
-                new_line = True
-                break
-    return data
-
-
-def get_message(fd, expr: str, timeout_sec: int):
-    while True:
-        data = _read_line(fd, timeout_sec)
-        logging.debug(f"new line: {data}")
-        if expr in data:
-            return data
-
-
 # TODO move CMSIS integration to microtvm_api_server.py
 # see https://discuss.tvm.apache.org/t/tvm-capturing-dependent-libraries-of-code-generated-tir-initially-for-use-in-model-library-format/11080
 def loadCMSIS(temp_dir):

tests/micro/zephyr/test_zephyr.py

@@ -374,8 +374,8 @@ def test_tensors(sess):
 @tvm.testing.requires_micro
 def test_autotune_conv2d(temp_dir, board, west_cmd, tvm_debug):
     """Test AutoTune for microTVM Zephyr"""
-    if board in ["qemu_riscv32", "qemu_riscv64"]:
-        pytest.xfail(f"Autotune fails on  {board}.")
+    if board != "qemu_x86":
+        pytest.xfail(f"Autotune fails on {board}.")
 
     model = test_utils.ZEPHYR_BOARDS[board]
     build_config = {"debug": tvm_debug}