From f6e4297c8b22b0a7a849a69c5220f12a95cfc54f Mon Sep 17 00:00:00 2001 From: Samuel Panijel Date: Mon, 14 Feb 2022 14:50:09 +0200 Subject: [PATCH] [ETHOSN] Implement tanh operator Adding compiler support for TANH operator, which is based on an underlying pattern matching scheme. One negative test is included as well. --- python/tvm/relay/op/contrib/ethosn.py | 12 +++ src/relay/backend/contrib/ethosn/codegen.cc | 30 +++++++ .../backend/contrib/ethosn/codegen_ethosn.h | 1 + .../backend/contrib/ethosn/ethosn_api.cc | 28 +++++++ src/relay/backend/contrib/ethosn/ethosn_api.h | 6 ++ tests/python/contrib/test_ethosn/test_tanh.py | 81 +++++++++++++++++++ 6 files changed, 158 insertions(+) create mode 100644 tests/python/contrib/test_ethosn/test_tanh.py diff --git a/python/tvm/relay/op/contrib/ethosn.py b/python/tvm/relay/op/contrib/ethosn.py index 3959748731428..106e5dd81e605 100644 --- a/python/tvm/relay/op/contrib/ethosn.py +++ b/python/tvm/relay/op/contrib/ethosn.py @@ -163,6 +163,11 @@ def qnn_mean_pattern(): pattern = is_op("qnn.requantize")( pattern, is_constant(), is_constant(), is_constant(), is_constant() ) + + 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): @@ -200,12 +205,19 @@ def check_sigmoid(extract): return support.sigmoid(extract) + def check_tanh(extract): + """Check if tanh is supported by Ethos(TM)-N NPU.""" + 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 2de36b5434143..3c66bea1786e8 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; @@ -767,6 +786,17 @@ 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); + std::cout << "hello" << std::endl; + }); + 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 4a1be5e597e75..b3b93ffb8bb7b 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 0c716d0a86db7..e2dc1450d160f 100644 --- a/src/relay/backend/contrib/ethosn/ethosn_api.cc +++ b/src/relay/backend/contrib/ethosn/ethosn_api.cc @@ -403,6 +403,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 8f0508626d4b7..459fad249b6e2 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); @@ -195,6 +199,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 0000000000000..664e74bd8e019 --- /dev/null +++ b/tests/python/contrib/test_ethosn/test_tanh.py @@ -0,0 +1,81 @@ +# 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 +def test_tanh(): + trials = [(1, 512, 512, 3)] + + 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, 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 +def test_tanh_failure(): + trials = [ + ( + (1, 16, 16, 16), + 120, + 0.0250629, + 64, + 0.0078125, + "uint8", + "output quantization params=(64, 0.0078125), must = (0, 1/256);", + ), + ] + + for shape, input_zp, input_sc, output_zp, output_sc, dtype, err_msg in trials: + model = _get_model(shape, input_zp, input_sc, output_zp, output_sc, dtype) + model = tei.make_ethosn_composite(model, "ethos-n.qnn_sigmoid") + mod = tei.make_ethosn_partition(model) + tei.test_error(mod, {}, err_msg)