[5/6] Arm(R) Ethos(TM)-U NPU codegen integration

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

@@ -19,6 +19,7 @@
 from . import legalize
 from . import preprocess
 from . import errors
+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/codegen.py

@@ -0,0 +1,83 @@
+# 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.
+"""Codegen for Arm(R) Ethos(TM)-U"""
+import tvm
+from tvm import relay
+from tvm.relay.backend.contrib.ethosu.tir.compiler import lower_to_tir
+from tvm.relay.backend.contrib.ethosu.tir.scheduler import copy_constants
+from tvm.relay.backend.contrib.ethosu.legalize import LegalizeEthosU
+from tvm.relay.backend.contrib.ethosu import tir_to_cs_translator
+from tvm.relay.backend.contrib.ethosu import util
+
+
+@tvm._ffi.register_func("relay.ext.ethosu.constant_updater")
+def constant_updater(expr, symbol):  # pylint: disable=unused-argument
+    """
+    We dont want the build process to extract constants to be loaded in
+    the runtime as we are embedding them inside the C runtime.Module.
+    """
+    return dict()
+
+
+@tvm._ffi.register_func("relay.ext.ethosu")
+def ethosu_compiler(ref):
+    """Main function to a compile a given relay function of
+    NPU compatible operators to generated command stream.
+    Such generated command stream would be loaded to the runtime
+    module that interfaces with NPU driver.
+    """
+    assert isinstance(ref, tvm.ir.function.BaseFunc)
+    func_name = ref.attrs["global_symbol"]
+    # There should only be a single input
+    assert len(ref.params) == 1
+    input_size = util.calculate_size_bytes(ref.params[0])
+    output_size = util.calculate_size_bytes(ref.body)
+    cmms, encoded_constants, scratch_size = _compile(ref)
+    ethosu_runtime = tvm._ffi.get_global_func("runtime.module.ethosu.create")
+    return ethosu_runtime(func_name, cmms, encoded_constants, scratch_size, input_size, output_size)
+
+
+def _compile(ext_func):
+    """
+    This is the main wrapper that accepts an external
+    relay function and runs all the passes to lower it down
+    to command stream
+    Parameters
+    ----------
+    ext_func : tvm.relay.function.Function
+        The partitioned relay function
+    Returns
+    -------
+    cs : str
+        An hex string of the bytes of command stream
+    encoded_constants : str
+        An hex string of the bytes that includes concat'd
+        encoded weights, encoded biases and scales.
+    scratch_size : int
+        The size of the scratch buffer needed.
+    """
+    mod = tvm.IRModule()
+    mod["main"] = ext_func
+    mod = LegalizeEthosU()(mod)
+    mod = relay.transform.InferType()(mod)
+    # We are currently using copy_constants scheduler In the long run,
+    # this should be a single intelligent and a composite scheduler
+    # that can perform scheduling based on user inputs such as
+    # scratch memory size.
+    tir_mod, params = lower_to_tir(mod["main"], copy_constants())
+    cmms, encoded_constants, scratch_size = tir_to_cs_translator.translate(tir_mod, params)
+    return cmms, encoded_constants, scratch_size

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

@@ -221,3 +221,9 @@ def transform_module(
         mod = LegalizeSplit()(mod)
         mod = LegalizeEthosUConv2D()(mod)
         return mod
+
+    def __call__(self, *args, **kwargs):
+        # pylint is unable figure out the decorated
+        # class is callable, thus adding this to
+        # suppress the warning.
+        pass

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

@@ -197,3 +197,15 @@ def get_dim_value(layout: str, dim: int):
         if dim_char == dim:
             return idx
     return None
+
+
+def calculate_size_bytes(expr):
+    """This is a helper function to calculate the number
+    of bytes required to hold the tensor/relay.expr"""
+    try:
+        type_info = np.iinfo(expr.checked_type.dtype)
+    except ValueError:
+        type_info = np.finfo(expr.checked_type.dtype)
+    element_size = type_info.bits // 8
+    elements = np.prod(list(expr.checked_type.shape))
+    return element_size * elements

src/relay/backend/aot_executor_codegen.cc

@@ -665,7 +665,6 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     // Apply storage rewrite pass to the runner function to do memory planning
     auto storage_rewrite = tir::transform::StorageRewrite();
     mod_run = storage_rewrite(mod_run);
-
     // The workspace for main function should be calculated after performing storage_rewrite for
     // the top level TIR function.
     auto workspace_byte_alignment =