[ETHOSN] Add support for non-default Ethos(TM)-N78 configurations (#9386

)

- Updated tvmc with new Ethos-N78 composite target.
- Added additional Ethos-N78 specific configuration options.

Co-authored-by: Tristan O'Connor <tristan.oconnor@arm.com>

Co-authored-by: Tristan O'Connor <tristan.oconnor@arm.com>

python/tvm/driver/tvmc/composite_target.py

@@ -23,7 +23,8 @@
 import tvm.contrib.target.vitis_ai  # pylint: disable=unused-import
 
 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.ethosn import partition_for_ethosn77
+from tvm.relay.op.contrib.ethosn import partition_for_ethosn78
 from tvm.relay.op.contrib.cmsisnn import partition_for_cmsisnn
 from tvm.relay.op.contrib.ethosu import partition_for_ethosu
 from tvm.relay.op.contrib.bnns import partition_for_bnns
@@ -57,7 +58,11 @@
     },
     "ethos-n77": {
         "config_key": "relay.ext.ethos-n.options",
-        "pass_pipeline": partition_for_ethosn,
+        "pass_pipeline": partition_for_ethosn77,
+    },
+    "ethos-n78": {
+        "config_key": "relay.ext.ethos-n.options",
+        "pass_pipeline": partition_for_ethosn78,
     },
     "ethos-u": {
         "config_key": "relay.ext.ethosu.options",

python/tvm/relay/op/contrib/ethosn.py

@@ -46,7 +46,7 @@ def ethosn_available():
     return Available.SW_AND_HW if hw else Available.SW_ONLY
 
 
-def partition_for_ethosn(mod, params=None, **opts):
+def partition_for_ethosn77(mod, params=None, **opts):
     """Partition the graph greedily offloading supported
     operators to Arm Ethos-N NPU.
 
@@ -61,6 +61,49 @@ def partition_for_ethosn(mod, params=None, **opts):
     -------
     ret : annotated and partitioned module.
     """
+    if opts:
+        tops = opts.get("tops", None)
+        ple_ratio = opts.get("ple_ratio", None)
+        sram_size = opts.get("sram_size", None)
+        if tops or ple_ratio or sram_size:
+            raise ValueError(
+                "Setting tops, ple_ratio or sram_size has no effect when targeting Ethos(TM)-N77"
+            )
+
+    if params:
+        mod["main"] = bind_params_by_name(mod["main"], params)
+
+    seq = tvm.transform.Sequential(
+        [
+            transform.InferType(),
+            transform.MergeComposite(pattern_table()),
+            transform.AnnotateTarget("ethos-n"),
+            transform.MergeCompilerRegions(),
+            transform.PartitionGraph(),
+        ]
+    )
+
+    return seq(mod)
+
+
+def partition_for_ethosn78(mod, params=None, **opts):
+    """Partition the graph greedily offloading supported
+    operators to Arm Ethos(TM)-N NPU.
+
+    Parameters
+    ----------
+    mod : Module
+        The module to run passes on.
+    params : Optional[Dict[str, NDArray]]
+        Constant input parameters.
+
+    Returns
+    -------
+    ret : annotated and partitioned module.
+    """
+    if not opts or opts.get("variant", "").lower() != "ethos-n78":
+        raise ValueError("When targeting Ethos(TM)-N78, -variant=Ethos-N78 should be set.")
+
     if params:
         mod["main"] = bind_params_by_name(mod["main"], params)
 

src/relay/backend/contrib/ethosn/codegen.cc

@@ -195,6 +195,20 @@ sl::TensorsAndId MakeOps(const sl::TensorAndId<sl::Operand>& op) {
   return ops;
 }
 
+String MakeVariant(auto configuration) {
+  String variant = configuration.value()->variant;
+  // Transform variant string to lowercase for comparison
+  std::string variant_string = variant.c_str();
+  std::transform(variant_string.begin(), variant_string.end(), variant_string.begin(), ::tolower);
+  std::string variant_n78 = "ethos-n78";
+  if (variant_string == variant_n78) {
+    String tops = configuration.value()->tops;
+    String ple_ratio = configuration.value()->ple_ratio;
+    variant = "Ethos-N78_" + tops + "TOPS_" + ple_ratio + "PLE_RATIO";
+  }
+  return variant;
+}
+
 NetworkWithIDs ConstructNetworkVisitor::Construct(const Function& func) {
   // Initialise everything
   auto ctx = transform::PassContext::Current();
@@ -203,8 +217,9 @@ NetworkWithIDs ConstructNetworkVisitor::Construct(const Function& func) {
     cfg = AttrsWithDefaultValues<EthosnCompilerConfig>();
   }
   NetworkWithIDs network_with_ids;
-  network_ = sl::CreateNetwork(sl::GetFwAndHwCapabilities(
-      sl::EthosNVariantFromString(cfg.value()->variant.c_str()), cfg.value()->sram_size_bytes));
+  network_ = sl::CreateNetwork(
+      sl::GetFwAndHwCapabilities(sl::EthosNVariantFromString(MakeVariant(cfg).c_str()),
+                                 static_cast<uint32_t>(std::stoul(cfg.value()->sram_size))));
   network_with_ids.network = network_;
   operand_table_.clear();
 
@@ -614,8 +629,9 @@ EthosnError EthosnCompiler::SupportedSetup() {
     auto cfg = ctx->GetConfig<EthosnCompilerConfig>("relay.ext.ethos-n.options").defined()
                    ? ctx->GetConfig<EthosnCompilerConfig>("relay.ext.ethos-n.options")
                    : AttrsWithDefaultValues<EthosnCompilerConfig>();
-    m_Queries = std::make_unique<sl::SupportQueries>(sl::GetFwAndHwCapabilities(
-        sl::EthosNVariantFromString(cfg.value()->variant.c_str()), cfg.value()->sram_size_bytes));
+    m_Queries = std::make_unique<sl::SupportQueries>(
+        sl::GetFwAndHwCapabilities(sl::EthosNVariantFromString(cfg.value()->variant.c_str()),
+                                   std::stoul(cfg.value()->sram_size)));
     if (m_Queries == nullptr) {
       return EthosnError("Could not initialise Ethos-N compiler isSupported");
     }

src/relay/backend/contrib/ethosn/codegen_ethosn.h

@@ -227,7 +227,9 @@ NetworkWithIDs ConstructNetwork(const IRModule& mod, const GlobalVar& var, const
 /*! \brief Attributes to store the compiler options for Ethos-N */
 struct EthosnCompilerConfigNode : public tvm::AttrsNode<EthosnCompilerConfigNode> {
   String variant;
-  int sram_size_bytes;
+  String sram_size;
+  String tops;
+  String ple_ratio;
   bool strategy0;
   bool strategy1;
   bool strategy3;
@@ -247,9 +249,15 @@ struct EthosnCompilerConfigNode : public tvm::AttrsNode<EthosnCompilerConfigNode
 
   TVM_DECLARE_ATTRS(EthosnCompilerConfigNode, "ext.attrs.EthosnCompilerConfigNode") {
     TVM_ATTR_FIELD(variant).describe("See Ethos-N documentation.").set_default("Ethos-N77");
-    TVM_ATTR_FIELD(sram_size_bytes)
-        .describe("Optionally override the default sram size. See Ethos-N documentation.")
-        .set_default(0);
+    TVM_ATTR_FIELD(sram_size)
+        .describe("Optionally override the default sram size. See Ethos(TM)-N documentation.")
+        .set_default("0");
+    TVM_ATTR_FIELD(tops)
+        .describe("Valid values 1, 2, 4 and 8. See Ethos(TM)-N documentation.")
+        .set_default("1");
+    TVM_ATTR_FIELD(ple_ratio)
+        .describe("Valid values 2 and 4. See Ethos(TM)-N documentation.")
+        .set_default("2");
     TVM_ATTR_FIELD(strategy0).set_default(true);
     TVM_ATTR_FIELD(strategy1).set_default(true);
     TVM_ATTR_FIELD(strategy3).set_default(true);
@@ -339,6 +347,10 @@ class EthosnCompiler {
   static std::pair<std::vector<uint32_t>, std::vector<uint32_t>> GetInputOutputOrder(
       NetworkWithIDs network, const std::unique_ptr<sl::CompiledNetwork>& compiled_network);
 
+  /*!
+   * \brief Query interface used to determine if the Ethos-N hardware supports an operation
+   * with the supplied parameters.
+   */
   static std::unique_ptr<sl::SupportQueries> m_Queries;
 };
 

tests/python/contrib/test_ethosn/infrastructure.py

@@ -254,7 +254,9 @@ def inference_result(outputs):
 
 def test_error(mod, params, err_msg):
     caught = None
-    with tvm.transform.PassContext(opt_level=3):
+    with tvm.transform.PassContext(
+        opt_level=3, config={"relay.ext.ethos-n.options": {"variant": get_ethosn_variant()}}
+    ):
         with tvm.target.Target("llvm"):
             try:
                 mod = relay.transform.InferType()(mod)
@@ -324,7 +326,4 @@ def get_ethosn_api_version():
 
 
 def get_ethosn_variant():
-    ethosn_variant_config = os.getenv("ETHOSN_VARIANT_CONFIG")
-    if ethosn_variant_config is not None:
-        return "Ethos-N78_1TOPS_2PLE_RATIO"
-    return "Ethos-N77"
+    return os.getenv("ETHOSN_VARIANT_CONFIG", default="Ethos-N77")

tests/python/contrib/test_ethosn/test_partition_params.py

@@ -0,0 +1,123 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Ethos(TM)-N partition parameter tests"""
+
+import pytest
+import tvm
+from tvm import relay
+import numpy as np
+
+from tvm.relay.op.contrib.ethosn import partition_for_ethosn77
+from tvm.relay.op.contrib.ethosn import partition_for_ethosn78
+from tvm.testing import requires_ethosn
+
+
+@requires_ethosn
+def test_ethosn78_partition_no_error():
+    a = relay.var("a", shape=[2, 7, 8, 8], dtype="uint8")
+    w = relay.const(np.random.uniform(-10, 10, (8, 7, 3, 3)).astype("uint8"))
+    res = relay.nn.conv2d(a, w, kernel_size=(3, 3), padding=(1, 1), channels=8, out_dtype="uint8")
+    b = relay.var("b", shape=[8], dtype="uint8")
+    res = relay.nn.bias_add(res, b, axis=1)
+
+    mod = tvm.IRModule.from_expr(res)
+    opts = {"variant": "Ethos-N78"}
+    partition_for_ethosn78(mod, **opts)
+
+
+@requires_ethosn
+def test_ethosn78_partition_undefined_variant():
+    with pytest.raises(
+        ValueError, match=r".*When targeting Ethos\(TM\)-N78, -variant=Ethos-N78 should be set.*"
+    ):
+        a = relay.var("a", shape=[2, 7, 8, 8], dtype="uint8")
+        w = relay.const(np.random.uniform(-10, 10, (8, 7, 3, 3)).astype("uint8"))
+        res = relay.nn.conv2d(
+            a, w, kernel_size=(3, 3), padding=(1, 1), channels=8, out_dtype="uint8"
+        )
+        b = relay.var("b", shape=[8], dtype="uint8")
+        res = relay.nn.bias_add(res, b, axis=1)
+
+        mod = tvm.IRModule.from_expr(res)
+        partition_for_ethosn78(mod)
+
+
+@requires_ethosn
+def test_ethosn78_partition_invalid_variant():
+    with pytest.raises(
+        ValueError, match=r".*When targeting Ethos\(TM\)-N78, -variant=Ethos-N78 should be set.*"
+    ):
+        a = relay.var("a", shape=[2, 7, 8, 8], dtype="uint8")
+        w = relay.const(np.random.uniform(-10, 10, (8, 7, 3, 3)).astype("uint8"))
+        res = relay.nn.conv2d(
+            a, w, kernel_size=(3, 3), padding=(1, 1), channels=8, out_dtype="uint8"
+        )
+        b = relay.var("b", shape=[8], dtype="uint8")
+        res = relay.nn.bias_add(res, b, axis=1)
+
+        mod = tvm.IRModule.from_expr(res)
+        opts = {"variant": "Ethos-N"}
+        partition_for_ethosn78(mod, **opts)
+
+
+@requires_ethosn
+def test_ethosn78_partition_error():
+    with pytest.raises(
+        ValueError, match=r".*When targeting Ethos\(TM\)-N78, -variant=Ethos-N78 should be set.*"
+    ):
+        a = relay.var("a", shape=[2, 7, 8, 8], dtype="uint8")
+        w = relay.const(np.random.uniform(-10, 10, (8, 7, 3, 3)).astype("uint8"))
+        res = relay.nn.conv2d(
+            a, w, kernel_size=(3, 3), padding=(1, 1), channels=8, out_dtype="uint8"
+        )
+        b = relay.var("b", shape=[8], dtype="uint8")
+        res = relay.nn.bias_add(res, b, axis=1)
+
+        mod = tvm.IRModule.from_expr(res)
+        opts = {"variant": "Ethos-N77"}
+        partition_for_ethosn78(mod, **opts)
+
+
+@requires_ethosn
+def test_ethosn77_partition_no_error():
+    a = relay.var("a", shape=[2, 7, 8, 8], dtype="uint8")
+    w = relay.const(np.random.uniform(-10, 10, (8, 7, 3, 3)).astype("uint8"))
+    res = relay.nn.conv2d(a, w, kernel_size=(3, 3), padding=(1, 1), channels=8, out_dtype="uint8")
+    b = relay.var("b", shape=[8], dtype="uint8")
+    res = relay.nn.bias_add(res, b, axis=1)
+
+    mod = tvm.IRModule.from_expr(res)
+    partition_for_ethosn77(mod)
+
+
+@requires_ethosn
+def test_ethosn77_partition_error():
+    with pytest.raises(
+        ValueError,
+        match=r".*Setting tops, ple_ratio or sram_size has no effect when targeting Ethos\(TM\)-N77.*",
+    ):
+        a = relay.var("a", shape=[2, 7, 8, 8], dtype="uint8")
+        w = relay.const(np.random.uniform(-10, 10, (8, 7, 3, 3)).astype("uint8"))
+        res = relay.nn.conv2d(
+            a, w, kernel_size=(3, 3), padding=(1, 1), channels=8, out_dtype="uint8"
+        )
+        b = relay.var("b", shape=[8], dtype="uint8")
+        res = relay.nn.bias_add(res, b, axis=1)
+
+        mod = tvm.IRModule.from_expr(res)
+        opts = {"tops": 4}
+        partition_for_ethosn77(mod, **opts)

tests/python/driver/tvmc/test_compiler.py

@@ -24,7 +24,7 @@
 import pytest
 
 import tvm
-import tvm.testing
+from tvm.testing.utils import ethosn_available
 
 from tvm.contrib.target.vitis_ai import vitis_ai_available
 
@@ -370,8 +370,11 @@ def test_compile_opencl(tflite_mobilenet_v1_0_25_128):
     assert os.path.exists(dumps_path)
 
 
-@tvm.testing.requires_ethosn
-def test_compile_tflite_module_with_external_codegen(tflite_mobilenet_v1_1_quant):
+@pytest.mark.skipif(
+    not ethosn_available(),
+    reason="--target=Ethos(TM)-N77 is not available. TVM built with 'USE_ETHOSN OFF'",
+)
+def test_compile_tflite_module_with_external_codegen_ethos_n77(tflite_mobilenet_v1_1_quant):
     pytest.importorskip("tflite")
     tvmc_model = tvmc.load(tflite_mobilenet_v1_1_quant)
     tvmc_package = tvmc.compile(tvmc_model, target="ethos-n77, llvm", dump_code="relay")
@@ -416,6 +419,26 @@ def test_compile_tflite_module_with_external_codegen_cmsisnn(
         assert len(c_source_files) == 3
 
 
+@pytest.mark.skipif(
+    not ethosn_available(),
+    reason="--target=Ethos(TM)-N78 is not available. TVM built with 'USE_ETHOSN OFF'",
+)
+def test_compile_tflite_module_with_external_codegen_ethos_n78(tflite_mobilenet_v1_1_quant):
+    pytest.importorskip("tflite")
+    tvmc_model = tvmc.load(tflite_mobilenet_v1_1_quant)
+    tvmc_package = tvmc.compile(
+        tvmc_model, target="ethos-n78 -variant=ethos-n78, llvm", dump_code="relay"
+    )
+    dumps_path = tvmc_package.package_path + ".relay"
+
+    # check for output types
+    assert type(tvmc_package) is TVMCPackage
+    assert type(tvmc_package.graph) is str
+    assert type(tvmc_package.lib_path) is str
+    assert type(tvmc_package.params) is bytearray
+    assert os.path.exists(dumps_path)
+
+
 @pytest.mark.skipif(
     not vitis_ai_available(),
     reason="--target=vitis-ai is not available. TVM built with 'USE_VITIS_AI OFF'",

tests/python/driver/tvmc/test_composite_target.py

@@ -34,6 +34,7 @@ def test_get_codegen_names():
     names = tvmc.composite_target.get_codegen_names()
 
     assert "ethos-n77" in names
+    assert "ethos-n78" in names
     assert "vitis-ai" in names
     assert len(names) > 0