From 25e7588a6e8a579b3184a0cc095b4229824365b3 Mon Sep 17 00:00:00 2001 From: Leandro Nunes Date: Tue, 18 May 2021 11:32:05 +0100 Subject: [PATCH] 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 --- python/tvm/driver/tvmc/composite_target.py | 5 +++ .../contrib/test_ethosu/test_vela_api.py | 16 +++------ tests/python/driver/tvmc/test_compiler.py | 35 +++++++++++++++++++ 3 files changed, 44 insertions(+), 12 deletions(-) diff --git a/python/tvm/driver/tvmc/composite_target.py b/python/tvm/driver/tvmc/composite_target.py index b1f00b7d1ddec..77eeb81c6d434 100644 --- 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 index a86dd919d5caf..1f9c72e1c9cce 100644 --- 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 index defd628c60c92..3993fdae86f4c 100644 --- 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")