[ETHOSN] Add support for mean on Ethos-N78 (#10130)

Adding the support of mean on Ethos-N78, which is based on
an underlying pattern matching scheme.
The operator is tested with 2 shapes: 4 and 3 dimensions.

Co-authored-by: Samuel Panijel <samuel.panijel@arm.com>

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

@@ -157,6 +157,14 @@ def qnn_sigmoid_pattern():
         pattern = is_op("qnn.quantize")(pattern, is_constant(), is_constant())
         return pattern
 
+    def qnn_mean_pattern():
+        pattern = is_op("cast")(wildcard())
+        pattern = is_op("mean")(pattern)
+        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():
@@ -178,6 +186,13 @@ def check_avg_pool2d(extract):
 
         return support.avg_pool2d(extract)
 
+    def check_mean(extract):
+        """Check if mean is supported by Ethos-N."""
+        if not ethosn_available():
+            return False
+
+        return support.mean(extract)
+
     def check_sigmoid(extract):
         """Check if a sigmoid is supported by Ethos-N."""
         if not ethosn_available():
@@ -190,6 +205,7 @@ def check_sigmoid(extract):
         ("ethos-n.qnn_avg_pool2d", qnn_avg_pool2d_pattern(), check_avg_pool2d),
         ("ethos-n.qnn_sigmoid", qnn_sigmoid_pattern(), check_sigmoid),
         ("ethos-n.qnn_fc", qnn_fc_pattern(), check_fc),
+        ("ethos-n.qnn_mean", qnn_mean_pattern(), check_mean),
     ]
 
 

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

