[ETHOSN] Implement tanh operator (apache#10486)

Adding compiler support for TANH operator, which is based on
an underlying pattern matching scheme.

One negative test is included as well.

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

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

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

@@ -165,6 +165,12 @@ def qnn_mean_pattern():
         )
         return pattern
 
+    def qnn_tanh_pattern():
+        pattern = is_op("qnn.dequantize")(wildcard(), is_constant(), is_constant())
+        pattern = is_op("tanh")(pattern)
+        pattern = is_op("qnn.quantize")(pattern, is_constant(), is_constant())
+        return pattern
+
     def check_conv2d(extract):
         """Check if a conv2d is supported by Ethos-N."""
         if not ethosn_available():
@@ -200,12 +206,20 @@ def check_sigmoid(extract):
 
         return support.sigmoid(extract)
 
+    def check_tanh(extract):
+        """Check if tanh is supported by Ethos-N."""
+        if not ethosn_available():
+            return False
+
+        return support.tanh(extract)
+
     return [
         ("ethos-n.qnn_conv2d", qnn_conv_pattern(), check_conv2d),
         ("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),
+        ("ethos-n.qnn_tanh", qnn_tanh_pattern(), check_tanh),
     ]
 
 

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

@@ -116,6 +116,10 @@ void InferTensorsVisitor::InferCall(const CallNode* cn) {
     MeanParams params;
     err += EthosnAPI::Mean(cn->op.as<FunctionNode>()->body, &params);
     tensor_table_[cn->args[0]] = {params.input_info};
+  } else if (IsEthosnFunc(call, "ethos-n.qnn_tanh")) {
+    TanhParams params;
+    err += EthosnAPI::Tanh(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);
@@ -283,6 +287,9 @@ sl::TensorsAndId ConstructNetworkVisitor::HandleCall(const CallNode* cn) {
   } else if (IsEthosnFunc(call, "ethos-n.qnn_mean")) {
     if ((err = MakeMeanLayer(call, &tensor))) ReportFatalError(call, err);
     return MakeOps(tensor);
+  } else if (IsEthosnFunc(call, "ethos-n.qnn_tanh")) {
+    if ((err = MakeTanhLayer(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);
@@ -473,6 +480,18 @@ EthosnError ConstructNetworkVisitor::MakeMeanLayer(const Call& call,
   return EthosnError();
 }
 
+EthosnError ConstructNetworkVisitor::MakeTanhLayer(const Call& call,
+                                                   sl::TensorAndId<sl::Operand>* out) {
+  auto input = operand_table_[call->args[0]][0];
+
+  try {
+    *out = AddTanh(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;
@@ -768,6 +787,15 @@ TVM_REGISTER_GLOBAL("relay.ethos-n.support.mean")
       err += EthosnError(reason);
     });
 
+TVM_REGISTER_GLOBAL("relay.ethos-n.support.tanh")
+    .set_body([](tvm::TVMArgs args, tvm::TVMRetValue* rv) {
+      Call call = args[0];
+      TanhParams params;
+      auto err = EthosnAPI::Tanh(call, &params);
+      err += EthosnCompiler::SupportedSetup();
+      *rv = !err && EthosnCompiler::GetSupported()->IsTanhSupported(params.input_info);
+    });
+
 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

@@ -206,6 +206,7 @@ class ConstructNetworkVisitor : public MixedModeVisitor, private ErrorReportingP
   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 MakeTanhLayer(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

@@ -411,6 +411,34 @@ EthosnError EthosnAPI::Mean(const Expr& expr, MeanParams* params) {
   return err;
 }
 
+EthosnError EthosnAPI::Tanh(const Expr& expr, TanhParams* params) {
+  Call quantize = Downcast<Call>(expr);
+  Call tanh = Downcast<Call>(quantize->args[0]);
+  Call dequantize = Downcast<Call>(tanh->args[0]);
+  // Create input info
+  const auto* input_dtype = quantize->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(dequantize->args[2], &input_zp);
+  err += AsConstant(dequantize->args[1], &input_sc);
+  float output_sc;
+  int output_zp;
+  err += AsConstant(quantize->args[2], &output_zp);
+  err += AsConstant(quantize->args[1], &output_sc);
+  auto test_zp = input_dtype->dtype.is_uint() ? 128 : 0;
+  if (output_zp != test_zp || output_sc != 0.0078125f) {
+    err += EthosnError(ErrStrm() << "output quantization params=(" << output_zp << ", " << output_sc
+                                 << "), must = (" << test_zp << ", 1/256)");
+  }
+  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);
   const auto& attrs = call->attrs.as<ConcatenateAttrs>();

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

@@ -96,6 +96,10 @@ struct MeanParams {
   sl::TensorInfo input_info;
 };
 
+struct TanhParams {
+  sl::TensorInfo input_info;
+};
+
 struct ConcatenateParams {
   sl::QuantizationInfo qInfo;
   sl::ConcatenationInfo concat_info = sl::ConcatenationInfo(1, qInfo);
@@ -198,6 +202,8 @@ class EthosnAPI {
   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 tanh params from a Relay an ethos-n tanh func */
+  static EthosnError Tanh(const Expr& expr, TanhParams* 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_tanh.py

@@ -0,0 +1,80 @@
+# 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 NPU integration tanh 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, input_zp, input_sc, output_zp, output_sc, dtype):
+    a = relay.var("a", shape=shape, dtype=dtype)
+    dequantize = relay.qnn.op.dequantize(
+        a,
+        input_scale=relay.const(input_sc, "float32"),
+        input_zero_point=relay.const(input_zp, "int32"),
+    )
+    tanh = relay.tanh(dequantize)
+    model = relay.qnn.op.quantize(
+        tanh,
+        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("shape", [(1, 512, 512, 3)])
+def test_tanh(shape):
+    np.random.seed(0)
+    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, 120, 0.0250629, 128, 0.0078125, "uint8")
+        mod = tei.make_module(model, [])
+        outputs.append(tei.build_and_run(mod, inputs, 1, {}, npu=npu))
+
+    tei.verify(outputs, "uint8", 1)
+
+
+@requires_ethosn
+@pytest.mark.parametrize(
+    "shape, input_zp, input_sc, output_zp, output_sc, dtype, err_msg",
+    [
+        (
+            (1, 16, 16, 16),
+            120,
+            0.0250629,
+            64,
+            0.0078125,
+            "uint8",
+            "output quantization params=(64, 0.0078125), must = (128, 1/256);",
+        )
+    ],
+)
+def test_tanh_failure(shape, input_zp, input_sc, output_zp, output_sc, dtype, err_msg):
+    model = _get_model(shape, input_zp, input_sc, output_zp, output_sc, dtype)
+    model = tei.make_ethosn_composite(model, "ethos-n.qnn_tanh")
+    mod = tei.make_ethosn_partition(model)
+    tei.test_error(mod, {}, err_msg)