diff --git a/python/tvm/relay/op/contrib/ethosn.py b/python/tvm/relay/op/contrib/ethosn.py index 395974873142..9ebef0fac87e 100644 --- a/python/tvm/relay/op/contrib/ethosn.py +++ b/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), ] diff --git a/src/relay/backend/contrib/ethosn/codegen.cc b/src/relay/backend/contrib/ethosn/codegen.cc index a9fe0aab02ec..37f66cb51c44 100644 --- a/src/relay/backend/contrib/ethosn/codegen.cc +++ b/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()->body, ¶ms); tensor_table_[cn->args[0]] = {params.input_info}; + } else if (IsEthosnFunc(call, "ethos-n.qnn_tanh")) { + TanhParams params; + err += EthosnAPI::Tanh(cn->op.as()->body, ¶ms); + tensor_table_[cn->args[0]] = {params.input_info}; } else if (IsEthosnOp(call, "qnn.concatenate")) { ConcatenateParams params; err = EthosnAPI::Concatenate(call, ¶ms); @@ -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* 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* 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, ¶ms); + 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]; diff --git a/src/relay/backend/contrib/ethosn/codegen_ethosn.h b/src/relay/backend/contrib/ethosn/codegen_ethosn.h index 4a1be5e597e7..b3b93ffb8bb7 100644 --- a/src/relay/backend/contrib/ethosn/codegen_ethosn.h +++ b/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* out); EthosnError MakeSigmoidLayer(const Call& call, sl::TensorAndId* out); EthosnError MakeMeanLayer(const Call& call, sl::TensorAndId* out); + EthosnError MakeTanhLayer(const Call& call, sl::TensorAndId* out); EthosnError MakeConcatenateLayer(const Call& call, sl::TensorAndId* out); EthosnError MakeSplitLayer(const Call& call, sl::TensorsAndId* outs); EthosnError MakeDepthToSpaceLayer(const Call& call, sl::TensorAndId* out); diff --git a/src/relay/backend/contrib/ethosn/ethosn_api.cc b/src/relay/backend/contrib/ethosn/ethosn_api.cc index 88cdc790a60c..2ed94fbc136a 100644 --- a/src/relay/backend/contrib/ethosn/ethosn_api.cc +++ b/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(expr); + Call tanh = Downcast(quantize->args[0]); + Call dequantize = Downcast(tanh->args[0]); + // Create input info + const auto* input_dtype = quantize->checked_type().as(); + 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(expr); const auto& attrs = call->attrs.as(); diff --git a/src/relay/backend/contrib/ethosn/ethosn_api.h b/src/relay/backend/contrib/ethosn/ethosn_api.h index b790f2daf00e..2d49fb235568 100644 --- a/src/relay/backend/contrib/ethosn/ethosn_api.h +++ b/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 */ diff --git a/tests/python/contrib/test_ethosn/test_tanh.py b/tests/python/contrib/test_ethosn/test_tanh.py new file mode 100644 index 000000000000..f4053cfd91f0 --- /dev/null +++ b/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)