From 8784c9339b043cfb28d0481b9d60cbd0959aa80a Mon Sep 17 00:00:00 2001
From: Tristan O'Connor <tristan.oconnor@arm.com>
Date: Thu, 12 Nov 2020 19:09:43 +0000
Subject: [PATCH] [ETHOSN] Add support for default Ethos-N78 configuration.
 Note: 'ETHOSN_VARIANT_CONFIG' must be set to test against Ethos-N78 and this
 adds support for one configuration of Ethos-N78 in TVM.

---
 .../backend/contrib/ethosn/capabilities.h     | 24 ++++++++++++++-----
 .../backend/contrib/ethosn/codegen_ethosn.h   |  4 +++-
 .../contrib/test_ethosn/infrastructure.py     |  9 ++++++-
 .../contrib/test_ethosn/test_networks.py      |  8 +++++++
 tests/scripts/task_python_ethosn_tests.sh     |  4 ++++
 tests/scripts/task_python_integration.sh      |  3 +++
 6 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/src/relay/backend/contrib/ethosn/capabilities.h b/src/relay/backend/contrib/ethosn/capabilities.h
index 77b2d911d38f..8c7ee6a0d009 100644
--- a/src/relay/backend/contrib/ethosn/capabilities.h
+++ b/src/relay/backend/contrib/ethosn/capabilities.h
@@ -20,7 +20,8 @@
 /*!
  * \file src/relay/backend/contrib/ethosn/capabilities.h
  * \brief The Ethos-N processor series has four variants, the Ethos-N37, Ethos-N57, Ethos-N77
- * and the Ethos-N78. This release of the integration supports the first three variants.
+ * and the Ethos-N78. This release of the integration supports the first three variants and
+ * the default configuration of the fourth variant.
  * Configuration information for each variant is stored as a blob in this file. These blobs
  * are passed into the Ethos-N support library, which in turn uses them to optimize the
  * generated command-stream appropriately for the specified variant.
@@ -38,13 +39,14 @@ namespace relay {
 namespace contrib {
 namespace ethosn {
 
-/* Ethos-N variants (N77, N57 and N37)
- * variant[0] - N77
- * variant[1] - N57
- * variant[2] - N37
+/* Ethos-N variants (Ethos-N77, Ethos-N57, Ethos-N37 and Ethos-N78)
+ * variant[0] - Ethos-N77
+ * variant[1] - Ethos-N57
+ * variant[2] - Ethos-N37
+ * variant[3] - Ethos-N78
  */
 #if _ETHOSN_API_VERSION_ == 2008