@@ -112,6 +112,10 @@ void InferTensorsVisitor::InferCall(const CallNode* cn) {
     SigmoidParams params;
     err += EthosnAPI::Sigmoid(cn->op.as<FunctionNode>()->body, &params);
     tensor_table_[cn->args[0]] = {params.input_info};
+  } else if (IsEthosnFunc(call, "ethos-n.qnn_mean")) {
+    MeanParams params;
+    err += EthosnAPI::Mean(cn->op.as<FunctionNode>()->body, &params);
+    tensor_table_[cn->args[0]] = {params.input_info};
   } else if (IsEthosnOp(call, "qnn.concatenate")) {
     ConcatenateParams params;
     err = EthosnAPI::Concatenate(call, &params);
@@ -276,6 +280,9 @@ sl::TensorsAndId ConstructNetworkVisitor::HandleCall(const CallNode* cn) {
   } else if (IsEthosnFunc(call, "ethos-n.qnn_sigmoid")) {
     if ((err = MakeSigmoidLayer(call, &tensor))) ReportFatalError(call, err);
     return MakeOps(tensor);
+  } else if (IsEthosnFunc(call, "ethos-n.qnn_mean")) {
+    if ((err = MakeMeanLayer(call, &tensor))) ReportFatalError(call, err);
+    return MakeOps(tensor);
   } else if (IsEthosnOp(call, "qnn.concatenate")) {
     if ((err = MakeConcatenateLayer(call, &tensor))) ReportFatalError(call, err);
     return MakeOps(tensor);
@@ -454,6 +461,18 @@ EthosnError ConstructNetworkVisitor::MakeSigmoidLayer(const Call& call,
   return EthosnError();
 }
 
+EthosnError ConstructNetworkVisitor::MakeMeanLayer(const Call& call,
+                                                   sl::TensorAndId<sl::Operand>* out) {
+  auto input = operand_table_[call->args[0]][0];
+
+  try {
+    *out = AddMeanXy(network_, *input);
+  } catch (const sl::NotSupportedException& e) {
+    return EthosnError(e.what());
+  }
+  return EthosnError();
+}
+
 EthosnError ConstructNetworkVisitor::MakeConcatenateLayer(const Call& call,
                                                           sl::TensorAndId<sl::Operand>* out) {
   ConcatenateParams params;
@@ -735,6 +754,19 @@ TVM_REGISTER_GLOBAL("relay.ethos-n.support.sigmoid")
       err += EthosnError(reason);
     });
 
+TVM_REGISTER_GLOBAL("relay.ethos-n.support.mean")
+    .set_body([](tvm::TVMArgs args, tvm::TVMRetValue* rv) {
+      Call call = args[0];
+      MeanParams params;
+      auto err = EthosnAPI::Mean(call, &params);
+      err += EthosnCompiler::SupportedSetup();
+      char reason[kReasonMaxLength];
+      reason[0] = '\0';
+      *rv = !err && EthosnCompiler::GetSupported()->IsMeanXySupported(params.input_info, nullptr,
+                                                                      reason, sizeof(reason));
+      err += EthosnError(reason);
+    });
+
 TVM_REGISTER_GLOBAL("relay.ethos-n.support.concatenate")
     .set_body([](tvm::TVMArgs args, tvm::TVMRetValue* rv) {
       Call call = args[0];

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

@@ -205,6 +205,7 @@ class ConstructNetworkVisitor : public MixedModeVisitor, private ErrorReportingP
   EthosnError MakeReshapeLayer(const Call& call, sl::TensorAndId<sl::Operand>* out);
   EthosnError MakeAdditionLayer(const Call& call, sl::TensorAndId<sl::Operand>* out);
   EthosnError MakeSigmoidLayer(const Call& call, sl::TensorAndId<sl::Operand>* out);
+  EthosnError MakeMeanLayer(const Call& call, sl::TensorAndId<sl::Operand>* out);
   EthosnError MakeConcatenateLayer(const Call& call, sl::TensorAndId<sl::Operand>* out);
   EthosnError MakeSplitLayer(const Call& call, sl::TensorsAndId* outs);
   EthosnError MakeDepthToSpaceLayer(const Call& call, sl::TensorAndId<sl::Operand>* out);

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

@@ -382,6 +382,26 @@ EthosnError EthosnAPI::Sigmoid(const Expr& expr, SigmoidParams* params) {
                                       sl::QuantizationInfo(input_zp, input_sc));
   return err;
 }
+EthosnError EthosnAPI::Mean(const Expr& expr, MeanParams* params) {
+  Call requantize = Downcast<Call>(expr);
+  Call mean = Downcast<Call>(requantize->args[0]);
+  Call cast_0 = Downcast<Call>(mean->args[0]);
+
+  // Create input info
+  const auto* input_dtype = cast_0->args[0]->checked_type().as<TensorTypeNode>();
+  sl::TensorShape input_tensor_shape = {1, 1, 1, 1};
+  sl::DataType input_tensor_dtype;
+  EthosnError err = Tvm2Npu(input_dtype->shape, &input_tensor_shape);
+  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);
+  params->input_info = sl::TensorInfo(input_tensor_shape, input_tensor_dtype, sl::DataFormat::NHWC,
+                                      sl::QuantizationInfo(input_zp, input_sc));
+
+  return err;
+}
 
 EthosnError EthosnAPI::Concatenate(const Expr& expr, ConcatenateParams* params) {
   Call call = Downcast<Call>(expr);

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

@@ -92,6 +92,10 @@ struct SigmoidParams {
   sl::TensorInfo input_info;
 };
 
+struct MeanParams {
+  sl::TensorInfo input_info;
+};
+
 struct ConcatenateParams {
   sl::QuantizationInfo qInfo;
   sl::ConcatenationInfo concat_info = sl::ConcatenationInfo(1, qInfo);
@@ -189,6 +193,8 @@ class EthosnAPI {
   static EthosnError Addition(const Expr& expr, AdditionParams* params);
   /*! \brief Extract the Support Library sigmoid params from a Relay an ethos-n.qnn_sigmoid func */
   static EthosnError Sigmoid(const Expr& expr, SigmoidParams* params);
+  /*! \brief Extract the Support Library mean params from a mean func */
+  static EthosnError Mean(const Expr& expr, MeanParams* params);
   /*! \brief Extract the Support Library concatenate params from a Relay qnn.concatenate call */
   static EthosnError Concatenate(const Expr& expr, ConcatenateParams* params);
   /*! \brief Extract the Support Library split params from a Relay split call */

tests/python/contrib/test_ethosn/test_mean.py

@@ -0,0 +1,57 @@
+# 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 mean tests"""
+
+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, axis, keepdims, input_zp, input_sc, output_zp, output_sc, dtype):
+    a = relay.var("a", shape=shape, dtype=dtype)
+    casted = relay.op.cast(a, "int32")
+    mean = relay.mean(casted, axis, keepdims)
+    model = relay.qnn.op.requantize(
+        mean,
+        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
+def test_mean():
+    trials = [(1, 7, 7, 2048), (1, 8, 8)]
+
+    np.random.seed(0)
+    for shape in trials:
+        inputs = {
+            "a": tvm.nd.array(np.random.randint(0, high=255, size=shape, dtype="uint8")),
+        }
+        outputs = []
+        for npu in [False, True]:
+            model = _get_model(shape, [1, 2], True, 128, 0.0784314, 128, 0.0784314, "uint8")
+            mod = tei.make_module(model, [])
+            outputs.append(tei.build_and_run(mod, inputs, 1, {}, npu=npu))
+
+        tei.verify(outputs, "uint8", 1)