Add Arm(R) Ethos(TM)-U codegen support on tvmc

* Include `ethos-u` as a new target for tvmc * Adds testing for the new target Co-authored-by: Manupa Karunaratne <[email protected]>
apache · Sep 27, 2021 · 25e7588 · 25e7588
1 parent 4f67889
commit 25e7588
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 12 deletions.
diff --git a/python/tvm/driver/tvmc/composite_target.py b/python/tvm/driver/tvmc/composite_target.py
@@ -25,6 +25,7 @@
 from tvm.relay.op.contrib.arm_compute_lib import partition_for_arm_compute_lib
 from tvm.relay.op.contrib.ethosn import partition_for_ethosn
 from tvm.relay.op.contrib.cmsisnn import partition_for_cmsisnn
+from tvm.relay.backend.contrib.ethosu import partition_for_ethosu
 from tvm.relay.op.contrib.bnns import partition_for_bnns
 from tvm.relay.op.contrib.vitis_ai import partition_for_vitis_ai
 
@@ -58,6 +59,10 @@
         "config_key": "relay.ext.ethos-n.options",
         "pass_pipeline": partition_for_ethosn,
     },
+    "ethos-u": {
+        "config_key": "relay.ext.ethosu.options",
+        "pass_pipeline": partition_for_ethosu,
+    },
     "bnns": {
         "config_key": None,
         "pass_pipeline": partition_for_bnns,

diff --git a/tests/python/contrib/test_ethosu/test_vela_api.py b/tests/python/contrib/test_ethosu/test_vela_api.py
@@ -347,9 +347,7 @@ def verify(test_vec, mock_obj):
         assert mock_obj.call_args[1]["block_traversal"] == test_vec["block_traversal"]
 
     def create_mock(test_vec):
-        with patch(
-            "tvm.relay.backend.contrib.ethosu.vela_api.vapi.npu_encode_weights"
-        ) as mock_npu_encode_weights:
+        with patch("ethosu.vela.api.npu_encode_weights") as mock_npu_encode_weights:
             ifm_bitdepth = np.iinfo(test_vec["ifm_dtype"]).bits
             ifm_dtype = test_vec["ifm_dtype"]
             max = np.iinfo(ifm_dtype).max
@@ -427,9 +425,7 @@ def verify(test_vec, mock_obj, packed_biases):
             assert test_vec["hw_shifts"][idx] == mock_obj.call_args_list[idx][0][2]
 
     def create_mock(test_vec):
-        with patch(
-            "tvm.relay.backend.contrib.ethosu.vela_api.vapi.npu_encode_bias"
-        ) as mock_npu_encode_bias:
+        with patch("ethosu.vela.api.npu_encode_bias") as mock_npu_encode_bias:
             mock_npu_encode_bias.return_value = bytearray(10)
             ifm_dtype = test_vec["ifm_dtype"]
             max = np.iinfo(ifm_dtype).max
@@ -507,12 +503,8 @@ def test_encode_weights(accel):
     ]
 
     def create_mock(test_vec):
-        with patch(
-            "tvm.relay.backend.contrib.ethosu.vela_api.vapi.npu_encode_weights"
-        ) as mock_enc_w:
-            with patch(
-                "tvm.relay.backend.contrib.ethosu.vela_api.vapi.npu_find_block_configs"
-            ) as mock_blk_cfg:
+        with patch("ethosu.vela.api.npu_encode_weights") as mock_enc_w:
+            with patch("ethosu.vela.api.npu_find_block_configs") as mock_blk_cfg:
                 mock_blk_cfg.return_value = [vapi.NpuShape3D(8, 8, 8)]
                 ethosu_conv2d_calls = extract_ethosu_conv2d_extern_calls(test_vec["tir_module"])
                 buffer_info = tirtocs.extract_buffer_info(

diff --git a/tests/python/driver/tvmc/test_compiler.py b/tests/python/driver/tvmc/test_compiler.py
@@ -361,6 +361,41 @@ def test_compile_tflite_module_with_external_codegen_vitis_ai(tflite_mobilenet_v
     assert os.path.exists(dumps_path)
 
 
+def test_compile_tflite_module_with_external_codegen_ethosu(
+    tmpdir_factory, tflite_mobilenet_v1_1_quant
+):
+    pytest.importorskip("tflite")
+    pytest.importorskip("ethosu.vela")
+    ACCEL_TYPES = ["ethos-u55-256", "ethos-u55-128", "ethos-u55-64", "ethos-u55-32"]
+
+    output_dir = tmpdir_factory.mktemp("mlf")
+
+    tvmc_model = tvmc.load(tflite_mobilenet_v1_1_quant)
+
+    for accel_type in ACCEL_TYPES:
+        output_file_name = f"{output_dir}/file_{accel_type}.tar"
+
+        tvmc_package = tvmc.compiler.compile_model(
+            tvmc_model,
+            target=f"ethos-u -accelerator_config={accel_type}, c -runtime=c --system-lib --link-params -mcpu=cortex-m55 --executor=aot",
+            output_format="mlf",
+            package_path=output_file_name,
+            pass_context_configs=["tir.disable_vectorize=true"],
+        )
+
+        # check whether an MLF package was created
+        assert os.path.exists(output_file_name)
+
+        # check whether the expected number of C sources are in the tarfile
+        with tarfile.open(output_file_name) as mlf_package:
+            c_source_files = [
+                name
+                for name in mlf_package.getnames()
+                if re.match(r"\./codegen/host/src/\D+\d+\.c", name)
+            ]
+            assert len(c_source_files) == 17
+
+
 @mock.patch("tvm.relay.build")
 @mock.patch("tvm.driver.tvmc.composite_target.get_codegen_by_target")
 @mock.patch("tvm.driver.tvmc.load")