[ETHOSN] Add support for resize (apache#12535)

This commit adds support for the `resize` operator for
Arm(R) Ethos(TM)-N NPU.

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

@@ -176,6 +176,13 @@ def qnn_requantize_pattern():
         )
         return pattern
 
+    def qnn_resize_pattern():
+        pattern = is_op("image.resize2d")(wildcard()).has_attr({"method": "nearest_neighbor"})
+        pattern = is_op("qnn.requantize")(
+            pattern, is_constant(), is_constant(), is_constant(), is_constant()
+        )
+        return pattern
+
     def check_conv2d(extract):
         """Check if a conv2d is supported by Ethos-N."""
         if not ethosn_available():
@@ -232,6 +239,13 @@ def check_requantize(extract):
 
         return support.requantize(extract)
 
+    def check_resize(extract):
+        """Check if resize (nearest neighbor) is supported."""
+        if not ethosn_available():
+            return False
+
+        return support.resize(extract)
+
     return [
         ("ethos-n.qnn_conv2d", qnn_conv_pattern(), check_conv2d),
         ("ethos-n.qnn_avg_pool2d", qnn_avg_pool2d_pattern(), check_avg_pool2d),
@@ -240,6 +254,7 @@ def check_requantize(extract):
         ("ethos-n.qnn_mean", qnn_mean_pattern(), check_mean),
         ("ethos-n.qnn_tanh", qnn_tanh_pattern(), check_tanh),
         ("ethos-n.qnn_leaky_relu", qnn_leaky_relu_pattern(), check_leaky_relu),
+        ("ethos-n.qnn_resize", qnn_resize_pattern(), check_resize),
         ("ethos-n.qnn_requantize", qnn_requantize_pattern(), check_requantize),
     ]
 

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

@@ -148,6 +148,10 @@ void InferTensorsVisitor::InferCall(const CallNode* cn) {
     RequantizeParams params;
     err += EthosnAPI::Requantize(cn->op.as<FunctionNode>()->body, &params);
     tensor_table_[cn->args[0]] = {params.input_info};
+  } else if (IsEthosnFunc(call, "ethos-n.qnn_resize")) {
+    ResizeParams params;
+    err += EthosnAPI::Resize(cn->op.as<FunctionNode>()->body, &params);
+    tensor_table_[cn->args[0]] = {params.input_info};
   } else {
     err = EthosnError("unknown operator");
   }
