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

* move partition_for_ethosu from tvm.relay.backend.contrib.ethosu
  to tvm.relay.op.contrib.ethosu

* lazy load ethos-u-vela dependencies and show an appropriate error
  message in case the dependency is not present

* Adjust test cases

Co-authored-by: Leandro Nunes <Leandro.Nunes@arm.com>

python/tvm/driver/tvmc/composite_target.py

@@ -25,7 +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.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
 

python/tvm/relay/backend/contrib/ethosu/__init__.py

@@ -22,4 +22,3 @@
 from . import codegen
 from . import vela_api
 from . import tir_to_cs_translator
-from .util import partition_for_ethosu

python/tvm/relay/backend/contrib/ethosu/util.py

@@ -23,13 +23,11 @@
 
 from inspect import signature
 from enum import Enum
-from typing import Union, Tuple, Dict, Optional
+from typing import Union, Tuple
 import numpy as np  # type: ignore
 
 import tvm  # type: ignore
 from tvm import relay
-from tvm.relay.build_module import bind_params_by_name  # type: ignore
-from tvm.relay.backend.contrib.ethosu import preprocess  # type: ignore
 
 
 class QConv2DArgs(Enum):
@@ -145,41 +143,6 @@ def get_accelerator_config():
     return compiler_attrs.accelerator_config
 
 
-# pylint: disable=unused-argument
-def partition_for_ethosu(
-    mod: tvm.ir.IRModule, params: Optional[Dict[str, tvm.runtime.NDArray]] = None, **opts
-):
-    """This helper function partition the relay graph as produced by the
-    relay frontend for a given model into external functions
-    to be presented to the codegen.
-
-    Parameters
-    ----------
-    mod : tvm.ir.IRModule
-        The IRModule that gets generated from a relay frontend
-    params : Optional[Dict[str, tvm.runtime.NDArray]]
-        Constant input parameters.
-
-    Returns
-    -------
-    mod : IRModule
-        The partitioned IRModule with external global functions
-    """
-    if params:
-        mod["main"] = bind_params_by_name(mod["main"], params)
-
-    pattern = relay.op.contrib.get_pattern_table("ethosu")
-    mod = relay.transform.InferType()(mod)
-    mod = relay.transform.MergeComposite(pattern)(mod)
-    mod = relay.transform.AnnotateTarget("ethosu")(mod)
-    mod = relay.transform.MergeCompilerRegions()(mod)
-    mod = relay.transform.InferType()(mod)
-    mod = relay.transform.PartitionGraph()(mod)
-    mod = relay.transform.InferType()(mod)
-    mod = preprocess.preprocess_ext_io()(mod)
-    return mod
-
-
 def get_arg_count(func):
     """Helper function to get the number of
     arguments in a python function"""

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

@@ -14,19 +14,50 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+# pylint: disable=ungrouped-imports
 """Arm(R) Ethos(TM)-U NPU supported operators."""
-from typing import List, Tuple, Callable
+import functools
+
+from typing import Dict, List, Tuple, Callable, Optional
 import numpy as np  # type: ignore
 
 import tvm  # type: ignore
+from tvm import relay
 from tvm.relay.expr import Constant  # type: ignore
 from tvm.relay.op.contrib.register import register_pattern_table  # type: ignore
 from tvm.relay.dataflow_pattern import wildcard, is_op, is_constant  # type: ignore
-from tvm.relay.backend.contrib.ethosu.util import QConv2DArgs  # type: ignore
-from tvm.relay.backend.contrib.ethosu.util import BiasAddArgs
-from tvm.relay.backend.contrib.ethosu.util import RequantArgs
-from tvm.relay.backend.contrib.ethosu.util import get_dim_value
-from ethosu.vela import api as vapi  # type: ignore
+
+try:
+    # As ethos-u-vela package is an optional TVM dependency, we want to lazy load it
+    # and check whether it is installed or not.
+    #
+    # In order to show the appropriate error messages when we try to invoke code that
+    # rely on imports from ethos-u-vela, we protect them with the decorator @requires_vela
+    # implemented below.
+    from ethosu.vela import api as vapi  # type: ignore
+    from tvm.relay.backend.contrib.ethosu import preprocess
+    from tvm.relay.backend.contrib.ethosu.util import QConv2DArgs  # type: ignore
+    from tvm.relay.backend.contrib.ethosu.util import BiasAddArgs
+    from tvm.relay.backend.contrib.ethosu.util import RequantArgs
+    from tvm.relay.backend.contrib.ethosu.util import get_dim_value
+except ImportError:
+    vapi = None
+
+
+def requires_vela(func):
+    """Decorator to check whether we have the required dependency ethos-u-vela
+    installed as a python package"""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        if not vapi:
+            raise ImportError(
+                "The 'ethos-u-vela' python package is required for the Arm(R) Ethos(TM)-U NPU "
+                "backend. Please install the dependency using your Python package manager."
+            ) from None
+        return func(*args, **kwargs)
+
+    return wrapper
 
 
 class TensorParams:
