Update uTVM code to work with the nRF5340DK dev board. (apache#7331)

* Various fixes to get nRF5340 working. Not yet there.

* nRF5340 test runs locally.

* Various fixes to get nRF5340 working. Not yet there.

* nRF5340 test runs locally.

* Add `nrfjprog --recover` for nRF5340DK

* Cleanup.

* Remove debugging code.

* Revert submodule update.

* Remove debugging code.

* Fix comment.

* Remove -keys argument.

* Adding some debugging code

* Fix passing west command to ZephyrFlasher.

* Various fixes to get nRF5340 working. Not yet there.

* nRF5340 test runs locally.

* Add `nrfjprog --recover` for nRF5340DK

* Cleanup.

* Various fixes to get nRF5340 working. Not yet there.

* nRF5340 test runs locally.

* Remove debugging code.

* Fix comment.

* Remove -keys argument.

* Fix merge.

apps/microtvm/reference-vm/zephyr/pyproject.toml

@@ -64,6 +64,9 @@ scipy = "^1.4"
 python = "^3.6"
 tornado = "^6"
 typed_ast = "^1.4"
+pyyaml = "^5.4.1"
+pyserial = "^3.5"
+
 
 # AutoTVM
 xgboost = {version = "^1.1", optional = true}

python/tvm/micro/contrib/zephyr.py

@@ -191,7 +191,7 @@ def library(self, output, sources, options=None):
         with open(os.path.join(output, "main.c"), "w"):
             pass
 
-        # expecetd not to exist after populate_tvm_libs
+        # expected not to exist after populate_tvm_libs
         build_dir = os.path.join(output, "__tvm_build")
         os.mkdir(build_dir)
         self._subprocess_env.run(
@@ -241,11 +241,12 @@ def binary(self, output, objects, options=None, link_main=True, main_options=Non
     def flasher_factory(self):
         return compiler.FlasherFactory(
             ZephyrFlasher,
-            (self._west_cmd,),
+            (self._board,),
             dict(
                 zephyr_base=self._zephyr_base,
                 project_dir=self._project_dir,
                 subprocess_env=self._subprocess_env.default_overrides,
+                west_cmd=self._west_cmd,
             ),
         )
 
@@ -291,7 +292,7 @@ class ZephyrFlasher(tvm.micro.compiler.Flasher):
 
     def __init__(
         self,
-        west_cmd,
+        board,
         zephyr_base=None,
         project_dir=None,
         subprocess_env=None,
@@ -300,6 +301,7 @@ def __init__(
         flash_args=None,
         debug_rpc_session=None,
         serial_timeouts=None,
+        west_cmd=None,
     ):
         zephyr_base = zephyr_base or os.environ["ZEPHYR_BASE"]
         sys.path.insert(0, os.path.join(zephyr_base, "scripts", "dts"))
@@ -310,6 +312,7 @@ def __init__(
         finally:
             sys.path.pop(0)
 
+        self._board = board
         self._zephyr_base = zephyr_base
         self._project_dir = project_dir
         self._west_cmd = west_cmd
@@ -414,6 +417,20 @@ def flash(self, micro_binary):
         build_dir = os.path.dirname(
             micro_binary.abspath(micro_binary.labelled_files["cmake_cache"][0])
         )
+
+        # The nRF5340DK requires an additional `nrfjprog --recover` before each flash cycle.
+        # This is because readback protection is enabled by default when this device is flashed.
+        # Otherwise, flashing may fail with an error such as the following:
+        #  ERROR: The operation attempted is unavailable due to readback protection in
+        #  ERROR: your device. Please use --recover to unlock the device.
+        if (
+            self._board.startswith("nrf5340dk")
+            and self._get_flash_runner(cmake_entries) == "nrfjprog"
+        ):
+            recover_args = ["nrfjprog", "--recover"]
+            recover_args.extend(self._get_nrf_device_args())
+            self._subprocess_env.run(recover_args, cwd=build_dir)
+
         west_args = (
             self._west_cmd
             + ["flash", "--build-dir", build_dir, "--skip-rebuild"]

python/tvm/target/target.py

@@ -234,7 +234,10 @@ def micro(model="unknown", options=None):
     trans_table = {
         "host": [],
         "stm32f746xx": ["-mcpu=cortex-m7", "-march=armv7e-m"],
+        "nrf5340dk": ["-mcpu=cortex-m33"],
     }
+    if model not in trans_table:
+        raise ValueError(f"Model {model} not supported by tvm.target.micro.")
     opts = _merge_opts(
         trans_table[model] + ["-runtime=c", "--system-lib", f"-model={model}"],
         options,

tests/micro/qemu/conftest.py

@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import pytest
 
 
 def pytest_addoption(parser):
@@ -25,8 +26,16 @@ def pytest_addoption(parser):
             "for microTVM tests."
         ),
     )
+    parser.addoption(
+        "--west-cmd", default="west", help="Path to `west` command for flashing device."
+    )
 
 
 def pytest_generate_tests(metafunc):
     if "platform" in metafunc.fixturenames:
         metafunc.parametrize("platform", metafunc.config.getoption("microtvm_platforms").split(","))
+
+
+@pytest.fixture
+def west_cmd(request):
+    return request.config.getoption("--west-cmd")

tests/micro/qemu/test_zephyr.py

@@ -43,15 +43,15 @@
 TARGET = None
 
 
-def _make_sess_from_op(model, zephyr_board, op_name, sched, arg_bufs):
+def _make_sess_from_op(model, zephyr_board, west_cmd, op_name, sched, arg_bufs):
     target = tvm.target.target.micro(model)
     with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
         mod = tvm.build(sched, arg_bufs, target, target_host=target, name=op_name)
 
-    return _make_session(model, target, zephyr_board, mod)
+    return _make_session(model, target, zephyr_board, west_cmd, mod)
 
 
-def _make_session(model, target, zephyr_board, mod):
+def _make_session(model, target, zephyr_board, west_cmd, mod):
     test_name = f"{os.path.splitext(os.path.abspath(__file__))[0]}-{model}"
     prev_build = f"{test_name}-last-build.micro-binary"
     workspace_root = (
@@ -65,8 +65,9 @@ def _make_session(model, target, zephyr_board, mod):
     project_dir = os.path.join(os.path.dirname(__file__) or ".", "zephyr-runtime")
     compiler = zephyr.ZephyrCompiler(
         project_dir=project_dir,
-        board="nucleo_f746zg" if "stm32f746" in str(target) else "qemu_x86",
+        board=zephyr_board,
         zephyr_toolchain_variant="zephyr",
+        west_cmd=west_cmd,
     )
 
     opts = tvm.micro.default_options(f"{project_dir}/crt")
@@ -106,24 +107,25 @@ def _make_session(model, target, zephyr_board, mod):
     return tvm.micro.Session(**session_kw)
 
 
-def _make_add_sess(model, zephyr_board):
+def _make_add_sess(model, zephyr_board, west_cmd):
     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, "add", sched, [A, B, C])
+    return _make_sess_from_op(model, zephyr_board, west_cmd, "add", sched, [A, B, C])
 
 
 # The models that should pass this configuration. Maps a short, identifying platform string to
 # (model, zephyr_board).
 PLATFORMS = {
     "host": ("host", "qemu_x86"),
     "stm32f746xx": ("stm32f746xx", "nucleo_f746zg"),
+    "nrf5340dk": ("nrf5340dk", "nrf5340dk_nrf5340_cpuapp"),
 }
 
 
 # The same test code can be executed on both the QEMU simulation and on real hardware.
-def test_compile_runtime(platform):
+def test_compile_runtime(platform, west_cmd):
     """Test compiling the on-device runtime."""
 
     model, zephyr_board = PLATFORMS[platform]
@@ -141,11 +143,11 @@ def test_basic_add(sess):
         system_lib.get_function("add")(A_data, B_data, C_data)
         assert (C_data.asnumpy() == np.array([6, 7])).all()
 
-    with _make_add_sess(model, zephyr_board) as sess:
+    with _make_add_sess(model, zephyr_board, west_cmd) as sess:
         test_basic_add(sess)
 
 
-def test_platform_timer(platform):
+def test_platform_timer(platform, west_cmd):
     """Test compiling the on-device runtime."""
 
     model, zephyr_board = PLATFORMS[platform]
@@ -168,11 +170,11 @@ def test_basic_add(sess):
         assert result.mean > 0
         assert len(result.results) == 3
 
-    with _make_add_sess(model, zephyr_board) as sess:
+    with _make_add_sess(model, zephyr_board, west_cmd) as sess:
         test_basic_add(sess)
 
 
-def test_relay(platform):
+def test_relay(platform, west_cmd):
     """Testing a simple relay graph"""
     model, zephyr_board = PLATFORMS[platform]
     shape = (10,)
@@ -188,7 +190,7 @@ def test_relay(platform):
     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, mod) as session:
+    with _make_session(model, target, zephyr_board, west_cmd, mod) as session:
         graph_mod = tvm.micro.create_local_graph_runtime(
             graph, session.get_system_lib(), session.context
         )
@@ -254,14 +256,14 @@ def visit_call(self, call):
         return super().visit_call(call)
 
 
-def check_result(relay_mod, model, zephyr_board, map_inputs, out_shape, result):
+def check_result(relay_mod, model, zephyr_board, west_cmd, map_inputs, out_shape, result):
     """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, mod) as session:
+    with _make_session(model, target, zephyr_board, west_cmd, mod) as session:
         rt_mod = tvm.micro.create_local_graph_runtime(
             graph, session.get_system_lib(), session.context
         )
@@ -280,7 +282,7 @@ def check_result(relay_mod, model, zephyr_board, map_inputs, out_shape, result):
             tvm.testing.assert_allclose(out.asnumpy(), results[idx], rtol=TOL, atol=TOL)
 
 
-def test_byoc_utvm(platform):
+def test_byoc_utvm(platform, west_cmd):
     """This is a simple test case to check BYOC capabilities of uTVM"""
     model, zephyr_board = PLATFORMS[platform]
     x = relay.var("x", shape=(10, 10))
@@ -335,6 +337,7 @@ def test_byoc_utvm(platform):
         ),
         model=model,
         zephyr_board=zephyr_board,
+        west_cmd=west_cmd,
     )