-static std::vector<char> variants[3] = {
+static std::vector<char> variants[4] = {
     {
         0x03, 0x00, 0x00, 0x00, 0x78, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
         0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x00, 0x10, 0x00,
@@ -74,6 +76,16 @@ static std::vector<char> variants[3] = {
         0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00,
         0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x01,
         0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    },
+    {
+        0x03, 0x00, 0x00, 0x00, 0x78, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x07, 0x00, 0x02, 0x00,
+        0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x00,
+        0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00,
+        0x01, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00,
+        0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00,
+        0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x01,
+        0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
     }};
 #else
 static std::vector<char> variants[3] = {
diff --git a/src/relay/backend/contrib/ethosn/codegen_ethosn.h b/src/relay/backend/contrib/ethosn/codegen_ethosn.h
index 4b3e1bc05367..9887a2b3ad78 100644
--- a/src/relay/backend/contrib/ethosn/codegen_ethosn.h
+++ b/src/relay/backend/contrib/ethosn/codegen_ethosn.h
@@ -251,7 +251,9 @@ struct EthosnCompilerConfigNode : public tvm::AttrsNode<EthosnCompilerConfigNode
 
   TVM_DECLARE_ATTRS(EthosnCompilerConfigNode, "ext.attrs.EthosnCompilerConfigNode") {
     TVM_ATTR_FIELD(variant)
-        .describe("0 for Ethos-N77, 1 for Ethos-N57, 2 for Ethos-N37. See Ethos-N documentation.")
+        .describe(
+            "0 for Ethos-N77, 1 for Ethos-N57, 2 for Ethos-N37,"
+            " 3 for Ethos-N78. See Ethos-N documentation.")
         .set_default(0);
     TVM_ATTR_FIELD(strategy0).set_default(true);
     TVM_ATTR_FIELD(strategy1).set_default(true);
diff --git a/tests/python/contrib/test_ethosn/infrastructure.py b/tests/python/contrib/test_ethosn/infrastructure.py
index 905d066ce7a3..cd9e9e91292d 100644
--- a/tests/python/contrib/test_ethosn/infrastructure.py
+++ b/tests/python/contrib/test_ethosn/infrastructure.py
@@ -151,7 +151,7 @@ def build(mod, params, npu=True, expected_host_ops=0, npu_partitions=1):
     """
     relay.backend.compile_engine.get().clear()
     with tvm.transform.PassContext(
-        opt_level=3, config={"relay.ext.ethos-n.options": {"variant": 0}}
+        opt_level=3, config={"relay.ext.ethos-n.options": {"variant": get_ethosn_variant()}}
     ):
         with tvm.target.Target("llvm"):
             if npu:
@@ -321,3 +321,10 @@ def get_conv2d_qnn_params(input_zp, input_sc, kernel_zp, kernel_sc, kernel_h, ke
 
 def get_ethosn_api_version():
     return tvm.get_global_func("relay.ethos-n.api.version")()
+
+
+def get_ethosn_variant():
+    ethosn_variant_config = os.getenv("ETHOSN_VARIANT_CONFIG")
+    if ethosn_variant_config is not None:
+        return 3
+    return 0
diff --git a/tests/python/contrib/test_ethosn/test_networks.py b/tests/python/contrib/test_ethosn/test_networks.py
index c9247884141b..e0eccdfb30f5 100644
--- a/tests/python/contrib/test_ethosn/test_networks.py
+++ b/tests/python/contrib/test_ethosn/test_networks.py
@@ -125,6 +125,8 @@ def test_mobilenet_v1():
     _compile_hash = {"81637c89339201a07dc96e3b5dbf836a"}
     if tei.get_ethosn_api_version() == 2008:
         _compile_hash = {"47e216d8ab2bf491708ccf5620bc0d02"}
+        if tei.get_ethosn_variant() == 3:
+            _compile_hash = {"2436f523e263f66a063cef902f2f43d7"}
     _test_image_network(
         model_url="https://storage.googleapis.com/download.tensorflow.org/"
         "models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz",
@@ -147,6 +149,8 @@ def test_inception_v3():
     _compile_hash = {"de0e175af610ebd45ccb03d170dc9664"}
     if tei.get_ethosn_api_version() == 2008:
         _compile_hash = {"8c9d75659cd7bc9ff6dd6d490d28f9b2"}
+        if tei.get_ethosn_variant() == 3:
+            _compile_hash = {"cdd4d7f6453d722ea73224ff9d6a115a"}
     _test_image_network(
         model_url="https://storage.googleapis.com/download.tensorflow.org/"
         "models/tflite_11_05_08/inception_v3_quant.tgz",
@@ -165,6 +169,8 @@ def test_inception_v4():
     # codegen, which could come about from either a change in Support Library
     # version or a change in the Ethos-N codegen. To update this requires running
     # on hardware that isn't available in CI.
+    if not tei.get_ethosn_variant() == 0:
+        pytest.skip("Ethos-N78 20.08 does not support inception_v4 in the default configuration.")
     _compile_hash = {"06bf6cb56344f3904bcb108e54edfe87"}
     if tei.get_ethosn_api_version() == 2008:
         _compile_hash = {"798292bfa596ca7c32086396b494b46c"}
@@ -189,6 +195,8 @@ def test_ssd_mobilenet_v1():
     _compile_hash = {"29aec6b184b09454b4323271aadf89b1", "6211d96103880b016baa85e638abddef"}
     if tei.get_ethosn_api_version() == 2008:
         _compile_hash = {"5999f26e140dee0d7866491997ef78c5", "24e3a690a7e95780052792d5626c85be"}
+        if tei.get_ethosn_variant() == 3:
+            _compile_hash = {"da871b3f03a93df69d704ed44584d6cd", "9f52411d301f3cba3f6e4c0f1c558e87"}
     _test_image_network(
         model_url="https://storage.googleapis.com/download.tensorflow.org/"
         "models/tflite/coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.zip",
diff --git a/tests/scripts/task_python_ethosn_tests.sh b/tests/scripts/task_python_ethosn_tests.sh
index 981d5715fac6..ae9b82b679ef 100755
--- a/tests/scripts/task_python_ethosn_tests.sh
+++ b/tests/scripts/task_python_ethosn_tests.sh
@@ -27,4 +27,8 @@ source tests/scripts/setup-pytest-env.sh
 find . -type f -path "*.pyc" | xargs rm -f
 make cython3
 
+# Note: Default behaviour is to assume the test target is Ethos-N77
+# but setting ETHOSN_VARIANT_CONFIG appropriately
+# (e.g. ETHOSN_VARIANT_CONFIG=ETHOSN78_1TOPS_4PLE_448KSRAM)
+# switches the target to an Ethos-N78 configuration.
 run_pytest ctypes python-ethosn tests/python/contrib/test_ethosn
diff --git a/tests/scripts/task_python_integration.sh b/tests/scripts/task_python_integration.sh
index e21aa065a024..613c7cbdf34f 100755
--- a/tests/scripts/task_python_integration.sh
+++ b/tests/scripts/task_python_integration.sh
@@ -60,6 +60,9 @@ run_pytest cython ${TVM_INTEGRATION_TESTSUITE_NAME}-dso_plugin_module apps/dso_p
 # TVM_FFI=ctypes sh prepare_and_test_tfop_module.sh
 
 run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME} tests/python/integration
+if python -c "import tvm; from tvm.relay.op.contrib.ethosn import ethosn_available; print(ethosn_available().name)" -eq "SW_ONLY"; then
+  ETHOSN_VARIANT_CONFIG=ETHOSN78_1TOPS_4PLE_448KSRAM run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME}-contrib-test_ethosn tests/python/contrib/test_ethosn
+fi
 run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME}-contrib tests/python/contrib
 
 TVM_TEST_TARGETS="${TVM_RELAY_TEST_TARGETS:-llvm;cuda}" \