@@ -36,6 +67,7 @@ class TensorParams:
     for the creation of tensors in Vela.
     """
 
+    @requires_vela
     def __init__(self, tensor, layout=None, scale=None, zero_point=None):
         self.tensor = tensor
         if isinstance(tensor, Constant):
@@ -148,6 +180,7 @@ class QnnConv2DParams:
     padding_bounds = [31, 31, 32, 32]
     activation_map = {"clip": "CLIP"}
 
+    @requires_vela
     def __init__(self, func_body: tvm.relay.Function):
         activation = None
         if str(func_body.op) in self.activation_map.keys():
@@ -247,3 +280,39 @@ def pattern_table() -> List[Tuple[str, tvm.relay.dataflow_pattern.DFPattern, Cal
             lambda pat: QnnConv2DParams(pat).is_valid(),
         )
     ]
+
+
+# pylint: disable=unused-argument
+@requires_vela
+def partition_for_ethosu(
+    mod: tvm.ir.IRModule, params: Optional[Dict[str, tvm.runtime.NDArray]] = None, **opts
+):
+    """This helper function partition the relay graph as produced by the
+    relay frontend for a given model into external functions
+    to be presented to the codegen.
+
+    Parameters
+    ----------
+    mod : tvm.ir.IRModule
+        The IRModule that gets generated from a relay frontend
+    params : Optional[Dict[str, tvm.runtime.NDArray]]
+        Constant input parameters.
+
+    Returns
+    -------
+    mod : IRModule
+        The partitioned IRModule with external global functions
+    """
+    if params:
+        mod["main"] = bind_params_by_name(mod["main"], params)
+
+    pattern = relay.op.contrib.get_pattern_table("ethosu")
+    mod = relay.transform.InferType()(mod)
+    mod = relay.transform.MergeComposite(pattern)(mod)
+    mod = relay.transform.AnnotateTarget("ethosu")(mod)
+    mod = relay.transform.MergeCompilerRegions()(mod)
+    mod = relay.transform.InferType()(mod)
+    mod = relay.transform.PartitionGraph()(mod)
+    mod = relay.transform.InferType()(mod)
+    mod = preprocess.preprocess_ext_io()(mod)
+    return mod

tests/python/contrib/test_ethosu/test_codegen.py

@@ -27,6 +27,7 @@
 from tvm import relay
 from tvm.relay.backend.contrib import ethosu
 from tvm.relay.backend.contrib.ethosu import util
+from tvm.relay.op.contrib.ethosu import partition_for_ethosu
 from tests.python.relay.aot.aot_test_utils import generate_ref_data
 
 from . import relay_ir_builder
@@ -139,7 +140,7 @@ def create_graph_activation(input_tensor_name, input_tensor_shape, input_tensor_
     for test_case in test_cases:
         relay_module, conv_params = test_case[0](*test_case[1])
         input_tensor, input_shape, input_dtype = test_case[1]
-        mod = ethosu.partition_for_ethosu(relay_module)
+        mod = partition_for_ethosu(relay_module)
 
         # Generate reference data
         in_min, in_max = util.get_range_for_dtype_str(input_dtype)

tests/python/contrib/test_ethosu/test_legalize.py

@@ -294,7 +294,7 @@ def verify_linear(ext_func, conv2d_params):
     ]
     for test_case in test_cases:
         mod, conv_params = test_case[0](*test_case[1])
-        mod = ethosu.partition_for_ethosu(mod)
+        mod = partition_for_ethosu(mod)
         mod = legalize.LegalizeEthosUConv2D()(mod)
         verify_linear(mod["tvmgen_default_ethosu_main_0"], conv_params)
 
@@ -327,7 +327,7 @@ def create_graph_single_unsupported_ifm_layout(
 
     for test_case in test_cases:
         mod, conv_params = test_case[0](*test_case[1])
-        mod = ethosu.partition_for_ethosu(mod)
+        mod = partition_for_ethosu(mod)
         with pytest.raises(
             tvm._ffi.base.TVMError, match="EthosUCodegenError: Unsupported Layout NCHW"
         ):

tests/python/contrib/test_ethosu/test_networks.py

@@ -27,8 +27,9 @@
 import tvm
 import tvm.micro as micro
 from tvm import relay
-from tvm.relay.backend.contrib import ethosu
 from tvm.relay.backend.contrib.ethosu import util
+from tvm.relay.op.contrib.ethosu import partition_for_ethosu
+
 import tvm.relay.testing.tf as tf_testing
 
 from . import infra
@@ -56,7 +57,7 @@ def test_forward_mobilenet_v1(accel_type="ethos-u55-256"):
     input_data = {input_tensor: input_data}
     output_data = generate_ref_data(relay_mod, input_data)
 
-    mod = ethosu.partition_for_ethosu(relay_mod, params)
+    mod = partition_for_ethosu(relay_mod, params)
     compiled_models = infra.build_source(mod, input_data, output_data, accel_type)
     infra.verify_source(compiled_models, accel_type)