diff --git a/BUILD.md b/BUILD.md index a93b0fbbdd..54b0f013c0 100644 --- a/BUILD.md +++ b/BUILD.md @@ -104,6 +104,7 @@ The complete list of build options can be found by running `./build.sh (or .\bui * [Nuphar Model Compiler](#Nuphar) * [DirectML](#DirectML) * [ARM Compute Library](#ARM-Compute-Library) +* [Rockchip RKNPU](#RKNPU) **Options** * [OpenMP](#OpenMP) @@ -449,6 +450,35 @@ export LD_LIBRARY_PATH=~/ComputeLibrary/build/ --- +### RKNPU +See more information on the RKNPU Execution Provider [here](./docs/execution_providers/RKNPU-ExecutionProvider.md). + +#### Pre-Requisites + +* Supported platform: RK1808 Linux +* See [Build ARM](#ARM) below for information on building for ARM devices +* Use gcc-linaro-6.3.1-2017.05-x86_64_aarch64-linux-gnu instead of gcc-linaro-6.3.1-2017.05-x86_64_arm-linux-gnueabihf, and modify CMAKE_CXX_COMPILER & CMAKE_C_COMPILER in tool.cmake: + ``` + set(CMAKE_CXX_COMPILER aarch64-linux-gnu-g++) + set(CMAKE_C_COMPILER aarch64-linux-gnu-gcc) + ``` + +#### Build Instructions +##### Linux +1. Download [rknpu_ddk](#https://github.com/airockchip/rknpu_ddk.git) to any directory. + +2. Build ONNX Runtime library and test: + ``` + ./build.sh --arm --use_rknpu --parallel --build_shared_lib --build_dir build_arm --config MinSizeRel --cmake_extra_defines RKNPU_DDK_PATH= CMAKE_TOOLCHAIN_FILE= ONNX_CUSTOM_PROTOC_EXECUTABLE= + ``` +3. Deploy ONNX runtime and librknpu_ddk.so on the RK1808 board: + ``` + libonnxruntime.so.1.2.0 + onnxruntime_test_all + rknpu_ddk/lib64/librknpu_ddk.so + ``` +--- + ## Options ### OpenMP #### Build Instructions diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 319371ecd7..776964f232 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -56,6 +56,7 @@ option(onnxruntime_USE_CUDA "Build with CUDA support" OFF) option(onnxruntime_USE_OPENVINO "Build with OpenVINO support" OFF) option(onnxruntime_USE_EIGEN_FOR_BLAS "Use eign for blas" ON) option(onnxruntime_USE_NNAPI "Build with DNNLibrary for Android NNAPI support" OFF) +option(onnxruntime_USE_RKNPU "Build with RKNPU support" OFF) option(onnxruntime_USE_DNNL "Build with DNNL support" OFF) option(onnxruntime_USE_MKLML "Build the default cpu provider with MKL-ML binary dependency" OFF) option(onnxruntime_USE_FEATURIZERS "Build ML Featurizers support" OFF) diff --git a/cmake/onnxruntime.cmake b/cmake/onnxruntime.cmake index 0558a20a59..9a79115edd 100644 --- a/cmake/onnxruntime.cmake +++ b/cmake/onnxruntime.cmake @@ -75,6 +75,7 @@ target_link_libraries(onnxruntime PRIVATE ${PROVIDERS_DNNL} ${PROVIDERS_NGRAPH} ${PROVIDERS_NNAPI} + ${PROVIDERS_RKNPU} ${PROVIDERS_TENSORRT} ${PROVIDERS_OPENVINO} ${PROVIDERS_NUPHAR} diff --git a/cmake/onnxruntime_providers.cmake b/cmake/onnxruntime_providers.cmake index 8b1fa68383..5d8d42a16b 100644 --- a/cmake/onnxruntime_providers.cmake +++ b/cmake/onnxruntime_providers.cmake @@ -55,6 +55,10 @@ if(onnxruntime_USE_NNAPI) set(PROVIDERS_NNAPI onnxruntime_providers_nnapi) list(APPEND ONNXRUNTIME_PROVIDER_NAMES nnapi) endif() +if(onnxruntime_USE_RKNPU) + set(PROVIDERS_RKNPU onnxruntime_providers_rknpu) + list(APPEND ONNXRUNTIME_PROVIDER_NAMES rknpu) +endif() if(onnxruntime_USE_DML) set(PROVIDERS_DML onnxruntime_providers_dml) list(APPEND ONNXRUNTIME_PROVIDER_NAMES dml) @@ -473,6 +477,40 @@ if (onnxruntime_USE_NNAPI) set_target_properties(onnxruntime_providers_nnapi PROPERTIES LINKER_LANGUAGE CXX) endif() +if (onnxruntime_USE_RKNPU) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-variable -Wno-unused-parameter") + add_definitions(-DUSE_RKNPU=1) + option(DNN_READ_ONNX "" ON) + set(DNN_CUSTOM_PROTOC_EXECUTABLE ${ONNX_CUSTOM_PROTOC_EXECUTABLE}) + option(DNN_CMAKE_INSTALL "" OFF) + option(DNN_BUILD_BIN "" OFF) + if (NOT RKNPU_DDK_PATH) + message(FATAL_ERROR "RKNPU_DDK_PATH required for onnxruntime_USE_RKNPU") + endif() + set(RKNPU_DDK_INCLUDE_DIR ${RKNPU_DDK_PATH}/include) + if (CMAKE_SIZEOF_VOID_P EQUAL 8) + set(RKNPU_DDK_LIB_DIR ${RKNPU_DDK_PATH}/lib64) + else() + set(RKNPU_DDK_LIB_DIR ${RKNPU_DDK_PATH}/lib) + endif() + file(GLOB_RECURSE + onnxruntime_providers_rknpu_cc_srcs CONFIGURE_DEPENDS + "${ONNXRUNTIME_ROOT}/core/providers/rknpu/*.h" + "${ONNXRUNTIME_ROOT}/core/providers/rknpu/*.cc" + ) + source_group(TREE ${ONNXRUNTIME_ROOT}/core FILES ${onnxruntime_providers_rknpu_cc_srcs}) + add_library(onnxruntime_providers_rknpu ${onnxruntime_providers_rknpu_cc_srcs}) + onnxruntime_add_include_to_target(onnxruntime_providers_rknpu onnxruntime_common onnxruntime_framework onnx onnx_proto protobuf::libprotobuf-lite) + target_link_libraries(onnxruntime_providers_rknpu PRIVATE -lrknpu_ddk) + add_dependencies(onnxruntime_providers_rknpu onnx ${onnxruntime_EXTERNAL_DEPENDENCIES}) + set_target_properties(onnxruntime_providers_rknpu PROPERTIES CXX_STANDARD 14) + set_target_properties(onnxruntime_providers_rknpu PROPERTIES CXX_STANDARD_REQUIRED ON) + set_target_properties(onnxruntime_providers_rknpu PROPERTIES FOLDER "ONNXRuntime") + target_include_directories(onnxruntime_providers_rknpu PRIVATE ${ONNXRUNTIME_ROOT} ${rknpu_INCLUDE_DIRS} ${RKNPU_DDK_INCLUDE_DIR}) + link_directories(onnxruntime_providers_rknpu ${RKNPU_DDK_LIB_DIR}) + set_target_properties(onnxruntime_providers_rknpu PROPERTIES LINKER_LANGUAGE CXX) +endif() + if (onnxruntime_USE_DML) file(GLOB_RECURSE onnxruntime_providers_dml_cc_srcs CONFIGURE_DEPENDS "${ONNXRUNTIME_ROOT}/core/providers/dml/*.h" diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 9701a2e1d4..cb4f4fc8d5 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -89,6 +89,7 @@ set(onnxruntime_pybind11_state_libs ${PROVIDERS_OPENVINO} ${PROVIDERS_NUPHAR} ${PROVIDERS_NNAPI} + ${PROVIDERS_RKNPU} ${PROVIDERS_DML} onnxruntime_optimizer onnxruntime_providers diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 6e9c7d6c12..cf3d9e3e70 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -211,6 +211,13 @@ if (onnxruntime_USE_NNAPI) list(APPEND onnxruntime_test_providers_src ${onnxruntime_test_providers_nnapi_src}) endif() +if (onnxruntime_USE_RKNPU) + file(GLOB_RECURSE onnxruntime_test_providers_rknpu_src CONFIGURE_DEPENDS + "${TEST_SRC_DIR}/providers/rknpu/*" + ) + list(APPEND onnxruntime_test_providers_src ${onnxruntime_test_providers_rknpu_src}) +endif() + set (ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR "${ONNXRUNTIME_ROOT}/test/shared_lib") set (ONNXRUNTIME_GLOBAL_THREAD_POOLS_TEST_SRC_DIR "${ONNXRUNTIME_ROOT}/test/global_thread_pools") @@ -293,6 +300,10 @@ if(onnxruntime_USE_NNAPI) list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_nnapi) endif() +if(onnxruntime_USE_RKNPU) + list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_rknpu) +endif() + if(onnxruntime_USE_FEATURIZERS) list(APPEND onnxruntime_test_providers_dependencies onnxruntime_featurizers) list(APPEND onnxruntime_test_providers_libs onnxruntime_featurizers re2) @@ -334,6 +345,7 @@ set(ONNXRUNTIME_TEST_LIBS ${PROVIDERS_OPENVINO} ${PROVIDERS_NUPHAR} ${PROVIDERS_NNAPI} + ${PROVIDERS_RKNPU} ${PROVIDERS_DML} ${PROVIDERS_ACL} onnxruntime_optimizer @@ -370,6 +382,13 @@ if(onnxruntime_USE_NNAPI) list(APPEND onnxruntime_test_providers_libs onnxruntime_providers_nnapi) endif() +if(onnxruntime_USE_RKNPU) + list(APPEND onnxruntime_test_framework_src_patterns ${TEST_SRC_DIR}/providers/rknpu/*) + list(APPEND onnxruntime_test_framework_libs onnxruntime_providers_rknpu) + list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_rknpu) + list(APPEND onnxruntime_test_providers_libs onnxruntime_providers_rknpu) +endif() + if(WIN32) if (onnxruntime_USE_TVM) list(APPEND disabled_warnings ${DISABLED_WARNINGS_FOR_TVM}) diff --git a/docs/execution_providers/RKNPU-ExecutionProvider.md b/docs/execution_providers/RKNPU-ExecutionProvider.md new file mode 100644 index 0000000000..7079559ad7 --- /dev/null +++ b/docs/execution_providers/RKNPU-ExecutionProvider.md @@ -0,0 +1,76 @@ +# RKNPU Execution Provider +RKNPU DDK is an advanced interface to access Rockchip NPU, currently support platform as follow: +- RK1808 Linux + + Note: RK3399Pro platform is not supported. + +RKNPU Execution Provider enables deep learning inference on Rockchip NPU via RKNPU DDK. + + +## Build +For build instructions, please see the [BUILD page](../../BUILD.md#RKNPU). + +## Using +### C/C++ +To use RKNPU as execution provider for inferencing, please register it as below. +``` +string log_id = "Foo"; +auto logging_manager = std::make_unique +(std::unique_ptr{new CLogSink{}}, + static_cast(lm_info.default_warning_level), + false, + LoggingManager::InstanceType::Default, + &log_id) +Environment::Create(std::move(logging_manager), env) +InferenceSession session_object{so,env}; +session_object.RegisterExecutionProvider(std::make_unique<::onnxruntime::RknpuExecutionProvider>()); +status = session_object.Load(model_file_name); +``` +The C API details are [here](../C_API.md#c-api). + + +## ONNX Op supported using RKNPU + +The table below shows the ONNX Ops supported using RKNPU Execution Provider and the mapping between ONNX Ops and RKNPU Ops. + +| **ONNX Ops** | **RKNPU Ops** | +| --- | --- | +| Add | ADD | +| Mul | MULTIPLY | +| Conv | CONV2D | +| QLinearConv | CONV2D | +| Gemm | FULLCONNECT | +| Softmax | SOFTMAX | +| AveragePool | POOL | +| GlobalAveragePool | POOL | +| MaxPool | POOL | +| GlobalMaxPool | POOL | +| LeakyRelu | LEAKY_RELU | +| Concat | CONCAT | +| BatchNormalization | BATCH_NORM | +| Reshape | RESHAPE | +| Flatten | RESHAPE | +| Squeeze | RESHAPE | +| Unsqueeze | RESHAPE | +| Transpose | PERMUTE | +| Relu | RELU | +| Sub | SUBTRACT | +| Clip(0~6)| RELU6 | +| DequantizeLinear | DATACONVERT | +| Clip | CLIP | + + +## Model Supported + +Below Models are supported from ONNX open model zoo using RKNPU Execution Provider + +### Image Classification +- squeezenet +- mobilenetv2-1.0 +- resnet50v1 +- resnet50v2 +- inception_v2 + +### Object Detection +- ssd +- yolov3 \ No newline at end of file diff --git a/include/onnxruntime/core/graph/constants.h b/include/onnxruntime/core/graph/constants.h index 6344c50774..ca4d7ee5ed 100644 --- a/include/onnxruntime/core/graph/constants.h +++ b/include/onnxruntime/core/graph/constants.h @@ -30,6 +30,7 @@ constexpr const char* kOpenVINOExecutionProvider = "OpenVINOExecutionProvider"; constexpr const char* kNupharExecutionProvider = "NupharExecutionProvider"; constexpr const char* kTensorrtExecutionProvider = "TensorrtExecutionProvider"; constexpr const char* kNnapiExecutionProvider = "NnapiExecutionProvider"; +constexpr const char* kRknpuExecutionProvider = "RknpuExecutionProvider"; constexpr const char* kDmlExecutionProvider = "DmlExecutionProvider"; constexpr const char* kAclExecutionProvider = "ACLExecutionProvider"; } // namespace onnxruntime diff --git a/include/onnxruntime/core/providers/rknpu/rknpu_provider_factory.h b/include/onnxruntime/core/providers/rknpu/rknpu_provider_factory.h new file mode 100644 index 0000000000..832171e09d --- /dev/null +++ b/include/onnxruntime/core/providers/rknpu/rknpu_provider_factory.h @@ -0,0 +1,14 @@ +// Copyright 2020 rock-chips.com Inc. + +#include "onnxruntime_c_api.h" + +#ifdef __cplusplus +extern "C" { +#endif + +ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_Rknpu, + _In_ OrtSessionOptions* options); + +#ifdef __cplusplus +} +#endif diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index f054d1b414..ac064e9e6d 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -100,7 +100,8 @@ bool ProviderIsCpuBased(const std::string& provider_type) { provider_type == onnxruntime::kNGraphExecutionProvider || provider_type == onnxruntime::kNupharExecutionProvider || provider_type == onnxruntime::kOpenVINOExecutionProvider || - provider_type == onnxruntime::kNnapiExecutionProvider; + provider_type == onnxruntime::kNnapiExecutionProvider || + provider_type == onnxruntime::kRknpuExecutionProvider; } common::Status AllocateHelper(const IExecutionProvider& execution_provider, const OrtDevice& device, diff --git a/onnxruntime/core/providers/rknpu/node_attr_helper.cc b/onnxruntime/core/providers/rknpu/node_attr_helper.cc new file mode 100644 index 0000000000..5796f0a5cf --- /dev/null +++ b/onnxruntime/core/providers/rknpu/node_attr_helper.cc @@ -0,0 +1,112 @@ +// +// Created by daquexian on 8/3/18. +// + +#include "node_attr_helper.h" + +#include +#include +#include + +namespace onnxruntime { +namespace rknpu { + +NodeAttrHelper::NodeAttrHelper(ONNX_NAMESPACE::NodeProto proto) + : node_(proto) { +} + +float NodeAttrHelper::get(const std::string& key, + const float def_val) const { + for (const auto& attr : node_.attribute()) { + if (attr.name() == key) { + return attr.f(); + } + } + + return def_val; +} + +int NodeAttrHelper::get(const std::string& key, + const int def_val) const { + for (const auto& attr : node_.attribute()) { + if (attr.name() == key) { + return static_cast(attr.i()); + } + } + + return def_val; +} + +std::string NodeAttrHelper::get(const std::string& key, + const std::string& def_val) const { + for (const auto& attr : node_.attribute()) { + if (attr.name() == key) { + return attr.s(); + } + } + + return def_val; +} + +std::vector NodeAttrHelper::get(const std::string& key, + const std::vector& def_val) const { + if (!has_attr(key)) { + return def_val; + } + std::vector v; + + for (const auto& attr : node_.attribute()) { + if (attr.name() == key) { + v.reserve(static_cast(attr.ints_size())); + for (int j = 0; j < attr.ints_size(); j++) { + v.push_back(static_cast(attr.ints(j))); + } + + break; + } + } + + if (v.empty()) { + return def_val; + } + + return v; +} + +std::vector NodeAttrHelper::get(const std::string& key, + const std::vector& def_val) const { + if (!has_attr(key)) { + return def_val; + } + std::vector v; + + for (const auto& attr : node_.attribute()) { + if (attr.name() == key) { + v.reserve(static_cast(attr.floats_size())); + for (int j = 0; j < attr.floats_size(); j++) { + v.push_back(attr.floats(j)); + } + + break; + } + } + + if (v.empty()) { + return def_val; + } + + return v; +} + +bool NodeAttrHelper::has_attr(const std::string& key) const { + for (const auto& attr : node_.attribute()) { + if (attr.name() == key) { + return true; + } + } + + return false; +} + +} // namespace rknpu +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/rknpu/node_attr_helper.h b/onnxruntime/core/providers/rknpu/node_attr_helper.h new file mode 100644 index 0000000000..6ab8f8c6bb --- /dev/null +++ b/onnxruntime/core/providers/rknpu/node_attr_helper.h @@ -0,0 +1,39 @@ +// +// Created by daquexian on 8/3/18. +// + +#pragma once + +#include +#include +#include + +namespace onnxruntime { +namespace rknpu { + +/** + * Wrapping onnx::NodeProto for retrieving attribute values + */ +class NodeAttrHelper { + public: + explicit NodeAttrHelper(ONNX_NAMESPACE::NodeProto proto); + + float get(const std::string& key, + const float def_val) const; + int get(const std::string& key, + const int def_val) const; + std::vector get(const std::string& key, + const std::vector& def_val) const; + std::vector get(const std::string& key, + const std::vector& def_val) const; + std::string get(const std::string& key, + const std::string& def_val) const; + + bool has_attr(const std::string& key) const; + + private: + ONNX_NAMESPACE::NodeProto node_; +}; + +} // namespace rknpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rknpu/onnx_converter.cc b/onnxruntime/core/providers/rknpu/onnx_converter.cc new file mode 100644 index 0000000000..30de1af8fa --- /dev/null +++ b/onnxruntime/core/providers/rknpu/onnx_converter.cc @@ -0,0 +1,1898 @@ +// Copyright 2020 rock-chips.com Inc. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "core/common/logging/logging.h" +#include "onnx_converter.h" +#include "node_attr_helper.h" + +using std::string; +using std::vector; + +namespace onnxruntime { +namespace rknpu { + +#define HAS(map, key) \ + (map.find(key) != map.end()) + +#define ADD_SHAPE(name) \ + if (HAS(rk_tensors_, name)) { \ + const auto& tensor = rk_tensors_.at(name); \ + shaper_.AddShape(name, tensor->GetDims()); \ + } + +#define GET_ATTR(attr, name, type) \ + std::vector attr; \ + if (HAS(rk_tensors_, name)) { \ + const auto& tensor = rk_tensors_.at(name); \ + const void* data = tensor->GetData(); \ + rk::nn::PrecisionType fmt = tensor->GetPrecision(); \ + uint32_t dim = 1; \ + if (tensor->GetDims().size() > 0) \ + dim = tensor->GetDims()[0]; \ + for (uint32_t i = 0; i < dim; i++) { \ + if (fmt == rk::nn::PrecisionType::UINT8 || \ + fmt == rk::nn::PrecisionType::INT8) { \ + attr.push_back(((type*)data)[i]); \ + } else if (fmt == rk::nn::PrecisionType::INT32 || \ + fmt == rk::nn::PrecisionType::UINT32) { \ + attr.push_back(((type*)data)[i]); \ + } else if (fmt == rk::nn::PrecisionType::INT64 || \ + fmt == rk::nn::PrecisionType::UINT64) { \ + int64_t val = ((int64_t*)data)[i]; \ + if (val > 0x7fffffff) val = 0x7fffffff; \ + attr.push_back((type)val); \ + } else if (fmt == rk::nn::PrecisionType::FLOAT32) { \ + attr.push_back((type)((float*)data)[i]); \ + } else { \ + LOGS_DEFAULT(FATAL) << "the format of " << name \ + << " is not support!"; \ + } \ + } \ + } + +std::string OnnxConverter::m(const std::string& str) const { + if (name_map_.find(str) != name_map_.end()) { + return name_map_.at(str); + } + return str; +} + +std::pair, OnnxConverter::FuseCode> +OnnxConverter::FindActivation(const ONNX_NAMESPACE::ModelProto& model_proto, + const std::string& output) { + std::pair, FuseCode> + activation{{}, FuseCode::FUSED_NONE}; + int i = 0; + for (const auto& _node : model_proto.graph().node()) { + if (!_node.input().empty() && output == _node.input(0) && + _node.op_type() == "Relu") { + // If there are two branches after a conv/pool and both branches has + // a relu on the top, we have to add two normal relu layers + if (activation.second != FuseCode::FUSED_NONE) { + return {{}, FuseCode::FUSED_NONE}; + } + const auto node_pair = std::make_pair(i, _node); + activation = std::make_pair(node_pair, FuseCode::FUSED_RELU); + } + i++; + } + if (activation.first.first != 0) { + skipped_act_.push_back(activation.first.first); + name_map_[activation.first.second.output(0)] = output; + } + return activation; +} + +std::shared_ptr +OnnxConverter::CreateRknnTensor(const std::string& name, + const std::vector& dims, + const void* data, + const rk::nn::TensorRole role, + const rk::nn::PrecisionType precision, + const rk::nn::DataLayoutType layout, + const rk::nn::QuantizationType qntType, + const uint8_t bits, + const float scale, + const uint32_t zero_point, + const int8_t fl) { + auto attr = std::make_shared(); + attr->name = name; + attr->dims = dims; + attr->precision = precision; + attr->layout = layout; + attr->qntType = qntType; + attr->role = role; + attr->qntBits = bits; + attr->qntParamDFP.fl.push_back(fl); + attr->qntParamAffineAsymmetric.zero_point.push_back(zero_point); + attr->qntParamAffineAsymmetric.scale.push_back(scale); + attr->qntParamSymmetric.scale.push_back(scale); + return graph_->CreateTensor(attr, (void*)data); +} + +void OnnxConverter::HandleInitializer() { + for (const auto& tensor : model_proto_.graph().initializer()) { + const std::string name = tensor.name(); + std::vector dims; + for (const auto dim : tensor.dims()) { + dims.push_back(static_cast(dim)); + } + if (tensor.data_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { + const char* ptr = tensor.float_data().empty() + ? tensor.raw_data().data() + : reinterpret_cast( + tensor.float_data().data()); + rk_tensors_[name] = CreateRknnTensor(name, dims, ptr, + rk::nn::TensorRole::CONST, + rk::nn::PrecisionType::FLOAT32); + } else if (tensor.data_type() == + ONNX_NAMESPACE::TensorProto_DataType_UINT8) { + const char* ptr = tensor.int32_data().empty() + ? tensor.raw_data().data() + : reinterpret_cast( + tensor.int32_data().data()); + rk_tensors_[name] = CreateRknnTensor(name, dims, ptr, + rk::nn::TensorRole::CONST, + rk::nn::PrecisionType::UINT8); + } else if (tensor.data_type() == + ONNX_NAMESPACE::TensorProto_DataType_INT32) { + const char* ptr = tensor.int32_data().empty() + ? tensor.raw_data().data() + : reinterpret_cast( + tensor.int32_data().data()); + rk_tensors_[name] = CreateRknnTensor(name, dims, ptr, + rk::nn::TensorRole::CONST, + rk::nn::PrecisionType::INT32); + } else if (tensor.data_type() == + ONNX_NAMESPACE::TensorProto_DataType_INT64) { + const char* ptr = tensor.int64_data().empty() + ? tensor.raw_data().data() + : reinterpret_cast( + tensor.int64_data().data()); + rk_tensors_[name] = CreateRknnTensor(name, dims, ptr, + rk::nn::TensorRole::CONST, + rk::nn::PrecisionType::INT64); + } else { + LOGS_DEFAULT(FATAL) << "tensor name = " << tensor.name() + << ", data_type = " << tensor.data_type() + << " is not support!"; + assert(0); + } + operands_.push_back(name); + } +} + +std::vector> OnnxConverter::GetInputOfOnnxModel( + const std::vector& input_bufs, + const std::unordered_map& input_maps) { + std::vector> inputs; + + for (const auto& input : model_proto_.graph().input()) { + if (std::find(operands_.begin(), operands_.end(), input.name()) != + operands_.end()) { + continue; + } + + Shaper::Shape shape; + for (const auto& dim : input.type().tensor_type().shape().dim()) { + if (dim.value_case() == + ONNX_NAMESPACE::TensorShapeProto_Dimension::kDimValue) { + shape.push_back(static_cast(dim.dim_value())); + } else { + throw std::invalid_argument( + "The input of graph doesn't have dim_value"); + } + } + shaper_.AddShape(input.name(), shape); + + const void* ptr = NULL; + const auto iter = input_maps.find(input.name()); + if (iter != input_maps.end()) { + if (iter->second < (int)input_bufs.size()) + ptr = input_bufs[iter->second]; + } + + rk::nn::PrecisionType type = rk::nn::PrecisionType::FLOAT32; + if (input.type().tensor_type().has_elem_type()) { + switch (input.type().tensor_type().elem_type()) { + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: + type = rk::nn::PrecisionType::FLOAT32; + break; + case ONNX_NAMESPACE::TensorProto_DataType_UINT8: + type = rk::nn::PrecisionType::UINT8; + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT32: + type = rk::nn::PrecisionType::INT32; + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + type = rk::nn::PrecisionType::INT64; + break; + default: + // TODO: support other type + throw std::invalid_argument( + "The input of graph doesn't have valid type"); + } + } + + auto rk_input = CreateRknnTensor(input.name(), shape, ptr, + rk::nn::TensorRole::DATA, type); + + inputs.push_back(rk_input); + + rk_tensors_[input.name()] = rk_input; + } + if (inputs.empty()) + throw std::runtime_error("GetInputOfOnnxModel fail!"); + return inputs; +} + +std::vector> +OnnxConverter::GetOutputOfOnnxModel() { + std::vector> outputs; + for (const auto& output : model_proto_.graph().output()) { + const auto name = m(output.name()); + if (HAS(rk_tensors_, name)) { + outputs.push_back(rk_tensors_[name]); + } + } + if (outputs.empty()) + throw std::runtime_error("GetOutputOfOnnxModel fail!"); + return outputs; +} + +Shaper::Shape GetShape(const ONNX_NAMESPACE::ModelProto& model_proto, + const std::map>& tensor_dims, + const std::string& name) { + Shaper::Shape shape; + for (const auto& value_info : model_proto.graph().value_info()) { + if (value_info.name() == name) { + if (!value_info.has_type()) { + break; + } else if (!value_info.type().has_tensor_type()) { + break; + } else if (!value_info.type().tensor_type().has_shape()) { + break; + } else if (value_info.type().tensor_type().shape().dim_size() == 0) { + break; + } + + for (const auto& dim : value_info.type().tensor_type().shape().dim()) { + if (dim.has_dim_value()) { + shape.push_back(dim.dim_value()); + } else { + break; + } + } + + return shape; + } + } + + if (HAS(tensor_dims, name)) { + const auto& dims = tensor_dims.at(name); + if (dims.size() > 0) + shape = dims; + } + + return shape; +} + +bool TypeSupport(const ONNX_NAMESPACE::ModelProto& model_proto, + const std::string& name) { + for (const auto& value_info : model_proto.graph().value_info()) { + if (value_info.name() == name) { + if (!value_info.has_type()) { + break; + } else if (!value_info.type().has_tensor_type()) { + break; + } else if (!value_info.type().tensor_type().has_elem_type()) { + break; + } + + switch (value_info.type().tensor_type().elem_type()) { + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: + return true; + case ONNX_NAMESPACE::TensorProto_DataType_UINT8: + return true; + case ONNX_NAMESPACE::TensorProto_DataType_INT32: + return true; + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + return false; + default: + break; + } + + return false; + } + } + + // if can't find, it's considered support. + return true; +} + +std::pair OnnxConverter::IsNodeSupported( + const ONNX_NAMESPACE::ModelProto& model_proto, + const ONNX_NAMESPACE::NodeProto& node) const { + NodeAttrHelper helper(node); + const auto& op = node.op_type(); + const std::vector supported_types{ + "Conv", "Relu", "Clip", "LeakyRelu", + "MaxPool", "AveragePool", "GlobalAveragePool", + "Concat", "Softmax", "BatchNormalization", "Gemm", + "Add", "Mul", "Sub", + "Reshape", "Squeeze", "Unsqueeze", + "Flatten", "Transpose", /*"Gather", "Slice",*/ + "QLinearConv", /*"QuantizeLinear",*/ "DequantizeLinear"}; + if (std::find(supported_types.begin(), supported_types.end(), op) == + supported_types.end()) { + return {false, "Unsupported operator"}; + } + + if (!TypeSupport(model_proto, node.input(0))) { + return {false, "Type of input(" + node.input(0) + ") is unsupported"}; + } + + if (op == "Conv") { + const auto strides = helper.get("strides", vector{1, 1}); + const auto pads = helper.get("pads", vector{0, 0, 0, 0}); + const auto dilations = helper.get("dilations", vector{1, 1}); + const auto group = helper.get("group", 1); + if (dilations != vector{1, 1} && strides != vector{1, 1}) { + return {false, "Both dilations and strides > 1 is not supported for now"}; + } + const auto weight = m(node.input(1)); + if (HAS(tensor_dims_, weight)) { + const auto& dims = tensor_dims_.at(weight); + if (group != 1 && dims[1] != 1) { + return {false, "group != 1 is not supported"}; + } + if (dims.size() != 4) { + return {false, "Only conv 2d is supported."}; + } + } else { + return {false, "The weight of convolution must be known"}; + } + } else if (op == "AveragePool" || op == "MaxPool") { + const auto count_include_pad = helper.get("count_include_pad", 0); + if (count_include_pad == 1) { + return {false, "count_include_pad == 1 is not supported"}; + } + const auto storage_order = helper.get("storage_order", 0); + if (storage_order == 1) { + return {false, "storage_order == 1 is not supported"}; + } + if (helper.get("auto_pad", "NOTSET") != "NOTSET") { + return {false, "auto_pad is not supported"}; + } + if (helper.get("kernel_shape", std::vector{1, 1}).size() != 2) { + return {false, "Only pooling 2d is supported"}; + } + if (helper.get("ceil_mode", 0) == 1) { + return {false, "ceil_mode == 1 is not supported for pooling"}; + } + if (helper.get("dilations", std::vector{1, 1}) != + std::vector{1, 1}) { + return {false, "Dilations of pooling is not supported"}; + } + if (node.output_size() != 1) { + return {false, "Argmax in maxpooling is not supported"}; + } + } else if (op == "GlobalAveragePool" || op == "GlobalMaxPool") { + const auto& input_shape = GetShape(model_proto, + tensor_dims_, node.input(0)); + if (input_shape.size() == 0 || input_shape.size() != 4) { + return {false, "Only rank-4 tensor is supported"}; + } + } else if (op == "PRelu") { + const auto slope = m(node.input(1)); + if (HAS(tensor_dims_, slope)) { + if (tensor_dims_.at(slope) != Shaper::Shape{1}) { + // TODO: support it + return {false, "PRelu only support one element slope."}; + } + } else { + return {false, "PRelu slope must be known"}; + } + } else if (op == "Gemm") { + const auto transA = helper.get("transA", 0); + const auto transB = helper.get("transB", 0); + const auto alpha = helper.get("alpha", 1.0f); + const auto beta = helper.get("beta", 1.0f); + if (!(transA == 0 && transB == 1 && alpha == 1.f && beta == 1.f)) { + return {false, + "Only transA == 0, transB == 1, alpha == 1.0 and beta == " + "1.0 is supported."}; + } + } else if (op == "BatchNormalization") { + if (node.output_size() != 1) { + return {false, + "Your onnx model may be in training mode, please export " + "it in test mode."}; + } + const auto scale = m(node.input(1)); + const auto b = m(node.input(2)); + const auto mean = m(node.input(3)); + const auto var = m(node.input(4)); + if (!HAS(tensor_dims_, scale)) { + return {false, "Scale of BN must be known"}; + } + if (!HAS(tensor_dims_, b)) { + return {false, "B of BN must be known"}; + } + if (!HAS(tensor_dims_, mean)) { + return {false, "Mean of BN must be known"}; + } + if (!HAS(tensor_dims_, var)) { + return {false, "Var of BN must be known"}; + } + } else if (op == "LRN") { + const auto size = helper.get("size", 1); + if (size % 2 == 0) { + return {false, "NNAPI only support odd size for LRN"}; + } + } else if (op == "Reshape") { + const auto output = node.output(0); + for (const auto& another_node : model_proto_.graph().node()) { + for (const auto& input : another_node.input()) { + if (input == output && + another_node.op_type() != "Gemm") { + return {false, + "Reshape can only be the last layer or precede a " + "gemm layer for now"}; + } + } + } + const auto& shape = GetShape(model_proto, tensor_dims_, node.input(1)); + int rank = shape.size(); + if (shape.size() != 1) { + return {false, "Wrong shape rank"}; + } + if (shape[0] <= 1) { // rknpu doesn't support dims of shape equal 1, + // but simulator support, why? (reproduce on + // "Reshape_699" of ssd.onnx) + return {false, "Only shape dims > 1 is supported"}; + } + } else if (op == "Softmax") { + const auto axis = helper.get("axis", 1); + if (axis != 1) { + return {false, "Only axis == 1 is supported"}; + } + } else if (op == "Flatten") { + const auto axis = helper.get("axis", 1); + const auto& input_shape = GetShape(model_proto, + tensor_dims_, node.input(0)); + int rank = input_shape.size(); + if (rank != 0) { + if (axis < 0 || axis > (int64_t)rank) { + return {false, "Only axis <= rank of input is supported"}; + } + } + } else if (op == "Squeeze") { + const auto& input_shape = GetShape(model_proto, + tensor_dims_, node.input(0)); + const auto axes = helper.get("axes", std::vector{}); + int rank = input_shape.size(); + if (rank != 0) { + for (auto axis : axes) { + if (axis >= rank || axis < -1) { + return {false, "Only axes <= rank of input is supported"}; + } else { + axis = (axis < 0) ? (axis + rank) : (axis); + if (input_shape[axis] != 1) { + return {false, "the input_shape[axis] must equal one"}; + } + } + } + } + } else if (op == "Unsqueeze") { + const auto& input_shape = GetShape(model_proto, + tensor_dims_, node.input(0)); + const auto axes = helper.get("axes", std::vector{}); + int rank = input_shape.size(); + if (rank != 0) { + for (const auto axis : axes) { + if (axis < 0) { + return {false, "Only axes >= 0 is supported"}; + } else if (axis >= (int)(rank + axes.size())) { + return {false, "Only axes <= rank of (input + axes) is supported"}; + } + } + } + } else if (op == "Gather") { + const auto& input_shape = GetShape(model_proto, + tensor_dims_, node.input(0)); + const auto axis = helper.get("axis", 1); + int rank = input_shape.size(); + if (rank != 0 && (axis >= rank || axis < -rank)) { + return {false, "Only axis <= rank of input is supported"}; + } + } else if (op == "Concat") { + const auto& input_shape = GetShape(model_proto, + tensor_dims_, node.input(0)); + const auto axis = helper.get("axis", 1); + int rank = input_shape.size(); + if (axis >= 4) { + return {false, "Only axis <= 4 of input is supported"}; + } + if (rank != 0 && rank > 4 && (axis >= rank || axis < -rank)) { + if (rank > 4) + return {false, "Only rank <= 4 of input is supported"}; + else + return {false, "Only axis <= rank of input is supported"}; + } + } else if (op == "Add") { + if (!TypeSupport(model_proto, node.input(1))) { + return {false, "Type of input1(" + node.input(1) + ") is unsupported"}; + } + } else if (op == "Mul") { + if (!TypeSupport(model_proto, node.input(1))) { + return {false, "Type of input1(" + node.input(1) + ") is unsupported"}; + } + } + + return {true, ""}; +} + +std::vector> OnnxConverter::GetSupportedNodes( + const ONNX_NAMESPACE::ModelProto& model_proto) { + for (const auto& tensor : model_proto.graph().initializer()) { + const std::string name = tensor.name(); + std::vector dims; + for (const auto dim : tensor.dims()) { + dims.push_back(static_cast(dim)); + } + tensor_dims_[name] = dims; + } + + std::vector> supported_node_vecs; + std::vector supported_node_vec; + for (int i = 0; i < model_proto.graph().node_size(); i++) { + bool supported; + std::string error_msg; + std::tie(supported, error_msg) = + IsNodeSupported(model_proto, model_proto.graph().node(i)); + if (supported) { + supported_node_vec.push_back(i); + } else { + const auto& op = model_proto.graph().node(i).op_type(); + LOGS_DEFAULT(INFO) << op << ": " << error_msg; + if (!supported_node_vec.empty()) { + supported_node_vecs.push_back(supported_node_vec); + supported_node_vec.clear(); + } + } + } + if (!supported_node_vec.empty()) { + supported_node_vecs.push_back(supported_node_vec); + } + + for (auto it = tensor_dims_.begin(); it != tensor_dims_.end(); ++it) { + it->second.clear(); + tensor_dims_.erase(it); + } + return supported_node_vecs; +} + +void OnnxConverter::Convert(const ONNX_NAMESPACE::ModelProto& model_proto, + rk::nn::Graph* graph, + const std::vector& input_bufs, + const std::unordered_map& input_maps) { + model_proto_ = model_proto; + graph_ = graph; + + HandleInitializer(); + + const auto inputs = GetInputOfOnnxModel(input_bufs, input_maps); + + for (int i = 0; i < model_proto_.graph().node_size(); i++) { + const auto& node = model_proto_.graph().node(i); + NodeAttrHelper helper(node); + const auto& op = node.op_type(); + LOGS_DEFAULT(VERBOSE) << "----------- Node=" << node.name() + << ", op=" << op << "-----------"; + if (std::find(skipped_act_.begin(), skipped_act_.end(), i) != + skipped_act_.end()) { + LOGS_DEFAULT(INFO) << "Skip layer " << node.name() << ", op=" << op; + continue; + } + + if (op == "Conv") { + const auto strides = helper.get("strides", vector{1, 1}); + const auto pads = helper.get("pads", vector{0, 0, 0, 0}); + const auto dilations = helper.get("dilations", vector{1, 1}); + const auto group = helper.get("group", 1); + std::string bias; + if (node.input_size() >= 3) { + bias = m(node.input(2)); + } + const auto auto_pad = helper.get("auto_pad", "NOTSET"); + + const auto ori_weight = m(node.input(1)); + AddConv(m(node.input(0)), strides, pads, dilations, group, + ori_weight, bias, auto_pad, m(node.output(0))); + } else if (op == "QLinearConv") { + const auto strides = helper.get("strides", vector{1, 1}); + const auto pads = helper.get("pads", vector{0, 0, 0, 0}); + const auto dilations = helper.get("dilations", vector{1, 1}); + const auto group = helper.get("group", 1); + const auto auto_pad = helper.get("auto_pad", "NOTSET"); + std::string bias; + if (node.input_size() >= 9) { + bias = m(node.input(8)); + } + AddQLinearConv(m(node.input(0)), m(node.input(1)), m(node.input(2)), + strides, pads, dilations, group, auto_pad, + m(node.input(3)), m(node.input(4)), m(node.input(5)), + bias, m(node.output(0)), m(node.input(6)), + m(node.input(7))); + } else if (op == "AveragePool" || op == "MaxPool" || + op == "GlobalAveragePool" || op == "GlobalMaxPool") { + const auto input = m(node.input(0)); + const auto output = m(node.output(0)); + vector strides, pads, kernel_shape; + int ceil_mode; + if (op == "AveragePool" || op == "MaxPool") { + strides = helper.get("strides", vector{1, 1}); + pads = helper.get("pads", vector{0, 0, 0, 0}); + kernel_shape = helper.get("kernel_shape", vector{0, 0}); + ceil_mode = helper.get("ceil_mode", 0); + const auto count_include_pad = + helper.get("count_include_pad", 0); + if (count_include_pad == 1) { + throw std::invalid_argument( + "count_include_pad == 1 is not supported"); + } + const auto storage_order = helper.get("storage_order", 0); + if (storage_order == 1) { + throw std::invalid_argument( + "storage_order == 1 is not supported"); + } + if (helper.get("auto_pad", "NOTSET") != "NOTSET") { + throw std::invalid_argument("auto_pad is not supported"); + } + } else { + strides = {0, 0}; + pads = {0, 0, 0, 0}; + kernel_shape = {-1, -1}; // -1 for global + ceil_mode = 0; + } + AddLayerPool(op, input, kernel_shape, pads, strides, ceil_mode, + output); + } else if (op == "Relu") { + const auto input = m(node.input(0)); + const auto output = m(node.output(0)); + AddLayerReLU(input, output); + } else if (op == "PRelu") { + } else if (op == "Add") { + const auto input1 = m(node.input(0)); + const auto input2 = m(node.input(1)); + const auto output = m(node.output(0)); + AddLayerAdd(input1, input2, output); + } else if (op == "Sub") { + const auto input1 = m(node.input(0)); + const auto input2 = m(node.input(1)); + const auto output = m(node.output(0)); + AddLayerSub(input1, input2, output); + } else if (op == "Mul") { + const auto input1 = m(node.input(0)); + const auto input2 = m(node.input(1)); + const auto output = m(node.output(0)); + AddLayerMul(input1, input2, output); + } else if (op == "Gemm") { + const auto input = m(node.input(0)); + const auto weight = m(node.input(1)); + const auto output = m(node.output(0)); + string bias; + if (node.input_size() >= 3) { + bias = m(node.input(2)); + } + const auto transA = helper.get("transA", 0); + const auto transB = helper.get("transB", 0); + const auto alpha = helper.get("alpha", 1.0f); + const auto beta = helper.get("beta", 1.0f); + if (transA == 0 && transB == 1 && alpha == 1.f && beta == 1.f) { + AddLayerFC(input, weight, bias, output); + } else { + throw std::invalid_argument( + "Only transA == 0, transB == 1, alpha == 1.0 and beta == " + "1.0 is " + "supported."); + } + } else if (op == "Softmax") { + const auto input = m(node.input(0)); + const auto output = m(node.output(0)); + AddLayerSoftmax(input, output); + } else if (op == "Concat") { + vector concat_inputs_str; + for (const auto& onnx_input : node.input()) { + concat_inputs_str.push_back(m(onnx_input)); + } + const auto axis = helper.get("axis", 1); + const auto output = m(node.output(0)); + AddLayerConcat(concat_inputs_str, axis, output); + } else if (op == "Dropout") { + } else if (op == "BatchNormalization") { + const auto input = m(node.input(0)); + const auto scale = m(node.input(1)); + const auto bias = m(node.input(2)); + const auto mean = m(node.input(3)); + const auto var = m(node.input(4)); + const auto eps = helper.get("epsilon", 1e-5f); + const auto output = m(node.output(0)); + AddLayerBatchNorm(input, scale, bias, mean, + var, eps, output); + } else if (op == "Reshape") { + const auto input = m(node.input(0)); + const auto shape = m(node.input(1)); + const auto output = m(node.output(0)); + AddLayerReshape(input, shape, output); + } else if (op == "LRN") { + } else if (op == "Tanh") { + } else if (op == "Floor") { + } else if (op == "Sigmoid") { + } else if (op == "Flatten") { + const auto input = m(node.input(0)); + const auto axis = helper.get("axis", 1); + const auto output = m(node.output(0)); + AddLayerFlatten(input, axis, output); + } else if (op == "Transpose") { + const auto input = m(node.input(0)); + const auto perm = helper.get("perm", vector{0, 1, 2}); + const auto output = m(node.output(0)); + AddLayerTranspose(input, perm, output); + } else if (op == "Slice") { + const auto input = m(node.input(0)); + const auto starts = m(node.input(1)); + const auto ends = m(node.input(2)); + const auto axes = m(node.input(3)); + const auto steps = m(node.input(4)); + const auto output = m(node.output(0)); + AddLayerSlice(input, starts, ends, axes, + steps, output); + } else if (op == "Squeeze") { + const auto input = m(node.input(0)); + const auto axes = helper.get("axes", std::vector{}); + const auto output = m(node.output(0)); + AddLayerSqueeze(input, axes, output); + } else if (op == "Unsqueeze") { + const auto input = m(node.input(0)); + const auto axes = helper.get("axes", std::vector{}); + const auto output = m(node.output(0)); + AddLayerUnsqueeze(input, axes, output); + } else if (op == "Gather") { + const auto input = m(node.input(0)); + const auto indices = m(node.input(1)); + const auto axis = helper.get("axis", 0); + const auto output = m(node.output(0)); + AddLayerGather(input, indices, axis, output); + } else if (op == "LeakyRelu") { + const auto input = m(node.input(0)); + const auto alpha = helper.get("alpha", (float)0.0); + const auto output = m(node.output(0)); + AddLayerLeakyRelu(input, alpha, output); + } else if (op == "Clip") { + const auto input = m(node.input(0)); + const auto min = helper.get("axis", 0); + const auto max = helper.get("axis", 6); + const auto output = m(node.output(0)); + AddLayerClip(input, min, max, output); + } else if (op == "QuantizeLinear") { + const auto input = m(node.input(0)); + const auto output_scale = m(node.input(1)); + const auto output_zp = m(node.input(2)); + const auto output = m(node.output(0)); + AddLayerQuantizeLinear(input, output_scale, output_zp, output); + } else if (op == "DequantizeLinear") { + const auto input = m(node.input(0)); + const auto input_scale = m(node.input(1)); + const auto input_zp = m(node.input(2)); + const auto output = m(node.output(0)); + AddLayerDequantizeLinear(input, input_scale, input_zp, output); + } else { + throw std::invalid_argument("Unsupported operator " + op); + } + } + + const auto outputs = GetOutputOfOnnxModel(); + + graph_->SetInputsOutputs(inputs, outputs); +} + +void OnnxConverter::Clear() { + skipped_act_.clear(); + operands_.clear(); + name_map_.clear(); + rk_tensors_.clear(); + shaper_.Clear(); + + for (const auto p : free_list_) { + if (p) free(p); + } + free_list_.clear(); +} + +void OnnxConverter::AddConv(const string& input, + const std::vector& strides, + const std::vector& pads, + const std::vector& dilations, + const int group, + const string& ori_weight, + const string& bias, + const string& auto_pad, + const string& output) { + if (dilations != vector{1, 1}) { + // TODO + throw std::invalid_argument("dilations != 1 is not supported yet"); + } + + if (!HAS(rk_tensors_, ori_weight)) { + throw std::invalid_argument("The weight of convolution must be known"); + } + const auto& weight = rk_tensors_.at(ori_weight); + if (group == 1) { + LOGS_DEFAULT(VERBOSE) << "Vanilla conv"; + AddLayerConvImpl(input, ori_weight, bias, pads, strides, + 1, auto_pad, output); + } else if (weight->GetDims()[1] == 1) { // depthwise + LOGS_DEFAULT(VERBOSE) << "Depthwise conv"; + AddLayerDepthwiseConvImpl(input, ori_weight, bias, pads, + strides, weight->GetDims()[0] / group, group, + output); + } else { + LOGS_DEFAULT(VERBOSE) << "Group conv"; + AddLayerConvImpl(input, ori_weight, bias, pads, strides, + group, auto_pad, output); + } +} + +void OnnxConverter::AddQLinearConv(const string& input, + const string& input_scale, + const string& input_zp, + const std::vector& strides, + const std::vector& pads, + const std::vector& dilations, + const int group, + const string& auto_pad, + const string& weight, + const string& weight_scale, + const string& weight_zp, + const string& bias, + const string& output, + const string& output_scale, + const string& output_zp) { + if (dilations != vector{1, 1}) { + return; + } + + if (!HAS(rk_tensors_, weight)) { + throw std::invalid_argument("The weight of convolution must be known"); + } + if (group == 1) { + LOGS_DEFAULT(VERBOSE) << "Vanilla QLinearConv"; + AddLayerQLinearConvImpl(input, input_scale, input_zp, + weight, weight_scale, weight_zp, + bias, pads, strides, 1, auto_pad, + output, output_scale, output_zp); + } else if (rk_tensors_.at(weight)->GetDims()[1] == 1) { // depthwise + LOGS_DEFAULT(VERBOSE) << "Depthwise QLinearConv"; + // TODO + throw std::invalid_argument("Depthwise QLinearConv is not supported yet"); + } else { + LOGS_DEFAULT(VERBOSE) << "Group QLinearConv"; + AddLayerQLinearConvImpl(input, input_scale, input_zp, + weight, weight_scale, weight_zp, + bias, pads, strides, group, auto_pad, + output, output_scale, output_zp); + } +} + +void OnnxConverter::AddLayerPool(const std::string& op, + const std::string& input, + const std::vector& kernel_shape, + const std::vector& pads, + const std::vector& strides, + const int32_t ceil_mode, + const std::string& output) { + if (op == "AveragePool" || op == "GlobalAveragePool") { + AddLayerAvePoolImpl(input, kernel_shape, pads, strides, + ceil_mode, output); + } else { + AddLayerMaxPoolImpl(input, kernel_shape, pads, strides, + ceil_mode, output); + } +} + +void OnnxConverter::AddLayerConvImpl(const std::string& input, + const std::string& weight, + const std::string& bias, + const std::vector& pads, + const std::vector& strides, + const int32_t group, + const std::string& auto_pad, + const std::string& output) { + const auto activation = FindActivation(model_proto_, output); + + ADD_SHAPE(input); + ADD_SHAPE(weight); + if (bias != "") { + ADD_SHAPE(bias); + } + + shaper_.Conv(input, weight, pads, strides, auto_pad, output); + + std::vector> inputs, outputs; + std::vector weight_dims; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, weight)) { + weight_dims = rk_tensors_.at(weight)->GetDims(); + inputs.push_back(rk_tensors_.at(weight)); + } + if (bias != "") { + if (HAS(rk_tensors_, bias)) { + inputs.push_back(rk_tensors_.at(bias)); + } + } else { + uint32_t dim = shaper_[weight][0]; + void* ptr = (void*)malloc(sizeof(float) * dim); + memset(ptr, 0, sizeof(float) * dim); + free_list_.push_back(ptr); + + std::vector dims = {dim}; + auto rk_bias = CreateRknnTensor(bias, dims, ptr, rk::nn::TensorRole::CONST); + inputs.push_back(rk_bias); + } + + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::PadType pad_type = rk::nn::PadType::AUTO; + if (auto_pad == "NOTSET") { + pad_type = rk::nn::PadType::AUTO; + } else if (auto_pad == "SAME_UPPER" || auto_pad == "SAME_LOWER") { + pad_type = rk::nn::PadType::SAME; + } else if (auto_pad == "VALID") { + pad_type = rk::nn::PadType::VALID; + } + + rk::nn::Conv2DAttr attr; + attr.ksize[0] = weight_dims[3]; + attr.ksize[1] = weight_dims[2]; + attr.stride[0] = strides[1]; + attr.stride[1] = strides[0]; + attr.pad[0] = pads[1]; + attr.pad[1] = pads[3]; + attr.pad[2] = pads[0]; + attr.pad[3] = pads[2]; + attr.group = group; + attr.weights = weight_dims[0]; // TODO + attr.dilation[0] = 1; + attr.dilation[1] = 1; + attr.pad_type = pad_type; + attr.multiplier = 0; // TODO + attr.has_relu = (activation.second == FuseCode::FUSED_RELU) ? true : false; + + graph_->AddOperator(rk::nn::OperatorType::CONV2D, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerQLinearConvImpl(const string& input, + const string& input_scale, + const string& input_zp, + const string& weight, + const string& weight_scale, + const string& weight_zp, + const string& bias, + const std::vector& pads, + const std::vector& strides, + const int group, + const string& auto_pad, + const string& output, + const string& output_scale, + const string& output_zp) { + const auto activation = FindActivation(model_proto_, output); + + ADD_SHAPE(input); + ADD_SHAPE(weight); + if (bias != "") { + ADD_SHAPE(bias); + } + + shaper_.Conv(input, weight, pads, strides, auto_pad, output); + + GET_ATTR(in_s, input_scale, float); + GET_ATTR(in_zp, input_zp, uint8_t); + GET_ATTR(w_s, weight_scale, float); + GET_ATTR(w_zp, weight_zp, uint8_t); + GET_ATTR(out_s, output_scale, float); + GET_ATTR(out_zp, output_zp, uint8_t); + + std::vector> inputs, outputs; + std::vector weight_dims; + if (HAS(rk_tensors_, input)) { + rk::nn::QuantizationParamAffineAsymmetric param; + param.scale.push_back(in_s[0]); + param.zero_point.push_back(in_zp[0]); + auto tensor = rk_tensors_.at(input); + tensor->SetQntParam(rk::nn::QuantizationType::AFFINE_ASYMMETRIC, 8, param); + inputs.push_back(tensor); + } + if (HAS(rk_tensors_, weight)) { + rk::nn::QuantizationParamAffineAsymmetric param; + param.scale.push_back(w_s[0]); + param.zero_point.push_back(w_zp[0]); + weight_dims = rk_tensors_.at(weight)->GetDims(); + auto tensor = rk_tensors_.at(weight); + tensor->SetQntParam(rk::nn::QuantizationType::AFFINE_ASYMMETRIC, 8, param); + inputs.push_back(tensor); + } + if (bias != "") { + if (HAS(rk_tensors_, bias)) { + rk::nn::QuantizationParamSymmetric param; + param.scale.push_back(in_s[0] * w_s[0]); + auto tensor = rk_tensors_.at(bias); + tensor->SetQntParam(rk::nn::QuantizationType::SYMMETRIC, 32, param); + inputs.push_back(tensor); + } + } else { + uint32_t dim = shaper_[weight][0]; + void* ptr = (void*)malloc(sizeof(int32_t) * dim); + memset(ptr, 0, sizeof(int32_t) * dim); + free_list_.push_back(ptr); + + std::vector dims = {dim}; + auto rk_bias = CreateRknnTensor(bias, dims, ptr, rk::nn::TensorRole::CONST, + rk::nn::PrecisionType::INT32, + rk::nn::DataLayoutType::NCHW, + rk::nn::QuantizationType::SYMMETRIC, + 32, in_s[0] * w_s[0]); + inputs.push_back(rk_bias); + } + + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], NULL, + rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision(), + rk::nn::DataLayoutType::NCHW, + rk::nn::QuantizationType::AFFINE_ASYMMETRIC, + 8, out_s[0], out_zp[0]); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::PadType pad_type = rk::nn::PadType::AUTO; + if (auto_pad == "NOTSET") { + pad_type = rk::nn::PadType::AUTO; + } else if (auto_pad == "SAME_UPPER" || auto_pad == "SAME_LOWER") { + pad_type = rk::nn::PadType::SAME; + } else if (auto_pad == "VALID") { + pad_type = rk::nn::PadType::VALID; + } + + rk::nn::Conv2DAttr attr; + attr.ksize[0] = weight_dims[3]; + attr.ksize[1] = weight_dims[2]; + attr.stride[0] = strides[1]; + attr.stride[1] = strides[0]; + attr.pad[0] = pads[1]; + attr.pad[1] = pads[3]; + attr.pad[2] = pads[0]; + attr.pad[3] = pads[2]; + attr.group = group; + attr.weights = weight_dims[0]; // TODO + attr.dilation[0] = 1; + attr.dilation[1] = 1; + attr.pad_type = pad_type; + attr.multiplier = 0; // TODO + attr.has_relu = (activation.second == FuseCode::FUSED_RELU) ? true : false; + + graph_->AddOperator(rk::nn::OperatorType::CONV2D, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerDepthwiseConvImpl( + const std::string& input, + const std::string& weight, + const std::string& bias, + const std::vector& pads, + const std::vector& strides, + const int32_t depth_multiplier, + const int32_t group, + const std::string& output) { + const auto activation = FindActivation(model_proto_, output); + + ADD_SHAPE(input); + ADD_SHAPE(weight); + if (bias != "") { + ADD_SHAPE(bias); + } + shaper_.DepthwiseConv(input, weight, pads, strides, output); + + std::vector> inputs, outputs; + std::vector weight_dims; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, weight)) { + weight_dims = rk_tensors_.at(weight)->GetDims(); + inputs.push_back(rk_tensors_.at(weight)); + } + if (bias != "") { + if (HAS(rk_tensors_, bias)) { + inputs.push_back(rk_tensors_.at(bias)); + } + } else { + uint32_t dim = shaper_[weight][0]; + void* ptr = (void*)malloc(sizeof(float) * dim); + memset(ptr, 0, sizeof(float) * dim); + free_list_.push_back(ptr); + + std::vector dims = {dim}; + auto rk_bias = CreateRknnTensor(bias, dims, ptr, rk::nn::TensorRole::CONST); + inputs.push_back(rk_bias); + } + + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::Conv2DAttr attr; + attr.ksize[0] = weight_dims[3]; + attr.ksize[1] = weight_dims[2]; + attr.stride[0] = strides[1]; + attr.stride[1] = strides[0]; + attr.pad[0] = pads[1]; + attr.pad[1] = pads[3]; + attr.pad[2] = pads[0]; + attr.pad[3] = pads[2]; + attr.group = group; + attr.weights = weight_dims[0]; // TODO + attr.dilation[0] = 1; + attr.dilation[1] = 1; + attr.pad_type = rk::nn::PadType::AUTO; + attr.multiplier = depth_multiplier; // TODO + attr.has_relu = (activation.second == FuseCode::FUSED_RELU) ? true : false; + + graph_->AddOperator(rk::nn::OperatorType::CONV2D, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerConcat(const std::vector& input, + const int32_t axis, + const std::string& output) { + for (const auto& name : input) { + ADD_SHAPE(name); + } + shaper_.Concat(input, axis, output); + + std::vector> inputs, outputs; + for (const auto& in : input) { + if (HAS(rk_tensors_, in)) { + inputs.push_back(rk_tensors_.at(in)); + } + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::ConcatAttr attr; + attr.axis = axis; + graph_->AddOperator(rk::nn::OperatorType::CONCAT, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerAvePoolImpl( + const std::string& input, + const std::vector& kernel_shape, + const std::vector& pads, + const std::vector& strides, + const int32_t ceil_mode, + const std::string& output) { + ADD_SHAPE(input); + shaper_.Pool(input, kernel_shape, pads, strides, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::PoolAttr attr; + attr.ksize[0] = kernel_shape[0]; + attr.ksize[1] = kernel_shape[1]; + attr.stride[0] = strides[0]; + attr.stride[1] = strides[1]; + attr.pad[0] = pads[0]; + attr.pad[1] = pads[2]; + attr.pad[2] = pads[1]; + attr.pad[3] = pads[3]; + attr.pad_type = rk::nn::PadType::AUTO; + attr.pool_type = rk::nn::PoolType::POOLING_AVG; + attr.round_type = + (ceil_mode == 1) ? rk::nn::RoundType::ROUND_CEIL : rk::nn::RoundType::ROUND_FLOOR; + attr.global_pooling = (kernel_shape[0] == -1 && kernel_shape[1] == -1); + + graph_->AddOperator(rk::nn::OperatorType::POOL, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerMaxPoolImpl( + const std::string& input, + const std::vector& kernel_shape, + const std::vector& pads, + const std::vector& strides, + const int32_t ceil_mode, + const std::string& output) { + ADD_SHAPE(input); + shaper_.Pool(input, kernel_shape, pads, strides, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::PoolAttr attr; + attr.ksize[0] = kernel_shape[0]; + attr.ksize[1] = kernel_shape[1]; + attr.stride[0] = strides[0]; + attr.stride[1] = strides[1]; + attr.pad[0] = pads[0]; + attr.pad[1] = pads[2]; + attr.pad[2] = pads[1]; + attr.pad[3] = pads[3]; + attr.pad_type = rk::nn::PadType::AUTO; + attr.pool_type = rk::nn::PoolType::POOLING_MAX; + attr.round_type = + (ceil_mode == 1) ? rk::nn::RoundType::ROUND_CEIL : rk::nn::RoundType::ROUND_FLOOR; + attr.global_pooling = (kernel_shape[0] == -1 && kernel_shape[1] == -1); + + graph_->AddOperator(rk::nn::OperatorType::POOL, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerReLU(const std::string& input, + const std::string& output) { + ADD_SHAPE(input); + shaper_.Relu(input, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + graph_->AddOperator(rk::nn::OperatorType::RELU, inputs, outputs, nullptr); +} + +void OnnxConverter::AddLayerSoftmax(const std::string& input, + const std::string& output) { + ADD_SHAPE(input); + shaper_.Softmax(input, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::SoftmaxAttr attr; + attr.axis = 1; + attr.beta = 1.0; + + graph_->AddOperator(rk::nn::OperatorType::SOFTMAX, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerFC(const std::string& input, + const std::string& weight, + const std::string& bias, + const std::string& output) { + const auto activation = FindActivation(model_proto_, output); + + ADD_SHAPE(input); + ADD_SHAPE(weight); + if (bias != "") { + ADD_SHAPE(bias); + } + + shaper_.FC(input, weight, output); + + std::vector> inputs, outputs; + std::vector weight_dims; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, weight)) { + weight_dims = rk_tensors_.at(weight)->GetDims(); + inputs.push_back(rk_tensors_.at(weight)); + } + if (bias != "") { + if (HAS(rk_tensors_, bias)) { + inputs.push_back(rk_tensors_.at(bias)); + } + } else { + uint32_t dim = shaper_[weight][0]; + void* ptr = (void*)malloc(sizeof(float) * dim); + memset(ptr, 0, sizeof(float) * dim); + free_list_.push_back(ptr); + + std::vector dims = {dim}; + auto rk_bias = CreateRknnTensor(bias, dims, ptr, rk::nn::TensorRole::CONST); + inputs.push_back(rk_bias); + } + + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::FCAttr attr; + attr.weights = weight_dims[1]; // TODO + attr.has_relu = (activation.second == FuseCode::FUSED_RELU) ? true : false; + + graph_->AddOperator(rk::nn::OperatorType::FULLCONNECT, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerAdd(const std::string& input1, + const std::string& input2, + const std::string& output) { + ADD_SHAPE(input1); + ADD_SHAPE(input2); + shaper_.Eltwise(input1, input2, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input1)) { + inputs.push_back(rk_tensors_.at(input1)); + } + if (HAS(rk_tensors_, input2)) { + inputs.push_back(rk_tensors_.at(input2)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + graph_->AddOperator(rk::nn::OperatorType::ADD, inputs, outputs, nullptr); +} + +void OnnxConverter::AddLayerSub(const std::string& input1, + const std::string& input2, + const std::string& output) { + ADD_SHAPE(input1); + ADD_SHAPE(input2); + shaper_.Eltwise(input1, input2, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input1)) { + inputs.push_back(rk_tensors_.at(input1)); + } + if (HAS(rk_tensors_, input2)) { + inputs.push_back(rk_tensors_.at(input2)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + graph_->AddOperator(rk::nn::OperatorType::SUBTRACT, inputs, outputs, nullptr); +} + +void OnnxConverter::AddLayerMul(const std::string& input1, + const std::string& input2, + const std::string& output) { + ADD_SHAPE(input1); + ADD_SHAPE(input2); + shaper_.Eltwise(input1, input2, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input1)) { + inputs.push_back(rk_tensors_.at(input1)); + } + if (HAS(rk_tensors_, input2)) { + inputs.push_back(rk_tensors_.at(input2)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + graph_->AddOperator(rk::nn::OperatorType::MULTIPLY, inputs, outputs, nullptr); +} + +void OnnxConverter::AddLayerBatchNorm(const string& input, + const string& scale, + const string& bias, + const string& mean, + const string& var, + const float eps, + const string& output) { + ADD_SHAPE(input); + shaper_.BatchNorm(input, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, mean)) { + inputs.push_back(rk_tensors_.at(mean)); + } + if (HAS(rk_tensors_, var)) { + inputs.push_back(rk_tensors_.at(var)); + } + if (HAS(rk_tensors_, scale)) { + inputs.push_back(rk_tensors_.at(scale)); + } + if (HAS(rk_tensors_, bias)) { + inputs.push_back(rk_tensors_.at(bias)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::BatchNormAttr attrs; + attrs.eps = eps; + + graph_->AddOperator(rk::nn::OperatorType::BATCH_NORM, + inputs, outputs, &attrs); +} + +void OnnxConverter::AddLayerReshape(const string& input, + const string& shape, + const string& output) { + ADD_SHAPE(input); + + GET_ATTR(v_shape, shape, int32_t); + + shaper_.Reshape(input, v_shape, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + // the "TFNodes/yolo_evaluation_layer_1/Reshape_5:0" of yolov3.onnx will + // wrong when role is VAR, modify it to DATA for workaround. + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::DATA, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::ReshapeAttr attr; + for (const auto dim : shaper_[output]) { + attr.shapes.push_back(static_cast(dim)); + } + graph_->AddOperator(rk::nn::OperatorType::RESHAPE, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerFlatten(const string& input, + const int32_t axis, + const string& output) { + ADD_SHAPE(input); + + std::vector shape = {1, -1}; // axis = 0 + if (axis > 0) { + const auto in_shape = shaper_[input]; + shape[0] = (int32_t)std::accumulate(in_shape.begin(), + in_shape.begin() + axis, 1, + std::multiplies()); + } + + shaper_.Reshape(input, shape, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::ReshapeAttr attr; + for (const auto dim : shaper_[output]) { + attr.shapes.push_back(static_cast(dim)); + } + graph_->AddOperator(rk::nn::OperatorType::RESHAPE, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerTranspose(const string& input, + const std::vector& perm, + const string& output) { + ADD_SHAPE(input); + + shaper_.Transpose(input, perm, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::PermuteAttr attr; + for (const auto val : perm) { + attr.perm.push_back(static_cast(val)); + } + graph_->AddOperator(rk::nn::OperatorType::PERMUTE, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerSlice(const string& input, + const string& starts, + const string& ends, + const string& axes, + const string& steps, + const string& output) { + ADD_SHAPE(input); + + GET_ATTR(v_starts, starts, int32_t); + GET_ATTR(v_ends, ends, int32_t); + GET_ATTR(v_axes, axes, int32_t); + GET_ATTR(v_steps, steps, int32_t); + + for (const auto step : v_steps) { + if (step != 1) { + LOGS_DEFAULT(FATAL) << "the steps of Slice must be 1!"; + return; + } + } + + shaper_.Slice(input, v_starts, v_ends, v_axes, v_steps, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::SliceAttr attr; + + for (const auto dim : shaper_[output]) { + attr.start.push_back(0); + attr.length.push_back(static_cast(dim)); + } + + const auto input_dims = shaper_[input]; + for (size_t i = 0; i < v_axes.size(); i++) { + int32_t dim = input_dims[v_axes[i]]; + if (dim > 0) { + int32_t start = v_starts[i] < 0 ? (v_starts[i] + dim) : v_starts[i]; + attr.start[v_axes[i]] = std::max(start, 0); + } + } + + graph_->AddOperator(rk::nn::OperatorType::SLICE, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerSqueeze(const string& input, + const std::vector& axes, + const string& output) { + ADD_SHAPE(input); + + shaper_.Squeeze(input, axes, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::ReshapeAttr attr; + for (const auto dim : shaper_[output]) { + attr.shapes.push_back(static_cast(dim)); + } + graph_->AddOperator(rk::nn::OperatorType::RESHAPE, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerUnsqueeze(const string& input, + const std::vector& axes, + const string& output) { + ADD_SHAPE(input); + + shaper_.Unsqueeze(input, axes, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::ReshapeAttr attr; + for (const auto dim : shaper_[output]) { + attr.shapes.push_back(static_cast(dim)); + } + graph_->AddOperator(rk::nn::OperatorType::RESHAPE, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerGather(const string& input, + const string& indices, + const int32_t axis, + const string& output) { + ADD_SHAPE(input); + ADD_SHAPE(indices); + + shaper_.Gather(input, indices, (int32_t)axis, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::GatherAttr attr; + attr.axis = axis; + graph_->AddOperator(rk::nn::OperatorType::GATHER, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerLeakyRelu(const std::string& input, + const float alpha, + const std::string& output) { + ADD_SHAPE(input); + shaper_.Relu(input, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + rk::nn::LeakyReluAttr attr; + attr.alpha = alpha; + graph_->AddOperator(rk::nn::OperatorType::LEAKY_RELU, + inputs, outputs, (void*)&attr); +} + +void OnnxConverter::AddLayerClip(const std::string& input, + const int32_t min, + const int32_t max, + const std::string& output) { + ADD_SHAPE(input); + shaper_.Relu(input, output); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], + NULL, rk::nn::TensorRole::VAR, + inputs[0]->GetPrecision()); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + /* if (min == 0 && max == 6) { // RELU6 will cause loss of accuracy + graph_->AddOperator(rk::nn::OperatorType::RELU6, + inputs, outputs, nullptr); + } + else */ + { + rk::nn::ClipAttr attr; + attr.min = min; + attr.max = max; + graph_->AddOperator(rk::nn::OperatorType::CLIP, + inputs, outputs, (void*)&attr); + } +} + +void OnnxConverter::AddLayerQuantizeLinear(const string& input, + const string& output_scale, + const string& output_zp, + const string& output) { + ADD_SHAPE(input); + shaper_.Identity(input, output); + + GET_ATTR(out_s, output_scale, float); + GET_ATTR(out_zp, output_zp, uint8_t); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + inputs.push_back(rk_tensors_.at(input)); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], NULL, + rk::nn::TensorRole::VAR, + rk::nn::PrecisionType::UINT8, + rk::nn::DataLayoutType::NCHW, + rk::nn::QuantizationType::AFFINE_ASYMMETRIC, + 8, out_s[0], out_zp[0]); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + graph_->AddOperator(rk::nn::OperatorType::DATACONVERT, + inputs, outputs, nullptr); +} + +void OnnxConverter::AddLayerDequantizeLinear(const string& input, + const string& input_scale, + const string& input_zp, + const string& output) { + ADD_SHAPE(input); + shaper_.Identity(input, output); + + GET_ATTR(in_s, input_scale, float); + GET_ATTR(in_zp, input_zp, uint8_t); + + std::vector> inputs, outputs; + if (HAS(rk_tensors_, input)) { + rk::nn::QuantizationParamAffineAsymmetric param; + param.scale.push_back(in_s[0]); + param.zero_point.push_back(in_zp[0]); + auto tensor = rk_tensors_.at(input); + tensor->SetQntParam(rk::nn::QuantizationType::AFFINE_ASYMMETRIC, 8, param); + inputs.push_back(tensor); + } + if (HAS(rk_tensors_, output)) { + outputs.push_back(rk_tensors_.at(output)); + } else { + auto rk_output = CreateRknnTensor(output, shaper_[output], NULL, + rk::nn::TensorRole::VAR, + rk::nn::PrecisionType::FLOAT32); + outputs.push_back(rk_output); + rk_tensors_[output] = rk_output; + } + + graph_->AddOperator(rk::nn::OperatorType::DATACONVERT, + inputs, outputs, nullptr); +} + +} // namespace rknpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rknpu/onnx_converter.h b/onnxruntime/core/providers/rknpu/onnx_converter.h new file mode 100644 index 0000000000..370187281d --- /dev/null +++ b/onnxruntime/core/providers/rknpu/onnx_converter.h @@ -0,0 +1,280 @@ +// Copyright 2020 rock-chips.com Inc. + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +#include "shaper.h" + +#include "rknpu/rknpu_pub.h" + +namespace onnxruntime { +namespace rknpu { + +/** + * For convert from onnx::ModelProto to rk::nn::Graph. + */ +class OnnxConverter { + public: + OnnxConverter() {} + ~OnnxConverter() { Clear(); } + + /** Get the supported subgraph. + */ + std::vector> GetSupportedNodes( + const ONNX_NAMESPACE::ModelProto& model_proto); + + /** Convert onnx::ModelProto to rk::nn::Graph. + * Because some attributes of rk::nn::Tensor are used as input in onnx and these attributes + * can't be found in onnx::ModelProto, so additional input-related information is required. + */ + void Convert(const ONNX_NAMESPACE::ModelProto& model, + rk::nn::Graph* graph, + const std::vector& input_bufs, + const std::unordered_map& input_maps); + + std::string m(const std::string& str) const; + + private: + Shaper shaper_; + + enum class FuseCode { FUSED_NONE, + FUSED_RELU, + FUSED_RELU1, + FUSED_RELU6 }; + + std::map name_map_; + + ONNX_NAMESPACE::ModelProto model_proto_; + rk::nn::Graph* graph_; + std::vector skipped_act_; + std::vector dequantize_after_; + + std::vector operands_; + std::map> rk_tensors_; + + // for GetSupportedNodes + std::map> tensor_dims_; + + std::vector free_list_; // remember free + + std::pair, FuseCode> + FindActivation(const ONNX_NAMESPACE::ModelProto& model_proto, + const std::string& output); + + std::shared_ptr + CreateRknnTensor(const std::string& name, + const std::vector& dims, + const void* data = NULL, + const rk::nn::TensorRole role = rk::nn::TensorRole::VAR, + const rk::nn::PrecisionType precision = rk::nn::PrecisionType::FLOAT32, + const rk::nn::DataLayoutType layout = rk::nn::DataLayoutType::NCHW, + const rk::nn::QuantizationType qntType = rk::nn::QuantizationType::NONE, + const uint8_t bits = 8, + const float scale = 1.0, + const uint32_t zero_point = 0, + const int8_t fl = 0); + + void HandleInitializer(); + std::vector> GetInputOfOnnxModel( + const std::vector& input_bufs, + const std::unordered_map& input_maps); + std::vector> GetOutputOfOnnxModel(); + + std::pair IsNodeSupported( + const ONNX_NAMESPACE::ModelProto& model_proto, + const ONNX_NAMESPACE::NodeProto& node_proto) const; + + void AddConv(const std::string& input, + const std::vector& strides, + const std::vector& pads, + const std::vector& dilations, + const int32_t group, + const std::string& ori_weight, + const std::string& bias, + const std::string& auto_pad, + const std::string& output); + void AddQLinearConv(const std::string& input, + const std::string& input_scale, + const std::string& input_zp, + const std::vector& strides, + const std::vector& pads, + const std::vector& dilations, + const int group, + const std::string& auto_pad, + const std::string& weight, + const std::string& weight_scale, + const std::string& weight_zp, + const std::string& bias, + const std::string& output, + const std::string& output_scale, + const std::string& output_zp); + void AddLayerPool(const std::string& op, + const std::string& input, + const std::vector& kernel_shape, + const std::vector& pads, + const std::vector& strides, + const int32_t ceil_mode, + const std::string& output); + void SetIdentity(const std::string& input, + const std::string& output); + void AddLayerConvImpl(const std::string& input, + const std::string& weight, + const std::string& bias, + const std::vector& pads, + const std::vector& strides, + const int32_t group, + const std::string& auto_pad, + const std::string& output); + void AddLayerQLinearConvImpl(const std::string& input, + const std::string& input_scale, + const std::string& input_zp, + const std::string& weight, + const std::string& weight_scale, + const std::string& weight_zp, + const std::string& bias, + const std::vector& pads, + const std::vector& strides, + const int group, + const std::string& auto_pad, + const std::string& output, + const std::string& output_scale, + const std::string& output_zp); + void AddLayerAvePoolImpl(const std::string& input, + const std::vector& kernel_shape, + const std::vector& pads, + const std::vector& strides, + const int32_t ceil_mode, + const std::string& output); + void AddLayerMaxPoolImpl(const std::string& input, + const std::vector& kernel_shape, + const std::vector& pads, + const std::vector& strides, + const int32_t ceil_mode, + const std::string& output); + void AddLayerReLU(const std::string& input, + const std::string& output); + void AddLayerSoftmax(const std::string& input, + const std::string& output); + void AddLayerFC(const std::string& input, + const std::string& weight, + const std::string& bias, + const std::string& output); + void AddLayerAdd(const std::string& input1, + const std::string& input2, + const std::string& output); + void AddLayerSub(const std::string& input1, + const std::string& input2, + const std::string& output); + void AddLayerConcat(const std::vector& inputs, + const int32_t axis, + const std::string& output); + void AddLayerDepthwiseConvImpl(const std::string& input, + const std::string& weight, + const std::string& bias, + const std::vector& pads, + const std::vector& strides, + const int32_t depth_multiplier, + const int32_t group, + const std::string& output); + void AddLayerBatchToSpaceND(const std::string& input, + const std::vector& block_sizes, + const std::string& output); + void AddLayerSpaceToBatchND(const std::string& input, + const std::vector& block_sizes, + const std::vector& pads, + const std::string& output); + void AddLayerSlice(const std::string& input, + const std::string& starts, + const std::string& ends, + const std::string& axes, + const std::string& steps, + const std::string& output); + void AddLayerStridedSlice(const std::string& input, + const std::vector& starts, + const std::vector& ends, + const std::vector& strides, + const int32_t begin_mask, + const int32_t end_mask, + const int32_t shrink_axis_mask, + const std::string& output); + void AddLayerMul(const std::string& input1, + const std::string& input2, + const std::string& output); + void AddLayerAdd(const std::string& input, + const float scalar, + const std::string& output); + void AddLayerMul(const std::string& input, + const float scalar, + const std::string& output); + void AddLayerDequantize(const std::string& input, + const std::string& output); + void AddLayerLRN(const std::string& input, + const int32_t radius, + const float bias, + const float alpha, + const float beta, + const std::string& output); + void AddLayerTanh(const std::string& input, + const std::string& output); + void AddLayerFloor(const std::string& input, + const std::string& output); + void AddLayerLogistic(const std::string& input, + const std::string& output); + void AddLayerBatchNorm(const std::string& input, + const std::string& scale, + const std::string& bias, + const std::string& mean, + const std::string& var, + const float eps, + const std::string& output); + void AddLayerReshape(const std::string& input, + const std::string& shape, + const std::string& output); + void AddLayerFlatten(const std::string& input, + const int32_t axis, + const std::string& output); + void AddLayerTranspose(const std::string& input, + const std::vector& perm, + const std::string& output); + void AddLayerSqueeze(const std::string& input, + const std::vector& axes, + const std::string& output); + void AddLayerUnsqueeze(const std::string& input, + const std::vector& axes, + const std::string& output); + void AddLayerGather(const std::string& input, + const std::string& indices, + const int32_t axis, + const std::string& output); + void AddLayerLeakyRelu(const std::string& input, + const float alpha, + const std::string& output); + void AddLayerClip(const std::string& input, + const int32_t min, + const int32_t max, + const std::string& output); + void AddLayerDequantizeLinear(const std::string& input, + const std::string& input_scale, + const std::string& input_zp, + const std::string& output); + void AddLayerQuantizeLinear(const std::string& input, + const std::string& output_scale, + const std::string& output_zp, + const std::string& output); + + void Clear(); + + OnnxConverter(const OnnxConverter&); + OnnxConverter& operator=(const OnnxConverter&); +}; + +} // namespace rknpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rknpu/rknpu_execution_provider.cc b/onnxruntime/core/providers/rknpu/rknpu_execution_provider.cc new file mode 100644 index 0000000000..55d0af7192 --- /dev/null +++ b/onnxruntime/core/providers/rknpu/rknpu_execution_provider.cc @@ -0,0 +1,578 @@ +// Copyright 2020 rock-chips.com Inc. + +#include +#include +#include +#include +#include +#include +#include +#include "rknpu_execution_provider.h" +#include "core/common/logging/logging.h" +#include "core/framework/allocatormgr.h" +#include "core/framework/compute_capability.h" +#include "core/session/onnxruntime_cxx_api.h" +#include "core/session/inference_session.h" +#include "core/graph/model.h" +#include "core/framework/memcpy.h" +#include "node_attr_helper.h" +#include "rknpu/rknpu_pub.h" +#include "onnx_converter.h" + +using std::string; +using std::vector; + +namespace onnxruntime { + +constexpr const char* RKNPU = "Rknpu"; + +struct RknpuFuncState { + std::string uniq_input_shape; + + std::unique_ptr exector; + ONNX_NAMESPACE::ModelProto model_proto; + std::unordered_map input_map; + std::unordered_map output_map; + std::vector input_indexes; + std::vector output_indexes; +}; + +RknpuExecutionProvider::RknpuExecutionProvider() + : IExecutionProvider{onnxruntime::kRknpuExecutionProvider} { + auto default_allocator_factory = [](int) { + auto memory_info = onnxruntime::make_unique(RKNPU, OrtAllocatorType::OrtDeviceAllocator); + return onnxruntime::make_unique(std::move(memory_info)); + }; + DeviceAllocatorRegistrationInfo default_memory_info{ + OrtMemTypeDefault, + std::move(default_allocator_factory), + std::numeric_limits::max()}; + InsertAllocator(CreateAllocator(default_memory_info)); + + auto cpu_allocator_factory = [](int) { + auto memory_info = onnxruntime::make_unique( + RKNPU, OrtAllocatorType::OrtDeviceAllocator, OrtDevice(), 0, OrtMemTypeCPUOutput); + return onnxruntime::make_unique(std::move(memory_info)); + }; + DeviceAllocatorRegistrationInfo cpu_memory_info{ + OrtMemTypeCPUOutput, + std::move(cpu_allocator_factory), + std::numeric_limits::max()}; + InsertAllocator(CreateAllocator(cpu_memory_info)); +} + +RknpuExecutionProvider::~RknpuExecutionProvider() {} + +std::vector> RknpuExecutionProvider::GetSupportedNodes( + const ONNX_NAMESPACE::ModelProto& model_proto) const { + rknpu::OnnxConverter converter; + return converter.GetSupportedNodes(model_proto); +} + +std::vector> +RknpuExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_viewer, + const std::vector& /*kernel_registries*/) const { + // Find inputs, initializers and outputs for each supported subgraph + std::vector> result; + + // Handle If and Loop operators + if (graph_viewer.IsSubgraph()) { + return result; + } + + // Need access to model_path_ + for (const auto& tensor : graph_viewer.GetAllInitializedTensors()) { + if (tensor.second->has_data_location() && + tensor.second->data_location() == ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL) { + LOGS_DEFAULT(WARNING) << "Rknpu: Initializers with external data" + " location are not currently supported"; + return result; + } + } + + // This method is based on that of TRT EP + // Construct modelproto from graph_viewer + onnxruntime::Model model(graph_viewer.Name(), true, ModelMetaData(), + PathString(), + IOnnxRuntimeOpSchemaRegistryList(), + graph_viewer.DomainToVersionMap(), + std::vector(), + *GetLogger()); + onnxruntime::Graph& graph_build = model.MainGraph(); + const std::vector& node_index = + graph_viewer.GetNodesInTopologicalOrder(); + std::set all_node_inputs; + for (const auto& node : graph_viewer.Nodes()) { + std::vector inputs, outputs; + for (const auto input : node.InputDefs()) { + auto& n_input = graph_build.GetOrCreateNodeArg( + input->Name(), input->TypeAsProto()); + inputs.push_back(&n_input); + all_node_inputs.insert(&n_input); + } + for (const auto output : node.OutputDefs()) { + auto& n_output = graph_build.GetOrCreateNodeArg( + output->Name(), output->TypeAsProto()); + outputs.push_back(&n_output); + } + graph_build.AddNode(node.Name(), node.OpType(), node.Description(), + inputs, outputs, &node.GetAttributes(), node.Domain()); + } + const auto graph_outputs = graph_viewer.GetOutputs(); + // Add initializer to graph_viewer + const auto& init_tensors = graph_viewer.GetAllInitializedTensors(); + for (const auto& tensor : init_tensors) { + graph_build.AddInitializedTensor(*(tensor.second)); + } + + ORT_ENFORCE(graph_build.Resolve().IsOK()); + ONNX_NAMESPACE::ModelProto model_proto = model.ToProto(); + model_proto.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); + + const auto supported_nodes_vector = GetSupportedNodes(model_proto); + + int counter = 0; + + for (const auto& group : supported_nodes_vector) { + if (!group.empty()) { + std::unordered_set node_set; + node_set.reserve(group.size()); + for (const auto& index : group) { + node_set.insert(node_index[index]); + } + std::unique_ptr sub_graph = + onnxruntime::make_unique(); + // Find inputs and outputs of the subgraph + std::unordered_map + fused_inputs, fused_outputs, fused_outputs_to_add; + std::unordered_set erased; + int input_order = 0; + int output_order = 0; + + for (const auto& index : group) { + sub_graph->nodes.push_back(node_index[index]); + const auto& node = graph_viewer.GetNode(node_index[index]); + + for (const auto& input : node->InputDefs()) { + const auto& it = fused_outputs.find(input); + + if (it != fused_outputs.end()) { + fused_outputs.erase(it); + erased.insert(input); + } + // only when input is neither in output list nor erased list, add the + // input to input list + else if (erased.find(input) == erased.end()) { + fused_inputs[input] = input_order++; + } + } + + // For output searching, there is a special case: + // If node's OutputEdges are more than its outputs, meaning certain + // output is used more than once, + // if the output is connected to nodes that don't belong to the + // subgraph, the output need to be added to the output list + if (node->GetOutputEdgesCount() > node->OutputDefs().size()) { + for (auto it = node->OutputEdgesBegin(), + end = node->OutputEdgesEnd(); + it != end; ++it) { + const auto& node_idx = it->GetNode().Index(); + const auto& output = (it->GetNode()).InputDefs()[it->GetDstArgIndex()]; + + if (node_set.find(node_idx) != node_set.end()) { + const auto& iter = fused_inputs.find(output); + + if (iter != fused_inputs.end()) { + fused_inputs.erase(iter); + erased.insert(output); + } else if (erased.find(output) == erased.end()) { + fused_outputs[output] = output_order++; + } + } else { + fused_outputs_to_add[output] = output_order++; + } + } + } else { + for (const auto& output : node->OutputDefs()) { + const auto& it = fused_inputs.find(output); + + if (it != fused_inputs.end()) { + fused_inputs.erase(it); + erased.insert(output); + } + // only when output is neither in input list nor erased list, + // add the output to output list + else if (erased.find(output) == erased.end()) { + fused_outputs[output] = output_order++; + } + } + } + } + + fused_outputs.insert( + fused_outputs_to_add.begin(), fused_outputs_to_add.end()); + + // Sort inputs and outputs by the order they were added + std::map inputs, outputs; + + for (auto it = fused_inputs.begin(), + end = fused_inputs.end(); + it != end; ++it) { + inputs.insert(std::pair(it->second, it->first)); + } + + for (auto it = fused_outputs.begin(), + end = fused_outputs.end(); + it != end; ++it) { + for (const auto& x : all_node_inputs) { + if (x->Name() == it->first->Name()) { + outputs.insert( + std::pair(it->second, it->first)); + break; + } + } + if (std::find(graph_outputs.begin(), + graph_outputs.end(), it->first) != graph_outputs.end()) { + outputs.insert(std::pair(it->second, it->first)); + } + } + + // Assign inputs and outputs to subgraph's meta_def + auto meta_def = + onnxruntime::make_unique<::onnxruntime::IndexedSubGraph::MetaDef>(); + meta_def->name = "RKNPU_" + std::to_string(counter++); + meta_def->domain = kMSDomain; + + for (const auto& input : inputs) { + meta_def->inputs.push_back(input.second->Name()); + } + + for (const auto& output : outputs) { + meta_def->outputs.push_back(output.second->Name()); + } + + meta_def->since_version = 1; + sub_graph->SetMetaDef(meta_def); + + result.push_back( + onnxruntime::make_unique(std::move(sub_graph))); + } + } + + return result; +} + +common::Status RknpuExecutionProvider::Compile( + const std::vector& fused_nodes, + std::vector& node_compute_funcs) { + for (const auto* fused_node : fused_nodes) { + // Reconstruct graph proto from fused node's function body + const auto* func_body = fused_node->GetFunctionBody(); + if (!func_body) { + return common::Status(common::ONNXRUNTIME, + common::INVALID_ARGUMENT, "Function body is empty"); + } + const Graph& graph_body = func_body->Body(); + onnxruntime::Model model(graph_body.Name(), true, ModelMetaData(), + PathString(), + IOnnxRuntimeOpSchemaRegistryList(), + graph_body.DomainToVersionMap(), + std::vector(), + *GetLogger()); + ONNX_NAMESPACE::ModelProto model_proto = model.ToProto(); + *(model_proto.mutable_graph()) = graph_body.ToGraphProto(); + model_proto.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); + + // Build map from input name to its index in input definitions + std::unordered_map input_map; + const auto& input_defs = fused_node->InputDefs(); + input_map.reserve(input_defs.size()); + for (int i = 0, end = input_defs.size(); i < end; ++i) { + input_map[input_defs[i]->Name()] = i; + } + + // Build map from output name to its index in output definitions + std::unordered_map output_map; + const auto& output_defs = fused_node->OutputDefs(); + output_map.reserve(output_defs.size()); + for (int i = 0, end = output_defs.size(); i < end; ++i) { + output_map[output_defs[i]->Name()] = i; + } + + model_proto_[fused_node->Name()] = model_proto; + input_info_[fused_node->Name()] = input_map; + output_info_[fused_node->Name()] = output_map; + + NodeComputeInfo compute_info; + compute_info.create_state_func = [&](ComputeContext* context, + FunctionState* state) { + std::unique_ptr p = + onnxruntime::make_unique(); + rk::nn::Graph* graph = new rk::nn::Graph(); + *p = {"", std::unique_ptr(new rk::nn::Exection(graph)), + model_proto_[context->node_name], input_info_[context->node_name], + output_info_[context->node_name], + std::vector{}, std::vector{}}; + *state = p.release(); + return 0; + }; + + compute_info.release_state_func = [](FunctionState state) { + if (state) { + RknpuFuncState* p = static_cast(state); + rk::nn::Graph* graph = p->exector->GetGraph(); + delete graph; + delete p; + } + }; + + compute_info.compute_func = [](FunctionState state, + const OrtCustomOpApi* api, + OrtKernelContext* context) { + Ort::CustomOpApi ort{*api}; + RknpuFuncState* rk_state = reinterpret_cast(state); + const size_t n_inputs = ort.KernelContext_GetInputCount(context); + const size_t n_outputs = ort.KernelContext_GetOutputCount(context); + + std::string input_shape; + input_shape.reserve(4 * sizeof(int64_t) * n_inputs + n_inputs); + for (size_t i = 0; i < n_inputs; i++) { + const OrtValue* input_tensor = ort.KernelContext_GetInput(context, i); + auto tensor_info = ort.GetTensorTypeAndShape(input_tensor); + auto tensor_shape = ort.GetTensorShape(tensor_info); + ort.ReleaseTensorTypeAndShapeInfo(tensor_info); + + const auto ndim = tensor_shape.size(); + input_shape.append(reinterpret_cast(&ndim), sizeof(ndim)); + input_shape.append(reinterpret_cast(tensor_shape.data()), + ndim * sizeof(int64_t)); + } + + bool rebuild = false; + std::vector input_bufs(n_inputs); + if (rk_state->uniq_input_shape == "") { + auto graph_proto = rk_state->model_proto.mutable_graph(); + for (size_t i = 0; i < n_inputs; i++) { + const OrtValue* input_tensor = ort.KernelContext_GetInput(context, i); + input_bufs[i] = const_cast( + ort.GetTensorData(input_tensor)); + auto tensor_info = ort.GetTensorTypeAndShape(input_tensor); + auto tensor_shape = ort.GetTensorShape(tensor_info); + ort.ReleaseTensorTypeAndShapeInfo(tensor_info); + + auto g_in_shape = + graph_proto->mutable_input((int)i)->mutable_type()->mutable_tensor_type()->mutable_shape(); + g_in_shape->clear_dim(); + for (size_t dim = 0; dim < tensor_shape.size(); dim++) { + g_in_shape->add_dim()->set_dim_value(tensor_shape[dim]); + } + } + rk_state->uniq_input_shape = input_shape; + rebuild = true; + } else if (rk_state->uniq_input_shape != input_shape) { + // TODO + throw std::invalid_argument( + "The input_shape is not match" + " the rk_state->uniq_input_shape!"); + } else { + LOGS_DEFAULT(INFO) << "input_shape equal to rk_state->uniq_input_shape," + " skip rebuild!"; + } + + rk::nn::Graph* graph = rk_state->exector->GetGraph(); + if (rebuild) { + rknpu::OnnxConverter converter; + converter.Convert(rk_state->model_proto, graph, input_bufs, rk_state->input_map); + + rk_state->exector->Build(); + + auto input_map = rk_state->input_map; + auto output_map = rk_state->output_map; + for (auto it = output_map.begin(); it != output_map.end(); it++) { + if (converter.m(it->first) != it->first) + output_map[converter.m(it->first)] = output_map[it->first]; + } + + int n_inputs = graph->GetInputs().size(); + rk_state->input_indexes.resize(n_inputs); + for (int i = 0; i < n_inputs; ++i) { + const auto input = graph->GetInputs()[i]; + const std::string& name = input->GetName(); + auto iter = input_map.find(name); + if (iter != input_map.end()) { + rk_state->input_indexes[i] = iter->second; + } + } + + int n_outputs = graph->GetOutputs().size(); + rk_state->output_indexes.resize(n_outputs); + for (int i = 0; i < n_outputs; ++i) { + const auto output = graph->GetOutputs()[i]; + const std::string& name = output->GetName(); + auto iter = output_map.find(name); + if (iter != output_map.end()) { + rk_state->output_indexes[i] = iter->second; + } + } + } + + ORT_ENFORCE(graph->GetInputs().size() <= n_inputs, + "Inconsistent input sizes"); + ORT_ENFORCE(graph->GetOutputs().size() == n_outputs, + "Inconsistent output sizes"); + + std::vector inputs; + inputs.resize(graph->GetInputs().size()); + for (size_t i = 0; i < graph->GetInputs().size(); i++) { + const OrtValue* input_tensor = + ort.KernelContext_GetInput(context, rk_state->input_indexes[i]); + float* input_buf = + const_cast(ort.GetTensorData(input_tensor)); + + const auto tensor_info = ort.GetTensorTypeAndShape(input_tensor); + const auto& tensor_shape = ort.GetTensorShape(tensor_info); + + auto type = graph->GetInputs()[i]->GetPrecision(); + size_t type_size = 4; + switch (type) { + case rk::nn::PrecisionType::FLOAT32: + type_size = 4; + break; + case rk::nn::PrecisionType::UINT8: + type_size = 1; + break; + case rk::nn::PrecisionType::INT32: + type_size = 4; + break; + case rk::nn::PrecisionType::INT64: + type_size = 8; + break; + default: + // TODO + throw std::invalid_argument( + "compute_func: unknow input data type!"); + break; + } + + inputs[i].index = i; + inputs[i].buf = (void*)input_buf; + inputs[i].size = std::accumulate(tensor_shape.begin(), + tensor_shape.end(), 1, + std::multiplies()) * + type_size; + inputs[i].pass_through = false; + inputs[i].type = type; + inputs[i].layout = rk::nn::DataLayoutType::NCHW; + } + + std::vector outputs; + outputs.resize(n_outputs); + for (size_t i = 0; i < n_outputs; i++) { + const auto output = graph->GetOutputs()[i]; + const auto output_shape = output->GetDims(); + std::vector + int64_output_shape(output_shape.begin(), output_shape.end()); + const auto* output_tensor = ort.KernelContext_GetOutput( + context, rk_state->output_indexes[i], + int64_output_shape.data(), + int64_output_shape.size()); + float* output_buf = + const_cast(ort.GetTensorData(output_tensor)); + + const auto type = output->GetPrecision(); + size_t type_size = 4; + switch (type) { + case rk::nn::PrecisionType::FLOAT32: + type_size = 4; + break; + case rk::nn::PrecisionType::UINT8: + type_size = 1; + break; + case rk::nn::PrecisionType::INT32: + type_size = 4; + break; + case rk::nn::PrecisionType::INT64: + type_size = 8; + break; + default: + // TODO + throw std::invalid_argument( + "compute_func: unknow output data type!"); + break; + } + + outputs[i].index = i; + outputs[i].buf = (void*)output_buf; + outputs[i].size = accumulate(output_shape.begin(), + output_shape.end(), 1, + std::multiplies()) * + type_size; + outputs[i].type = type; + outputs[i].layout = rk::nn::DataLayoutType::NCHW; + outputs[i].want_float = false; + } + + rk_state->exector->SetInputs(inputs); + rk_state->exector->Run(); + rk_state->exector->GetOutputs(outputs); + + return Status::OK(); + }; + + node_compute_funcs.push_back(compute_info); + } + return Status::OK(); +} + +ONNX_OPERATOR_KERNEL_EX( + MemcpyFromHost, + kOnnxDomain, + 1, + kRknpuExecutionProvider, + KernelDefBuilder() + .InputMemoryType(0) + .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), + Memcpy); + +ONNX_OPERATOR_KERNEL_EX( + MemcpyToHost, + kOnnxDomain, + 1, + kRknpuExecutionProvider, + KernelDefBuilder() + .OutputMemoryType(0) + .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), + Memcpy); + +class ONNX_OPERATOR_KERNEL_CLASS_NAME( + kRknpuExecutionProvider, kOnnxDomain, 1, MemcpyFromHost); +class ONNX_OPERATOR_KERNEL_CLASS_NAME( + kRknpuExecutionProvider, kOnnxDomain, 1, MemcpyToHost); + +static void RegisterRknpuKernels(KernelRegistry& kernel_registry) { + static const BuildKernelCreateInfoFn function_table[] = { + BuildKernelCreateInfo, + BuildKernelCreateInfo, + }; + + for (auto& function_table_entry : function_table) { + kernel_registry.Register(function_table_entry()); + } +} + +std::shared_ptr GetRknpuKernelRegistry() { + std::shared_ptr kernel_registry = + std::make_shared(); + RegisterRknpuKernels(*kernel_registry); + + return kernel_registry; +} + +std::shared_ptr +RknpuExecutionProvider::GetKernelRegistry() const { + static std::shared_ptr kernel_registry = + onnxruntime::GetRknpuKernelRegistry(); + return kernel_registry; +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rknpu/rknpu_execution_provider.h b/onnxruntime/core/providers/rknpu/rknpu_execution_provider.h new file mode 100644 index 0000000000..1316af5951 --- /dev/null +++ b/onnxruntime/core/providers/rknpu/rknpu_execution_provider.h @@ -0,0 +1,34 @@ +// Copyright 2020 rock-chips.com Inc. + +#pragma once + +#include +#include +#include +#include + +#include "core/framework/execution_provider.h" +#include "core/graph/onnx_protobuf.h" + +namespace onnxruntime { + +class RknpuExecutionProvider : public IExecutionProvider { + public: + RknpuExecutionProvider(); + virtual ~RknpuExecutionProvider(); + + std::vector> + GetCapability(const onnxruntime::GraphViewer& graph, + const std::vector& /*kernel_registries*/) const override; + common::Status Compile(const std::vector& fused_nodes, + std::vector& node_compute_funcs) override; + std::shared_ptr GetKernelRegistry() const override; + + private: + std::unordered_map model_proto_; + std::unordered_map> input_info_; + std::unordered_map> output_info_; + std::vector> GetSupportedNodes( + const ONNX_NAMESPACE::ModelProto& model_proto) const; +}; +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rknpu/rknpu_provider_factory.cc b/onnxruntime/core/providers/rknpu/rknpu_provider_factory.cc new file mode 100644 index 0000000000..f7ff3b8fc4 --- /dev/null +++ b/onnxruntime/core/providers/rknpu/rknpu_provider_factory.cc @@ -0,0 +1,33 @@ +// Copyright 2020 rock-chips.com Inc. + +#include "core/providers/rknpu/rknpu_provider_factory.h" +#include "rknpu_execution_provider.h" +#include "core/session/abi_session_options_impl.h" + +using namespace onnxruntime; + +namespace onnxruntime { + +struct RknpuProviderFactory : IExecutionProviderFactory { + RknpuProviderFactory() {} + ~RknpuProviderFactory() override {} + + std::unique_ptr CreateProvider() override; +}; + +std::unique_ptr RknpuProviderFactory::CreateProvider() { + return onnxruntime::make_unique(); +} + +std::shared_ptr +CreateExecutionProviderFactory_Rknpu() { + return std::make_shared(); +} +} // namespace onnxruntime + +ORT_API_STATUS_IMPL(OrtSessionOptionsAppendExecutionProvider_Rknpu, + _In_ OrtSessionOptions* options) { + options->provider_factories.push_back( + onnxruntime::CreateExecutionProviderFactory_Rknpu()); + return nullptr; +} diff --git a/onnxruntime/core/providers/rknpu/shaper.cc b/onnxruntime/core/providers/rknpu/shaper.cc new file mode 100644 index 0000000000..18144aa248 --- /dev/null +++ b/onnxruntime/core/providers/rknpu/shaper.cc @@ -0,0 +1,561 @@ +#include "shaper.h" +#include +#include +#include + +namespace onnxruntime { +namespace rknpu { + +template +T Product(const std::vector& v) { + return static_cast( + accumulate(v.begin(), v.end(), 1, std::multiplies())); +} + +Shaper::len_t Shaper::total(const Shape& shape) { + return Product(shape); +} + +/** + * strides: [stride_y, stride_x] + * paddings: [top, left, bottom, right] + */ +void Shaper::Conv(const std::string& input, + const std::string& weight, + const std::vector& paddings, + const std::vector& strides, + const std::string& auto_pad, + const std::string& output) { + Shaper::Conv(input, strides[1], strides[0], 1, 1, paddings[1], + paddings[3], paddings[0], paddings[2], weight, auto_pad, + output); +} + +void Shaper::Conv(const std::string& input, + const std::vector& paddings, + const std::vector& strides, + const std::vector& dilations, + const std::string& weight, + const std::string& auto_pad, + const std::string& output) { + Shaper::Conv(input, strides[1], strides[0], dilations[1], dilations[0], + paddings[1], paddings[3], paddings[0], paddings[2], + weight, auto_pad, output); +} +void Shaper::Conv(const std::string& input, + const std::string& weight, + const int32_t padding_left, + const int32_t padding_right, + const int32_t padding_top, + const int32_t padding_bottom, + const int32_t stride_x, + const int32_t stride_y, + const std::string& auto_pad, + const std::string& output) { + Shaper::Conv(input, stride_x, stride_y, 1, 1, padding_left, padding_right, + padding_top, padding_bottom, weight, auto_pad, output); +} + +void Shaper::Conv(const std::string& input, + const int32_t strideX, + const int32_t strideY, + const int32_t dilationX, + const int32_t dilationY, + const int32_t paddingLeft, + const int32_t paddingRight, + const int32_t paddingTop, + const int32_t paddingBottom, + const std::string& weight, + const std::string& auto_pad, + const std::string& output) { + Shape weightDimen = + shape_map_.at(weight); // num_output, height, width, num_input + // NCHW + Shape inputDimen = shape_map_.at(input); + + if (auto_pad == "VALID") { + Shape outputDimen{inputDimen[0], + weightDimen[0], + (inputDimen[2] - ((weightDimen[2] - 1) * dilationY + 1)) / strideY + 1, + (inputDimen[3] - ((weightDimen[3] - 1) * dilationX + 1)) / strideX + 1}; + shape_map_[output] = outputDimen; + } else if (auto_pad == "SAME_UPPER" || auto_pad == "SAME_LOWER") { + int32_t legacy_target_size_Y = (inputDimen[2] + strideY - 1) / strideY; + int32_t legacy_target_size_X = (inputDimen[3] + strideX - 1) / strideX; + int32_t pad_needed_Y = + (legacy_target_size_Y - 1) * strideY + weightDimen[2] - inputDimen[2]; + int32_t pad_needed_X = + (legacy_target_size_X - 1) * strideX + weightDimen[3] - inputDimen[3]; + Shape outputDimen{inputDimen[0], + weightDimen[0], + (inputDimen[2] - ((weightDimen[2] - 1) * dilationY + 1) + pad_needed_Y) / strideY + 1, + (inputDimen[3] - ((weightDimen[3] - 1) * dilationX + 1) + pad_needed_X) / strideX + 1}; + shape_map_[output] = outputDimen; + } else { // default: NOTSET + Shape outputDimen{inputDimen[0], + weightDimen[0], + (inputDimen[2] - ((weightDimen[2] - 1) * dilationY + 1) + paddingTop + paddingBottom) / strideY + 1, + (inputDimen[3] - ((weightDimen[3] - 1) * dilationX + 1) + paddingLeft + paddingRight) / strideX + 1}; + shape_map_[output] = outputDimen; + } +} + +void Shaper::DepthwiseConv(const std::string& input, + const std::vector& paddings, + const std::vector& strides, + const std::vector& dilations, + const std::string& weight, + const std::string& output) { + Shaper::DepthwiseConv(input, strides[1], strides[0], dilations[1], + dilations[0], paddings[1], paddings[3], paddings[0], + paddings[2], weight, output); +} + +void Shaper::DepthwiseConv(const std::string& input, + const std::string& weight, + const int32_t padding_left, + const int32_t padding_right, + const int32_t padding_top, + const int32_t padding_bottom, + const int32_t stride_x, + const int32_t stride_y, + const std::string& output) { + DepthwiseConv(input, stride_x, stride_y, 1, 1, padding_left, + padding_right, padding_top, padding_bottom, weight, + output); +} + +void Shaper::DepthwiseConv(const std::string& input, + const int32_t strideX, + const int32_t strideY, + const int32_t dilationX, + const int32_t dilationY, + const int32_t paddingLeft, + const int32_t paddingRight, + const int32_t paddingTop, + const int32_t paddingBottom, + const std::string& weight, + const std::string& output) { + Shape weightDimen = + shape_map_.at(weight); // 1, height, width, num_output + // NCHW + Shape inputDimen = shape_map_.at(input); + Shape outputDimen{ + inputDimen[0], + weightDimen[0], + (inputDimen[2] - ((weightDimen[2] - 1) * dilationY + 1) + + paddingTop + paddingBottom) / + strideY + + 1, + (inputDimen[3] - ((weightDimen[3] - 1) * dilationX + 1) + + paddingLeft + paddingRight) / + strideX + + 1, + }; + shape_map_[output] = outputDimen; +} + +void Shaper::DepthwiseConv(const std::string& input, + const std::string& weight, + const std::vector& paddings, + const std::vector& strides, + const std::string& output) { + DepthwiseConv(input, weight, paddings[1], paddings[3], + paddings[0], paddings[2], strides[1], strides[0], + output); +} + +void Shaper::Slice(const std::string& input, + const std::vector& starts, + const std::vector& ends, + const std::vector& axes, + const std::vector& steps, + const std::string& output) { + std::vector inputDimen = shape_map_.at(input); + std::vector outputDimen = inputDimen; + for (size_t i = 0; i < axes.size(); i++) { + int32_t axis = + (axes[i] < 0) ? (axes[i] + (int32_t)inputDimen.size()) : axes[i]; + int32_t dim = outputDimen[axis]; + if (dim > 0) { + int32_t start = starts[i] < 0 ? (starts[i] + dim) : starts[i]; + int32_t end = ends[i] < 0 ? (ends[i] + dim) : ends[i]; + start = std::max(start, 0); + end = std::max(end, 0); + end = std::min(end, dim); + outputDimen[axis] = end - start; + } + } + + shape_map_[output] = outputDimen; +} + +void Shaper::StridedSlice(const std::string& input, + const std::vector& starts, + const std::vector& ends, + const std::vector& strides, + const int32_t beginMask, + const int32_t endMask, + const int32_t shrinkAxisMask, + const std::string& output) { + // NHWC + std::vector inputDimen = shape_map_.at(input); + std::vector outputDimen; + for (size_t i = 0; i < inputDimen.size(); i++) { + if (shrinkAxisMask & (1 << i)) { + continue; + } + int32_t start = starts[i], end = ends[i], stride = strides[i]; + if (beginMask & (1 << i)) { + start = 0; + } + if (endMask & (1 << i)) { + end = inputDimen[i]; + } + outputDimen.emplace_back((end - start) / stride); + } + shape_map_[output] = outputDimen; +} + +void Shaper::Gather(const std::string& input, + const std::string& indices, + const int32_t axis, + const std::string& output) { + auto inputDimen = shape_map_.at(input); + auto indicesDimen = shape_map_.at(indices); + int32_t input_rank = indicesDimen.size(); + int32_t axis_new = (axis < 0) ? (axis + input_rank) : axis; + + std::vector outputDimen; + outputDimen.reserve(input_rank - 1 + indicesDimen.size()); + + // replace the dimension for axis with the shape from the indices + for (int32_t i = 0; i < axis_new; ++i) + outputDimen.push_back(inputDimen[i]); + + for (const auto dim : indicesDimen) + outputDimen.push_back(dim); + + for (int32_t i = axis_new + 1; i < input_rank; ++i) + outputDimen.push_back(inputDimen[i]); + + shape_map_[output] = outputDimen; +} + +void Shaper::Pool(const std::string& input, + const int32_t padding_left, + const int32_t padding_right, + const int32_t padding_top, + const int32_t padding_bottom, + const int32_t stride_x, + const int32_t stride_y, + const int32_t width, + const int32_t height, + const std::string& output) { + auto inputDimen = shape_map_.at(input); + + // NCHW + Shape outputDimen; + if (height == -1 && width == -1) { + outputDimen = {inputDimen[0], inputDimen[1], 1, 1}; + } else { + outputDimen = { + inputDimen[0], + inputDimen[1], + (inputDimen[2] - height + padding_top + padding_bottom) / stride_y + + 1, + (inputDimen[3] - width + padding_left + padding_right) / stride_x + + 1}; + } + shape_map_[output] = outputDimen; +} + +/** + * kernel_shape: [height, width] + * strides: [stride_y, stride_x] + * pads: [top, left, bottom, right] + */ +void Shaper::Pool(const std::string& input, + const std::vector& kernel_shape, + const std::vector& pads, + const std::vector& strides, + const std::string& output) { + Shaper::Pool(input, pads[1], pads[3], pads[0], pads[2], strides[1], + strides[0], kernel_shape[1], kernel_shape[0], output); +} + +void Shaper::Softmax(const std::string& input, + const std::string& output) { + shape_map_[output] = shape_map_.at(input); +} + +void Shaper::Relu(const std::string& input, + const std::string& output) { + shape_map_[output] = shape_map_.at(input); +} + +void Shaper::Concat(const std::vector& inputs, + const int32_t axis, + const std::string& output) { + std::vector dimens; + for (const auto& input : inputs) { + auto& dimen = shape_map_.at(input); + if (!dimens.empty()) { + for (size_t i = 0; i < dimens[0].size(); i++) { + if ((int32_t)i == axis) continue; + if (dimen[i] != dimens[0][i]) { + throw std::invalid_argument("Wrong input for concat"); + } + } + } + dimens.push_back(shape_map_.at(input)); + } + + auto outputDimen = dimens[0]; + for (size_t i = 1; i < dimens.size(); i++) { + outputDimen[axis] += dimens[i][axis]; + } + shape_map_[output] = outputDimen; +} + +void Shaper::LRN(const std::string& input, + const std::string& output) { + shape_map_[output] = shape_map_.at(input); +} + +void Shaper::FC(const std::string& input, + const std::string& weight, + const std::string& output) { + Shape weightDimen = shape_map_.at(weight); // num_units, input_size + auto input_dimen = shape_map_.at(input); + Shape outputDimen{input_dimen[0], weightDimen[0]}; + shape_map_[output] = outputDimen; +} + +void Shaper::Eltwise(const std::string& input1, + const std::string& input2, + const std::string& output) { + auto shape1 = shape_map_.at(input1); + auto shape2 = shape_map_.at(input2); + + // broadcasting support + auto max_shape = shape1.size() >= shape2.size() ? shape1 : shape2; + auto min_shape = shape1.size() < shape2.size() ? shape1 : shape2; + for (int i = (int)max_shape.size() - 1, j = (int)min_shape.size() - 1; + i >= 0 && j >= 0; + i--, j--) { + if (max_shape[i] < min_shape[j]) + max_shape[i] = min_shape[j]; + } + + shape_map_[output] = max_shape; +} + +void Shaper::Eltwise(const std::string& input1, + const std::string& output) { + shape_map_[output] = shape_map_.at(input1); +} + +void Shaper::BatchToSpace(const std::string& input, + const std::vector& block_sizes, + const std::string& output) { + auto input_dimen = shape_map_.at(input); + auto output_dimen = {input_dimen[0] / Product(block_sizes), + input_dimen[1] * block_sizes[0], + input_dimen[2] * block_sizes[1], input_dimen[3]}; + shape_map_[output] = output_dimen; +} + +void Shaper::SpaceToBatch(const std::string& input, + const std::vector& block_sizes, + const std::vector& pads, + const std::string& output) { + auto input_dimen = shape_map_.at(input); + auto output_dimen = {input_dimen[0] * Product(block_sizes), + (input_dimen[1] + pads[0] + pads[1]) / block_sizes[0], + (input_dimen[2] + pads[2] + pads[3]) / block_sizes[1], + input_dimen[3]}; + shape_map_[output] = output_dimen; +} + +void Shaper::BatchNorm(const std::string& input, + const std::string& output) { + shape_map_[output] = shape_map_.at(input); +} + +void Shaper::Reshape(const std::string& input, + const std::vector& shape, + const std::string& output) { + auto input_dimen = shape_map_.at(input); + int64_t input_size = std::accumulate( + input_dimen.begin(), input_dimen.end(), 1, std::multiplies()); + std::vector output_dimen(shape.size()); + + int64_t capacity = 1; + int unk_dim_idx = -1; + for (size_t i = 0; i < shape.size(); i++) { + if (shape[i] == -1) { + if (unk_dim_idx != -1) + throw std::invalid_argument( + "Only one input dimension of Attr(shape) can be unknown!"); + unk_dim_idx = i; + } else if (shape[i] == 0) { + if (i >= input_dimen.size()) + throw std::invalid_argument( + "The index of dimension to copy from input" + " shape must be less than the size of input shape!"); + } else { + if (shape[i] < 0) + throw std::invalid_argument( + "Each input dimension of Attr(shape) must" + " not be negtive except one unknown dimension!"); + } + int32_t output_dim_i = shape[i] ? shape[i] : input_dimen[i]; + output_dimen[i] = output_dim_i; + capacity *= output_dim_i; + } + + if (unk_dim_idx != -1) { + output_dimen[unk_dim_idx] = -input_size / capacity; + if ((output_dimen[unk_dim_idx] * capacity) != (-input_size)) + throw std::invalid_argument("Invalid shape is given!"); + } else { + if (capacity != input_size) + throw std::invalid_argument("Invalid shape is given!"); + } + + Shape final_dimen(shape.size()); + for (size_t i = 0; i < shape.size(); i++) { + final_dimen[i] = (uint32_t)output_dimen[i]; + } + shape_map_[output] = final_dimen; +} + +void Shaper::Transpose(const std::string& input, + const std::vector& perm, + const std::string& output) { + auto input_dimen = shape_map_.at(input); + + Shape outputDimen(perm.size()); + + for (size_t i = 0; i < perm.size(); i++) { + outputDimen[i] = input_dimen[perm[i]]; + } + + shape_map_[output] = outputDimen; +} + +void Shaper::Squeeze(const std::string& input, + const std::vector& axes, + const std::string& output) { + std::vector inputDimen = shape_map_.at(input); + size_t n_axes = axes.size(); + int cnt_squeezed_dims = 0; + bool should_squeeze[9] = {false}; + + if (n_axes == 0) { + for (size_t idx = 0; idx < inputDimen.size(); ++idx) { + if (inputDimen[idx] == 1) { + should_squeeze[idx] = true; + ++cnt_squeezed_dims; + } + } + } else { + for (size_t idx = 0; idx < n_axes; ++idx) { + int32_t current = + axes[idx] < 0 ? axes[idx] + inputDimen.size() : axes[idx]; + if (!(should_squeeze[current])) { + ++cnt_squeezed_dims; + } + should_squeeze[current] = true; + } + } + + // Make output dimensions + std::vector outputDimen(inputDimen.size() - cnt_squeezed_dims, 0); + for (size_t in_idx = 0, out_idx = 0; in_idx < inputDimen.size(); ++in_idx) { + if (!should_squeeze[in_idx]) { + outputDimen[out_idx++] = inputDimen[in_idx]; + } + } + + shape_map_[output] = outputDimen; +} + +void Shaper::Unsqueeze(const std::string& input, + const std::vector& axes, + const std::string& output) { + std::vector inputDimen = shape_map_.at(input); + + int output_size = inputDimen.size() + axes.size(); + int cur_output_size = inputDimen.size(); + std::vector outputDimen(output_size, 0); + + for (int axis : axes) { + int cur = axis; + + // Move old axis, and insert new axis. + for (int i = cur_output_size; i >= cur; --i) { + if (outputDimen[i] == 1) { + // Move axis + outputDimen[i + 1] = 1; + outputDimen[i] = 0; + } + } + + outputDimen[cur] = 1; + // Add the output size. + cur_output_size++; + } + + // Make output shape + for (int in_idx = 0, out_idx = 0; out_idx < output_size; ++out_idx) { + if (outputDimen[out_idx] == 0) { + outputDimen[out_idx] = inputDimen[in_idx++]; + } + } + + shape_map_[output] = outputDimen; +} + +void Shaper::Affine(const std::string& input, + const std::string& output) { + shape_map_[output] = shape_map_.at(input); +} +void Shaper::Affine(const std::string& input, + const std::string& a, + const std::string& b, + const std::string& output) { + (void)a; + (void)b; + Shaper::Affine(input, output); +} + +void Shaper::Identity(const std::string& input, + const std::string& output) { + shape_map_[output] = shape_map_.at(input); +} + +void Shaper::AddShape(const std::string& name, + const Shape& shape) { + shape_map_[name] = shape; +} + +size_t Shaper::GetSize(const std::string& name) { + return static_cast(Product(shape_map_.at(name))); +} + +void Shaper::Clear() { + shape_map_.clear(); +} + +std::ostream& operator<<(std::ostream& os, + const Shaper& shaper) { + for (const auto& p : shaper.shape_map_) { + os << (p.first + ": ") /*<< p.second*/ << std::endl; + } + return os; +} + +} // namespace rknpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rknpu/shaper.h b/onnxruntime/core/providers/rknpu/shaper.h new file mode 100644 index 0000000000..ac093e1aa1 --- /dev/null +++ b/onnxruntime/core/providers/rknpu/shaper.h @@ -0,0 +1,183 @@ +#pragma once + +#include +#include +#include + +#include + +namespace onnxruntime { +namespace rknpu { + +/** + * Help to caculate the output shape of nodes. + */ +class Shaper { + public: + using len_t = uint32_t; + using Shape = std::vector; + + static len_t total(const Shape& shape); + + void Conv(const std::string& input, + const std::string& weight, + const std::vector& strides, + const std::vector& paddings, + const std::string& auto_pad, + const std::string& output); + void Conv(const std::string& input, + const std::vector& strides, + const std::vector& dilations, + const std::vector& paddings, + const std::string& weight, + const std::string& auto_pad, + const std::string& output); + void Conv(const std::string& input, + const std::string& weight, + const int32_t padding_left, + const int32_t padding_right, + const int32_t padding_top, + const int32_t padding_bottom, + const int32_t stride_x, + const int32_t stride_y, + const std::string& auto_pad, + const std::string& output); + void Conv(const std::string& input, + const int32_t strideX, + const int32_t strideY, + const int32_t dilationX, + const int32_t dilationY, + const int32_t paddingLeft, + const int32_t paddingRight, + const int32_t paddingTop, + const int32_t paddingBottom, + const std::string& weight, + const std::string& auto_pad, + const std::string& output); + void DepthwiseConv(const std::string& input, + const std::vector& strides, + const std::vector& dilations, + const std::vector& paddings, + const std::string& weight, + const std::string& output); + void DepthwiseConv(const std::string& input, + const std::string& weight, + const int32_t padding_left, + const int32_t padding_right, + const int32_t padding_top, + const int32_t padding_bottom, + const int32_t stride_x, + const int32_t stride_y, + const std::string& output); + void DepthwiseConv(const std::string& input, + const int32_t strideX, + const int32_t strideY, + const int32_t dilationX, + const int32_t dilationY, + const int32_t paddingLeft, + const int32_t paddingRight, + const int32_t paddingTop, + const int32_t paddingBottom, + const std::string& weight, + const std::string& output); + void DepthwiseConv(const std::string& input, + const std::string& weight, + const std::vector& paddings, + const std::vector& strides, + const std::string& output); + void Slice(const std::string& input, + const std::vector& starts, + const std::vector& ends, + const std::vector& axes, + const std::vector& steps, + const std::string& output); + void StridedSlice(const std::string& input, + const std::vector& starts, + const std::vector& ends, + const std::vector& strides, + const int32_t beginMask, + const int32_t endMask, + const int32_t shrinkAxisMask, + const std::string& output); + void Gather(const std::string& input, + const std::string& indices, + const int32_t axis, + const std::string& output); + void Pool(const std::string& input, + const int32_t padding_left, + const int32_t padding_right, + const int32_t padding_top, + const int32_t padding_bottom, + const int32_t stride_x, + const int32_t stride_y, + const int32_t width, + const int32_t height, + const std::string& output); + void Pool(const std::string& input, + const std::vector& kernel_shape, + const std::vector& pads, + const std::vector& strides, + const std::string& output); + void Softmax(const std::string& input, + const std::string& output); + void Relu(const std::string& input, + const std::string& output); + void Concat(const std::vector& inputs, + const int32_t axis, + const std::string& output); + void LRN(const std::string& input, + const std::string& output); + void FC(const std::string& input, + const std::string& weight, + const std::string& output); + void Eltwise(const std::string& input1, + const std::string& input2, + const std::string& output); + void Eltwise(const std::string& input1, + const std::string& output); + void Affine(const std::string& input, + const std::string& output); + void Affine(const std::string& input, + const std::string& a, + const std::string& b, + const std::string& output); + void Identity(const std::string& input, + const std::string& output); + void BatchToSpace(const std::string& input, + const std::vector& block_sizes, + const std::string& output); + void SpaceToBatch(const std::string& input, + const std::vector& block_sizes, + const std::vector& pads, + const std::string& output); + void BatchNorm(const std::string& input, + const std::string& output); + void Reshape(const std::string& input, + const std::vector& shape, + const std::string& output); + void Transpose(const std::string& input, + const std::vector& perm, + const std::string& output); + void Squeeze(const std::string& input, + const std::vector& axes, + const std::string& output); + void Unsqueeze(const std::string& input, + const std::vector& axes, + const std::string& output); + void AddShape(const std::string& name, + const Shape& shape); + size_t GetSize(const std::string& name); + void Clear(); + + inline const Shape& operator[](const std::string& key) { + return shape_map_[key]; + } + friend std::ostream& operator<<(std::ostream& os, + const Shaper& shaper); + + private: + std::map shape_map_; +}; + +} // namespace rknpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rknpu/symbols.txt b/onnxruntime/core/providers/rknpu/symbols.txt new file mode 100644 index 0000000000..ae6eb6a4db --- /dev/null +++ b/onnxruntime/core/providers/rknpu/symbols.txt @@ -0,0 +1 @@ +OrtSessionOptionsAppendExecutionProvider_Rknpu diff --git a/onnxruntime/test/framework/test_utils.cc b/onnxruntime/test/framework/test_utils.cc index 053d75f230..e9186dcffc 100644 --- a/onnxruntime/test/framework/test_utils.cc +++ b/onnxruntime/test/framework/test_utils.cc @@ -43,6 +43,13 @@ IExecutionProvider* TestNnapiExecutionProvider() { } #endif +#ifdef USE_RKNPU +IExecutionProvider* TestRknpuExecutionProvider() { + static RknpuExecutionProvider rknpu_provider; + return &rknpu_provider; +} +#endif + static void CountOpsInGraphImpl( const Graph& graph, bool recurse_into_subgraphs, std::map& ops) { for (auto& node : graph.Nodes()) { diff --git a/onnxruntime/test/framework/test_utils.h b/onnxruntime/test/framework/test_utils.h index 03fc286ee1..d7dd37f83a 100644 --- a/onnxruntime/test/framework/test_utils.h +++ b/onnxruntime/test/framework/test_utils.h @@ -24,6 +24,9 @@ #ifdef USE_NNAPI #include "core/providers/nnapi/nnapi_execution_provider.h" #endif +#ifdef USE_RKNPU +#include "core/providers/rknpu/rknpu_execution_provider.h" +#endif namespace onnxruntime { class Graph; @@ -50,6 +53,10 @@ IExecutionProvider* TestOpenVINOExecutionProvider(); IExecutionProvider* TestNnapiExecutionProvider(); #endif +#ifdef USE_RKNPU +IExecutionProvider* TestRknpuExecutionProvider(); +#endif + template inline void CopyVectorToTensor(const std::vector& value, Tensor& tensor) { gsl::copy(gsl::make_span(value), tensor.MutableDataAsSpan()); diff --git a/onnxruntime/test/providers/provider_test_utils.cc b/onnxruntime/test/providers/provider_test_utils.cc index ee31544f96..3721dfa3a6 100644 --- a/onnxruntime/test/providers/provider_test_utils.cc +++ b/onnxruntime/test/providers/provider_test_utils.cc @@ -774,6 +774,8 @@ void OpTester::Run( execution_provider = DefaultTensorrtExecutionProvider(); else if (provider_type == onnxruntime::kNnapiExecutionProvider) execution_provider = DefaultNnapiExecutionProvider(); + else if (provider_type == onnxruntime::kRknpuExecutionProvider) + execution_provider = DefaultRknpuExecutionProvider(); else if (provider_type == onnxruntime::kAclExecutionProvider) execution_provider = DefaultAclExecutionProvider(); // skip if execution provider is disabled diff --git a/onnxruntime/test/providers/rknpu/rknpu_basic_test.cc b/onnxruntime/test/providers/rknpu/rknpu_basic_test.cc new file mode 100644 index 0000000000..cdbc826feb --- /dev/null +++ b/onnxruntime/test/providers/rknpu/rknpu_basic_test.cc @@ -0,0 +1,102 @@ +#include "core/session/inference_session.h" +#include "test/providers/provider_test_utils.h" +#include "test/framework/test_utils.h" +#include "gtest/gtest.h" +#include "core/providers/rknpu/rknpu_execution_provider.h" +#include "core/common/logging/logging.h" + +using namespace std; +using namespace ONNX_NAMESPACE; +using namespace ::onnxruntime::logging; + +namespace onnxruntime { + +namespace test { +void VerifyOutputs(const std::vector& fetches, const std::vector& expected_dims, + const std::vector& expected_values) { + ASSERT_EQ(1, fetches.size()); + auto& rtensor = fetches.front().Get(); + TensorShape expected_shape(expected_dims); + ASSERT_EQ(expected_shape, rtensor.Shape()); + const std::vector found(rtensor.template Data(), rtensor.template Data() + expected_values.size()); + ASSERT_EQ(expected_values, found); +} + +TEST(RknpuExecutionProviderTest, FunctionTest) { + onnxruntime::Model model("graph_1", false, DefaultLoggingManager().DefaultLogger()); + auto& graph = model.MainGraph(); + std::vector inputs; + std::vector outputs; + + // FLOAT tensor. + ONNX_NAMESPACE::TypeProto float_tensor; + float_tensor.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + float_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + float_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(1); + float_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(3); + float_tensor.mutable_tensor_type()->mutable_shape()->add_dim()->set_dim_value(2); + + auto& input_arg_1 = graph.GetOrCreateNodeArg("X", &float_tensor); + auto& input_arg_2 = graph.GetOrCreateNodeArg("Y", &float_tensor); + inputs.push_back(&input_arg_1); + inputs.push_back(&input_arg_2); + auto& output_arg = graph.GetOrCreateNodeArg("node_1_out_1", &float_tensor); + outputs.push_back(&output_arg); + graph.AddNode("node_1", "Add", "node 1.", inputs, outputs); + + auto& input_arg_3 = graph.GetOrCreateNodeArg("Z", &float_tensor); + inputs.clear(); + inputs.push_back(&output_arg); + inputs.push_back(&input_arg_3); + auto& output_arg_2 = graph.GetOrCreateNodeArg("M", &float_tensor); + outputs.clear(); + outputs.push_back(&output_arg_2); + graph.AddNode("node_2", "Add", "node 2.", inputs, outputs); + + auto status = graph.Resolve(); + ASSERT_TRUE(status.IsOK()); + std::string model_file_name = "rknpu_execution_provider_test_graph.onnx"; + status = onnxruntime::Model::Save(model, model_file_name); + + std::vector dims_mul_x = {1, 1, 3, 2}; + std::vector values_mul_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + OrtValue ml_value_x; + CreateMLValue(TestRknpuExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x, &ml_value_x); + OrtValue ml_value_y; + CreateMLValue(TestRknpuExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x, &ml_value_y); + OrtValue ml_value_z; + CreateMLValue(TestRknpuExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x, &ml_value_z); + NameMLValMap feeds; + feeds.insert(std::make_pair("X", ml_value_x)); + feeds.insert(std::make_pair("Y", ml_value_y)); + feeds.insert(std::make_pair("Z", ml_value_z)); + + // prepare outputs + std::vector output_names; + output_names.push_back("M"); + std::vector fetches; + + // prepare expected inputs and outputs + std::vector expected_dims_mul_m = {1, 1, 3, 2}; + std::vector expected_values_mul_m = {3.0f, 6.0f, 9.0f, 12.0f, 15.0f, 18.0f}; + + SessionOptions so; + so.session_logid = "RknpuExecutionProviderTest.FunctionTest"; + RunOptions run_options; + run_options.run_tag = so.session_logid; + + InferenceSession session_object{so, GetEnvironment()}; + status = session_object.RegisterExecutionProvider(onnxruntime::make_unique<::onnxruntime::RknpuExecutionProvider>()); + ASSERT_TRUE(status.IsOK()); + status = session_object.Load(model_file_name); + ASSERT_TRUE(status.IsOK()); + status = session_object.Initialize(); + ASSERT_TRUE(status.IsOK()); + + // Now run + status = session_object.Run(run_options, feeds, output_names, &fetches); + ASSERT_TRUE(status.IsOK()); + VerifyOutputs(fetches, expected_dims_mul_m, expected_values_mul_m); +} +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/util/default_providers.cc b/onnxruntime/test/util/default_providers.cc index 0b23204d73..d28b9bfd08 100644 --- a/onnxruntime/test/util/default_providers.cc +++ b/onnxruntime/test/util/default_providers.cc @@ -18,6 +18,7 @@ std::shared_ptr CreateExecutionProviderFactory_NGraph std::shared_ptr CreateExecutionProviderFactory_OpenVINO(const char* device_id); std::shared_ptr CreateExecutionProviderFactory_Nuphar(bool, const char*); std::shared_ptr CreateExecutionProviderFactory_Nnapi(); +std::shared_ptr CreateExecutionProviderFactory_Rknpu(); std::shared_ptr CreateExecutionProviderFactory_Tensorrt(int device_id); std::shared_ptr CreateExecutionProviderFactory_ACL(int use_arena); @@ -86,6 +87,14 @@ std::unique_ptr DefaultNnapiExecutionProvider() { #endif } +std::unique_ptr DefaultRknpuExecutionProvider() { +#ifdef USE_RKNPU + return CreateExecutionProviderFactory_Rknpu()->CreateProvider(); +#else + return nullptr; +#endif +} + std::unique_ptr DefaultAclExecutionProvider(bool enable_arena) { #ifdef USE_ACL return CreateExecutionProviderFactory_ACL(enable_arena)->CreateProvider(); diff --git a/onnxruntime/test/util/include/default_providers.h b/onnxruntime/test/util/include/default_providers.h index 4112710ce8..83a16d3f67 100644 --- a/onnxruntime/test/util/include/default_providers.h +++ b/onnxruntime/test/util/include/default_providers.h @@ -15,6 +15,7 @@ std::unique_ptr DefaultNupharExecutionProvider(bool allow_un std::unique_ptr DefaultTensorrtExecutionProvider(); std::unique_ptr DefaultOpenVINOExecutionProvider(); std::unique_ptr DefaultNnapiExecutionProvider(); +std::unique_ptr DefaultRknpuExecutionProvider(); std::unique_ptr DefaultAclExecutionProvider(bool enable_arena = true); } // namespace test diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index d362f4085f..20a5789c68 100755 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -257,6 +257,8 @@ def parse_arguments(): help="Build with OpenVINO for specific hardware.") parser.add_argument( "--use_dnnlibrary", action='store_true', help="Build with DNNLibrary.") + parser.add_argument( + "--use_rknpu", action='store_true', help="Build with RKNPU.") parser.add_argument( "--use_preinstalled_eigen", action='store_true', help="Use pre-installed Eigen.") @@ -560,10 +562,11 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, "-Donnxruntime_USE_OPENVINO_BINARY=" + ( "ON" if args.use_openvino else "OFF"), "-Donnxruntime_USE_NNAPI=" + ("ON" if args.use_dnnlibrary else "OFF"), + "-Donnxruntime_USE_RKNPU=" + ("ON" if args.use_rknpu else "OFF"), "-Donnxruntime_USE_OPENMP=" + ( "ON" if args.use_openmp and not ( args.use_dnnlibrary or args.use_mklml or args.use_ngraph or - args.android or (args.ios and is_macOS())) + args.android or (args.ios and is_macOS()) or args.use_rknpu) else "OFF"), "-Donnxruntime_USE_TVM=" + ("ON" if args.use_tvm else "OFF"), "-Donnxruntime_USE_LLVM=" + ("ON" if args.use_llvm else "OFF"),