@@ -322,6 +326,9 @@ sl::TensorsAndId ConstructNetworkVisitor::HandleCall(const CallNode* cn) {
   } else if (IsEthosnFunc(call, "ethos-n.qnn_requantize")) {
     if ((err = MakeRequantizeLayer(call, &tensor))) ReportFatalError(call, err);
     return MakeOps(tensor);
+  } else if (IsEthosnFunc(call, "ethos-n.qnn_resize")) {
+    if ((err = MakeResizeLayer(call, &tensor))) ReportFatalError(call, err);
+    return MakeOps(tensor);
   } else {
     ReportFatalError(call, EthosnError("unknown operator"));
     return {};
@@ -622,6 +629,24 @@ EthosnError ConstructNetworkVisitor::MakeRequantizeLayer(const Call& call,
   return EthosnError();
 }
 
+EthosnError ConstructNetworkVisitor::MakeResizeLayer(const Call& call,
+                                                     sl::TensorAndId<sl::Operand>* out) {
+  ResizeParams params;
+  params.input_info = GetTensorInfo(tensor_table_, call);
+  if (auto err = EthosnAPI::Resize(call->op.as<FunctionNode>()->body, &params)) {
+    return err;
+  }
+
+  auto input = operand_table_[call->args[0]][0];
+
+  try {
+    *out = AddResize(network_, *input, params.resize_info);
+  } catch (const sl::NotSupportedException& e) {
+    return EthosnError(e.what());
+  }
+  return EthosnError();
+}
+
 runtime::Module EthosnCompiler::CreateRuntimeModule(const ObjectRef& ref) {
   std::vector<runtime::ethosn::OrderedCompiledNetwork> cmms;
   if (ref->IsInstance<FunctionNode>()) {
@@ -958,6 +983,20 @@ TVM_REGISTER_GLOBAL("relay.ethos-n.support.requantize")
       err += EthosnError(reason);
     });
 
+TVM_REGISTER_GLOBAL("relay.ethos-n.support.resize")
+    .set_body([](tvm::TVMArgs args, tvm::TVMRetValue* rv) {
+      Call call = args[0];
+      ResizeParams params;
+      auto err = EthosnAPI::Resize(call, &params);
+      err += EthosnCompiler::SupportedSetup();
+      char reason[kReasonMaxLength];
+      reason[0] = '\0';
+      *rv = !err &&
+            EthosnCompiler::GetSupported()->IsResizeSupported(
+                params.resize_info, params.input_info, &params.output_info, reason, sizeof(reason));
+      err += EthosnError(reason);
+    });
+
 TVM_REGISTER_GLOBAL("relay.ethos-n.query").set_body([](tvm::TVMArgs args, tvm::TVMRetValue* rv) {
 #if defined ETHOSN_HW
   *rv = true;

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

@@ -212,6 +212,7 @@ class ConstructNetworkVisitor : public MixedModeVisitor, private ErrorReportingP
   EthosnError MakeReluLayer(const Call& call, sl::TensorAndId<sl::Operand>* out);
   EthosnError MakeLeakyReLULayer(const Call& call, sl::TensorAndId<sl::Operand>* out);
   EthosnError MakeRequantizeLayer(const Call& call, sl::TensorAndId<sl::Operand>* out);
+  EthosnError MakeResizeLayer(const Call& call, sl::TensorAndId<sl::Operand>* out);
 
   /*! \brief A look-up table from Expr to layers. */
   std::map<Expr, std::vector<std::shared_ptr<sl::Operand>>> operand_table_;

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

@@ -23,6 +23,7 @@
 
 #include "ethosn_api.h"
 
+#include <tvm/relay/attrs/image.h>
 #include <tvm/relay/attrs/nn.h>
 #include <tvm/relay/expr.h>
 #include <tvm/relay/expr_functor.h>
@@ -684,6 +685,45 @@ EthosnError EthosnAPI::Requantize(const Expr& expr, RequantizeParams* params) {
   return err;
 }
 
+EthosnError EthosnAPI::Resize(const Expr& expr, ResizeParams* params) {
+  Call requantize = Downcast<Call>(expr);
+  Call resize = Downcast<Call>(requantize->args[0]);
+
+  const auto* input_dtype = resize->args[0]->checked_type().as<TensorTypeNode>();
+  sl::TensorShape input_tensor_shape = {1, 1, 1, 1};
+  EthosnError err = Tvm2Npu(input_dtype->shape, &input_tensor_shape);
+  sl::DataType input_tensor_dtype;
+  err += Tvm2Npu(input_dtype->dtype, &input_tensor_dtype);
+  float input_sc;
+  int input_zp;
+  err += AsConstant(requantize->args[2], &input_zp);
+  err += AsConstant(requantize->args[1], &input_sc);
+  sl::QuantizationInfo input_q_info;
+  err += Tvm2Npu(input_zp, input_sc, &input_q_info);
+  params->input_info =
+      sl::TensorInfo(input_tensor_shape, input_tensor_dtype, sl::DataFormat::NHWC, input_q_info);
+
+  float output_sc;
+  int output_zp;
+  err += AsConstant(requantize->args[3], &output_sc);
+  err += AsConstant(requantize->args[4], &output_zp);
+  sl::QuantizationInfo resize_q_info;
+  err += Tvm2Npu(output_zp, output_sc, &resize_q_info);
+  const auto* attrs = resize->attrs.as<Resize2DAttrs>();
+  uint32_t height, width;
+  err += Tvm2Npu(attrs->size, &height, &width);
+  params->resize_info =
+      sl::ResizeInfo{sl::ResizeAlgorithm::NEAREST_NEIGHBOUR, height, width, resize_q_info};
+
+  sl::TensorInfo output_info = params->input_info;
+  output_info.m_Dimensions[1] = params->resize_info.m_NewHeight;
+  output_info.m_Dimensions[2] = params->resize_info.m_NewWidth;
+  output_info.m_QuantizationInfo = params->resize_info.m_OutputQuantizationInfo;
+  params->output_info = output_info;
+
+  return err;
+}
+
 EthosnError EthosnAPI::Tvm2Npu(const Array<IndexExpr>& padding, sl::Padding* npu_padding) {
   std::array<uint32_t, 4> dim;
   if (EthosnError err = AsArray<IndexExpr, uint32_t>(padding, &dim)) {

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

@@ -146,6 +146,12 @@ struct RequantizeParams {
   sl::TensorInfo output_info;
 };
 
+struct ResizeParams {
+  sl::ResizeInfo resize_info;
+  sl::TensorInfo input_info;
+  sl::TensorInfo output_info;
+};
+
 /*!
  * \brief A wrapper around std::stringstream to build an EthosnError.
  */
@@ -241,6 +247,8 @@ class EthosnAPI {
   static EthosnError Relu(const Expr& expr, ReluParams* params);
   /*! \brief Extract the Support Library requantize params from a Relay qnn.requantize call */
   static EthosnError Requantize(const Expr& expr, RequantizeParams* params);
+  /*! \brief Extract the Support Library resize params from a Relay resize call */
+  static EthosnError Resize(const Expr& expr, ResizeParams* params);
 
  private:
   /*! \brief Convert a TVM IndexExpr array to a SL tensor shape */

tests/python/contrib/test_ethosn/test_resize.py

@@ -0,0 +1,134 @@
+# 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.
+
+"""Arm(R) Ethos(TM)-N integration resize tests"""
+
+import pytest
+import numpy as np
+import tvm
+from tvm import relay
+from tvm.testing import requires_ethosn
+from . import infrastructure as tei
+
+
+def _get_model(
+    shape,
+    dtype,
+    size,
+    input_zp,
+    input_sc,
+    output_zp,
+    output_sc,
+    coordinate_transformation_mode,
+    rounding_method,
+):
+    x = relay.var("x", shape=shape, dtype=dtype)
+    resize = relay.image.resize2d(
+        data=x,
+        size=size,
+        layout="NHWC",
+        method="nearest_neighbor",
+        coordinate_transformation_mode=coordinate_transformation_mode,
+        rounding_method=rounding_method,
+    )
+    model = relay.qnn.op.requantize(
+        resize,
+        input_scale=relay.const(input_sc, "float32"),
+        input_zero_point=relay.const(input_zp, "int32"),
+        output_scale=relay.const(output_sc, "float32"),
+        output_zero_point=relay.const(output_zp, "int32"),
+        out_dtype=dtype,
+    )
+    return model
+
+
+@requires_ethosn
+@pytest.mark.parametrize("dtype", ["uint8", "int8"])
+@pytest.mark.parametrize(
+    "shape, size, coordinate_transformation_mode, rounding_method",
+    [
+        ((1, 4, 4, 2), (8, 8), "half_pixel", "round_prefer_ceil"),
+        ((1, 4, 4, 2), (7, 7), "asymmetric", "floor"),
+        ((1, 4, 8, 3), (8, 16), "half_pixel", "round_prefer_ceil"),
+        ((1, 4, 8, 3), (7, 15), "asymmetric", "floor"),
+    ],
+)
+def test_resize(dtype, shape, size, coordinate_transformation_mode, rounding_method):
+    np.random.seed(0)
+    zp_min = np.iinfo(dtype).min
+    zp_max = np.iinfo(dtype).max
+    inputs = {
+        "x": tvm.nd.array(np.random.randint(zp_min, high=zp_max + 1, size=shape, dtype=dtype)),
+    }
+    outputs = []
+    for npu in [False, True]:
+        model = _get_model(
+            shape=shape,
+            dtype=dtype,
+            size=size,
+            input_zp=zp_min + 128,
+            input_sc=0.0784314,
+            output_zp=zp_min + 128,
+            output_sc=0.0784314,
+            coordinate_transformation_mode=coordinate_transformation_mode,
+            rounding_method=rounding_method,
+        )
+        mod = tei.make_module(model, {})
+        x = tei.build_and_run(mod, inputs, 1, {}, npu=npu)
+        outputs.append(x)
+
+    tei.verify(outputs, dtype, 1)
+
+
+@requires_ethosn
+def test_resize_failure():
+    trials = [
+        (
+            (30, 20),
+            "Requested height isn't supported",
+        ),
+        (
+            (20, 30),
+            "Requested width isn't supported",
+        ),
+        (
+            (19, 20),
+            "Requested width and height must be both even or both odd",
+        ),
+        (
+            (20, 19),
+            "Requested width and height must be both even or both odd",
+        ),
+    ]
+    dtype = "int8"
+    zp_min = np.iinfo(dtype).min
+
+    for size, err_msg in trials:
+        model = _get_model(
+            shape=(1, 10, 10, 1),
+            dtype=dtype,
+            size=size,
+            input_zp=zp_min + 128,
+            input_sc=0.0784314,
+            output_zp=zp_min + 128,
+            output_sc=0.0784314,
+            coordinate_transformation_mode="half_pixel",
+            rounding_method="round_prefer_ceil",
+        )
+        model = tei.make_ethosn_composite(model, "ethos-n.qnn_resize")
+        mod = tei.make_ethosn_partition(model)
+        tei.test_error(mod, {}, err_msg)