[Fix][microTVM] QEMU RPC issue (#8021)

* add test

* fix test

* add parameter to test

* cleanup

* format

* address comments

* address comments

* direct read/write from/to ring buffer

* merge fix

* add comment

apps/microtvm/zephyr/host_driven/crt/crt_config.h

@@ -42,7 +42,7 @@
 #define TVM_CRT_MAX_REGISTERED_MODULES 2
 
 /*! Maximum packet size, in bytes, including the length header. */
-#define TVM_CRT_MAX_PACKET_SIZE_BYTES 8192
+#define TVM_CRT_MAX_PACKET_SIZE_BYTES (4 * 1024)
 
 /*! Maximum supported string length in dltype, e.g. "int8", "int16", "float32" */
 #define TVM_CRT_MAX_STRLEN_DLTYPE 10

apps/microtvm/zephyr/host_driven/src/main.c

@@ -61,6 +61,7 @@ static const struct device* led0_pin;
 
 static size_t g_num_bytes_requested = 0;
 static size_t g_num_bytes_written = 0;
+static size_t g_num_bytes_in_rx_buffer = 0;
 
 // Called by TVM to write serial data to the UART.
 ssize_t write_serial(void* unused_context, const uint8_t* data, size_t size) {
@@ -99,6 +100,7 @@ size_t TVMPlatformFormatMessage(char* out_buf, size_t out_buf_size_bytes, const
 
 // Called by TVM when an internal invariant is violated, and execution cannot continue.
 void TVMPlatformAbort(tvm_crt_error_t error) {
+  TVMLogf("TVMError: %x", error);
   sys_reboot(SYS_REBOOT_COLD);
 #ifdef CONFIG_LED
   gpio_pin_set(led0_pin, LED0_PIN, 1);
@@ -214,33 +216,37 @@ tvm_crt_error_t TVMPlatformTimerStop(double* elapsed_time_seconds) {
 }
 
 // Ring buffer used to store data read from the UART on rx interrupt.
-#define RING_BUF_SIZE_BYTES 4 * 1024
-RING_BUF_DECLARE(uart_rx_rbuf, RING_BUF_SIZE_BYTES);
-
-// Small buffer used to read data from the UART into the ring buffer.
-static uint8_t uart_data[8];
+// This ring buffer size is only required for testing with QEMU and not for physical hardware.
+#define RING_BUF_SIZE_BYTES (TVM_CRT_MAX_PACKET_SIZE_BYTES + 100)
+RING_BUF_ITEM_DECLARE_SIZE(uart_rx_rbuf, RING_BUF_SIZE_BYTES);
 
 // UART interrupt callback.
 void uart_irq_cb(const struct device* dev, void* user_data) {
-  while (uart_irq_update(dev) && uart_irq_is_pending(dev)) {
+  uart_irq_update(dev);
+  if (uart_irq_is_pending(dev)) {
     struct ring_buf* rbuf = (struct ring_buf*)user_data;
     if (uart_irq_rx_ready(dev) != 0) {
-      for (;;) {
-        // Read a small chunk of data from the UART.
-        int bytes_read = uart_fifo_read(dev, uart_data, sizeof(uart_data));
-        if (bytes_read < 0) {
-          TVMPlatformAbort((tvm_crt_error_t)0xbeef1);
-        } else if (bytes_read == 0) {
-          break;
-        }
-        // Write it into the ring buffer.
-        int bytes_written = ring_buf_put(rbuf, uart_data, bytes_read);
-        if (bytes_read != bytes_written) {
-          TVMPlatformAbort((tvm_crt_error_t)0xbeef2);
-        }
-        // CHECK_EQ(bytes_read, bytes_written, "bytes_read: %d; bytes_written: %d", bytes_read,
-        //         bytes_written);
+      uint8_t* data;
+      uint32_t size;
+      size = ring_buf_put_claim(rbuf, &data, RING_BUF_SIZE_BYTES);
+      int rx_size = uart_fifo_read(dev, data, size);
+      // Write it into the ring buffer.
+      g_num_bytes_in_rx_buffer += rx_size;
+
+      if (g_num_bytes_in_rx_buffer > RING_BUF_SIZE_BYTES) {
+        TVMPlatformAbort((tvm_crt_error_t)0xbeef3);
       }
+
+      if (rx_size < 0) {
+        TVMPlatformAbort((tvm_crt_error_t)0xbeef1);
+      }
+
+      int err = ring_buf_put_finish(rbuf, rx_size);
+      if (err != 0) {
+        TVMPlatformAbort((tvm_crt_error_t)0xbeef2);
+      }
+      // CHECK_EQ(bytes_read, bytes_written, "bytes_read: %d; bytes_written: %d", bytes_read,
+      // bytes_written);
     }
   }
 }
@@ -251,17 +257,6 @@ void uart_rx_init(struct ring_buf* rbuf, const struct device* dev) {
   uart_irq_rx_enable(dev);
 }
 
-// Used to read data from the UART.
-int uart_rx_buf_read(struct ring_buf* rbuf, uint8_t* data, size_t data_size_bytes) {
-  unsigned int key = irq_lock();
-  int bytes_read = ring_buf_get(rbuf, data, data_size_bytes);
-  irq_unlock(key);
-  return bytes_read;
-}
-
-// Buffer used to read from the UART rx ring buffer and feed it to the UTvmRpcServerLoop.
-static uint8_t main_rx_buf[RING_BUF_SIZE_BYTES];
-
 // The main function of this application.
 extern void __stdout_hook_install(int (*hook)(int));
 void main(void) {
@@ -299,13 +294,15 @@ void main(void) {
   // The main application loop. We continuously read commands from the UART
   // and dispatch them to UTvmRpcServerLoop().
   while (true) {
-    int bytes_read = uart_rx_buf_read(&uart_rx_rbuf, main_rx_buf, sizeof(main_rx_buf));
+    uint8_t* data;
+    unsigned int key = irq_lock();
+    uint32_t bytes_read = ring_buf_get_claim(&uart_rx_rbuf, &data, RING_BUF_SIZE_BYTES);
     if (bytes_read > 0) {
+      g_num_bytes_in_rx_buffer -= bytes_read;
       size_t bytes_remaining = bytes_read;
-      uint8_t* cursor = main_rx_buf;
       while (bytes_remaining > 0) {
         // Pass the received bytes to the RPC server.
-        tvm_crt_error_t err = UTvmRpcServerLoop(server, &cursor, &bytes_remaining);
+        tvm_crt_error_t err = UTvmRpcServerLoop(server, &data, &bytes_remaining);
         if (err != kTvmErrorNoError && err != kTvmErrorFramingShortPacket) {
           TVMPlatformAbort(err);
         }
@@ -317,7 +314,12 @@ void main(void) {
           g_num_bytes_requested = 0;
         }
       }
+      int err = ring_buf_get_finish(&uart_rx_rbuf, bytes_read);
+      if (err != 0) {
+        TVMPlatformAbort((tvm_crt_error_t)0xbeef6);
+      }
     }
+    irq_unlock(key);
   }
 
 #ifdef CONFIG_ARCH_POSIX

tests/micro/zephyr/conftest.py

@@ -44,6 +44,17 @@ def pytest_addoption(parser):
     parser.addoption(
         "--west-cmd", default="west", help="Path to `west` command for flashing device."
     )
+    parser.addoption(
+        "--skip-build",
+        action="store_true",
+        help="If set true, reuses build from the previous test run. Otherwise, build from the scratch.",
+    )
+    parser.addoption(
+        "--tvm-debug",
+        action="store_true",
+        default=False,
+        help="If set true, enable a debug session while the test is running. Before running the test, in a separate shell, you should run: <python -m tvm.exec.microtvm_debug_shell>",
+    )
 
 
 def pytest_generate_tests(metafunc):
@@ -54,3 +65,13 @@ def pytest_generate_tests(metafunc):
 @pytest.fixture
 def west_cmd(request):
     return request.config.getoption("--west-cmd")
+
+
+@pytest.fixture
+def skip_build(request):
+    return request.config.getoption("--skip-build")
+
+
+@pytest.fixture
+def tvm_debug(request):
+    return request.config.getoption("--tvm-debug")

tests/micro/zephyr/test_zephyr.py

@@ -42,30 +42,21 @@
 
 import conftest
 
-# If set, build the uTVM binary from scratch on each test.
-# Otherwise, reuses the build from the previous test run.
-BUILD = True
-
-# If set, enable a debug session while the test is running.
-# Before running the test, in a separate shell, you should run:
-#   python -m tvm.exec.microtvm_debug_shell
-DEBUG = False
-
 _LOG = logging.getLogger(__name__)
 
 PLATFORMS = conftest.PLATFORMS
 
 
-def _make_sess_from_op(model, zephyr_board, west_cmd, op_name, sched, arg_bufs):
+def _make_sess_from_op(model, zephyr_board, west_cmd, op_name, sched, arg_bufs, build_config):
     target = tvm.target.target.micro(model)
     target = tvm.target.Target(target=target, host=target)
     with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
         mod = tvm.build(sched, arg_bufs, target=target, name=op_name)
 
-    return _make_session(model, target, zephyr_board, west_cmd, mod)
+    return _make_session(model, target, zephyr_board, west_cmd, mod, build_config)
 
 
-def _make_session(model, target, zephyr_board, west_cmd, mod):
+def _make_session(model, target, zephyr_board, west_cmd, mod, build_config):
     parent_dir = os.path.dirname(__file__)
     filename = os.path.splitext(os.path.basename(__file__))[0]
     prev_build = f"{os.path.join(parent_dir, 'archive')}_{filename}_{zephyr_board}_last_build.micro"
@@ -94,14 +85,14 @@ def _make_session(model, target, zephyr_board, west_cmd, mod):
     opts["lib_opts"]["ccflags"] = ["-std=gnu++14"]
 
     flasher_kw = {}
-    if DEBUG:
+    if build_config["debug"]:
         flasher_kw["debug_rpc_session"] = tvm.rpc.connect("127.0.0.1", 9090)
 
     session_kw = {
         "flasher": compiler.flasher(**flasher_kw),
     }
 
-    if BUILD:
+    if not build_config["skip_build"]:
         session_kw["binary"] = tvm.micro.build_static_runtime(
             # the x86 compiler *expects* you to give the exact same dictionary for both
             # lib_opts and bin_opts. so the library compiler is mutating lib_opts and
@@ -124,19 +115,20 @@ def _make_session(model, target, zephyr_board, west_cmd, mod):
     return tvm.micro.Session(**session_kw)
 
 
-def _make_add_sess(model, zephyr_board, west_cmd):
+def _make_add_sess(model, zephyr_board, west_cmd, build_config):
     A = tvm.te.placeholder((2,), dtype="int8")
     B = tvm.te.placeholder((1,), dtype="int8")
     C = tvm.te.compute(A.shape, lambda i: A[i] + B[0], name="C")
     sched = tvm.te.create_schedule(C.op)
-    return _make_sess_from_op(model, zephyr_board, west_cmd, "add", sched, [A, B, C])
+    return _make_sess_from_op(model, zephyr_board, west_cmd, "add", sched, [A, B, C], build_config)
 
 
 # The same test code can be executed on both the QEMU simulation and on real hardware.
-def test_compile_runtime(platform, west_cmd):
+def test_compile_runtime(platform, west_cmd, skip_build, tvm_debug):
     """Test compiling the on-device runtime."""
 
     model, zephyr_board = PLATFORMS[platform]
+    build_config = {"skip_build": skip_build, "debug": tvm_debug}
 
     # NOTE: run test in a nested function so cPython will delete arrays before closing the session.
     def test_basic_add(sess):
@@ -151,14 +143,15 @@ def test_basic_add(sess):
         system_lib.get_function("add")(A_data, B_data, C_data)
         assert (C_data.numpy() == np.array([6, 7])).all()
 
-    with _make_add_sess(model, zephyr_board, west_cmd) as sess:
+    with _make_add_sess(model, zephyr_board, west_cmd, build_config) as sess:
         test_basic_add(sess)
 
 
-def test_platform_timer(platform, west_cmd):
+def test_platform_timer(platform, west_cmd, skip_build, tvm_debug):
     """Test compiling the on-device runtime."""
 
     model, zephyr_board = PLATFORMS[platform]
+    build_config = {"skip_build": skip_build, "debug": tvm_debug}
 
     # NOTE: run test in a nested function so cPython will delete arrays before closing the session.
     def test_basic_add(sess):
@@ -178,13 +171,14 @@ def test_basic_add(sess):
         assert result.mean > 0
         assert len(result.results) == 3
 
-    with _make_add_sess(model, zephyr_board, west_cmd) as sess:
+    with _make_add_sess(model, zephyr_board, west_cmd, build_config) as sess:
         test_basic_add(sess)
 
 
-def test_relay(platform, west_cmd):
+def test_relay(platform, west_cmd, skip_build, tvm_debug):
     """Testing a simple relay graph"""
     model, zephyr_board = PLATFORMS[platform]
+    build_config = {"skip_build": skip_build, "debug": tvm_debug}
     shape = (10,)
     dtype = "int8"
 
@@ -198,7 +192,7 @@ def test_relay(platform, west_cmd):
     with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
         graph, mod, params = tvm.relay.build(func, target=target)
 
-    with _make_session(model, target, zephyr_board, west_cmd, mod) as session:
+    with _make_session(model, target, zephyr_board, west_cmd, mod, build_config) as session:
         graph_mod = tvm.micro.create_local_graph_executor(
             graph, session.get_system_lib(), session.device
         )
@@ -210,9 +204,10 @@ def test_relay(platform, west_cmd):
         tvm.testing.assert_allclose(result, x_in * x_in + 1)
 
 
-def test_onnx(platform, west_cmd):
+def test_onnx(platform, west_cmd, skip_build, tvm_debug):
     """Testing a simple ONNX model."""
     model, zephyr_board = PLATFORMS[platform]
+    build_config = {"skip_build": skip_build, "debug": tvm_debug}
 
     # Load test images.
     this_dir = os.path.dirname(__file__)
@@ -239,7 +234,7 @@ def test_onnx(platform, west_cmd):
         lowered = relay.build(relay_mod, target, params=params)
         graph = lowered.get_graph_json()
 
-    with _make_session(model, target, zephyr_board, west_cmd, lowered.lib) as session:
+    with _make_session(model, target, zephyr_board, west_cmd, lowered.lib, build_config) as session:
         graph_mod = tvm.micro.create_local_graph_executor(
             graph, session.get_system_lib(), session.device
         )
@@ -311,14 +306,16 @@ def visit_call(self, call):
         return super().visit_call(call)
 
 
-def check_result(relay_mod, model, zephyr_board, west_cmd, map_inputs, out_shape, result):
+def check_result(
+    relay_mod, model, zephyr_board, west_cmd, map_inputs, out_shape, result, build_config
+):
     """Helper function to verify results"""
     TOL = 1e-5
     target = tvm.target.target.micro(model)
     with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
         graph, mod, params = tvm.relay.build(relay_mod, target=target)
 
-    with _make_session(model, target, zephyr_board, west_cmd, mod) as session:
+    with _make_session(model, target, zephyr_board, west_cmd, mod, build_config) as session:
         rt_mod = tvm.micro.create_local_graph_executor(
             graph, session.get_system_lib(), session.device
         )
@@ -337,9 +334,10 @@ def check_result(relay_mod, model, zephyr_board, west_cmd, map_inputs, out_shape
             tvm.testing.assert_allclose(out.numpy(), results[idx], rtol=TOL, atol=TOL)
 
 
-def test_byoc_utvm(platform, west_cmd):
+def test_byoc_utvm(platform, west_cmd, skip_build, tvm_debug):
     """This is a simple test case to check BYOC capabilities of uTVM"""
     model, zephyr_board = PLATFORMS[platform]
+    build_config = {"skip_build": skip_build, "debug": tvm_debug}
     x = relay.var("x", shape=(10, 10))
     w0 = relay.var("w0", shape=(10, 10))
     w1 = relay.var("w1", shape=(10, 10))
@@ -393,8 +391,42 @@ def test_byoc_utvm(platform, west_cmd):
         model=model,
         zephyr_board=zephyr_board,
         west_cmd=west_cmd,
+        build_config=build_config,
     )
 
 
+def _make_add_sess_with_shape(model, zephyr_board, west_cmd, shape, build_config):
+    A = tvm.te.placeholder(shape, dtype="int8")
+    C = tvm.te.compute(A.shape, lambda i: A[i] + A[i], name="C")
+    sched = tvm.te.create_schedule(C.op)
+    return _make_sess_from_op(model, zephyr_board, west_cmd, "add", sched, [A, C], build_config)
+
+
+@pytest.mark.parametrize(
+    "shape,",
+    [
+        pytest.param((1 * 1024,), id="(1*1024)"),
+        pytest.param((4 * 1024,), id="(4*1024)"),
+        pytest.param((16 * 1024,), id="(16*1024)"),
+    ],
+)
+def test_rpc_large_array(platform, west_cmd, skip_build, tvm_debug, shape):
+    """Test large RPC array transfer."""
+    model, zephyr_board = PLATFORMS[platform]
+    build_config = {"skip_build": skip_build, "debug": tvm_debug}
+
+    # NOTE: run test in a nested function so cPython will delete arrays before closing the session.
+    def test_tensors(sess):
+        a_np = np.random.randint(low=-128, high=127, size=shape, dtype="int8")
+
+        A_data = tvm.nd.array(a_np, device=sess.device)
+        assert (A_data.asnumpy() == a_np).all()
+        C_data = tvm.nd.array(np.zeros(shape, dtype="int8"), device=sess.device)
+        assert (C_data.asnumpy() == np.zeros(shape)).all()
+
+    with _make_add_sess_with_shape(model, zephyr_board, west_cmd, shape, build_config) as sess:
+        test_tensors(sess)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__] + sys.argv[1:]))