From 5a28bb86872e135136a20dbfc209878457e4ae6b Mon Sep 17 00:00:00 2001 From: chenjianping Date: Mon, 1 Nov 2021 19:26:30 +0800 Subject: [PATCH] support get attr from kernel --- include/api/kernel.h | 28 +++++++-- .../registry/register_kernel_interface.h | 7 ++- mindspore/lite/src/cxx_api/kernel.cc | 30 ++++++++++ .../src/registry/kernel_interface_registry.cc | 57 +++++++++++-------- .../src/registry/kernel_interface_registry.h | 6 +- .../src/registry/register_kernel_interface.cc | 7 ++- mindspore/lite/src/runtime/infer_manager.cc | 22 +++++-- mindspore/lite/src/runtime/infer_manager.h | 4 +- .../runtime/kernel/opencl/opencl_fusion.cc | 2 +- mindspore/lite/src/sub_graph_kernel.cc | 2 +- 10 files changed, 122 insertions(+), 43 deletions(-) create mode 100644 mindspore/lite/src/cxx_api/kernel.cc diff --git a/include/api/kernel.h b/include/api/kernel.h index 6ec62dec020..d3b855b2444 100644 --- a/include/api/kernel.h +++ b/include/api/kernel.h @@ -19,13 +19,14 @@ #include #include #include +#include #include "schema/model_generated.h" #include "include/api/types.h" #include "include/api/context.h" namespace mindspore::kernel { /// \brief The Kernel class is used to define a MindSpore Kernel. -class Kernel { +class MS_API Kernel { public: Kernel() = default; /// \brief Constructor. @@ -37,9 +38,7 @@ class Kernel { Kernel(const std::vector &inputs, const std::vector &outputs, const schema::Primitive *primitive, const mindspore::Context *ctx) : context_(ctx), inputs_(std::move(inputs)), outputs_(std::move(outputs)), primitive_(primitive) { - if (primitive != nullptr) { - type_ = primitive->value_type(); - } + Initialize(); } /// \brief Destructor. virtual ~Kernel() = default; @@ -101,6 +100,23 @@ class Kernel { /// /// \return the primitive of kernel generated by flatbuffers. const schema::Primitive *primitive() const { return this->primitive_; } + /// \brief get kernel's attribute + /// + /// \param[in] key define the kernel's attribute key. + std::string GetAttr(const std::string &key) const { + auto iter = attrs_.find(key); + if (iter != attrs_.end()) { + return iter->second; + } + return ""; + } + + protected: + /// \brief set kernel's attribute + /// + /// \param[in] key define the kernel's attribute key. + /// \param[in] value define the kernel's attribute value. + void SetAttr(const std::string &key, const std::string &value) { attrs_[key] = value; } protected: std::string name_; @@ -109,6 +125,10 @@ class Kernel { std::vector outputs_; schema::PrimitiveType type_ = schema::PrimitiveType_NONE; const schema::Primitive *primitive_ = nullptr; + std::map attrs_; + + private: + void Initialize(); }; } // namespace mindspore::kernel diff --git a/mindspore/lite/include/registry/register_kernel_interface.h b/mindspore/lite/include/registry/register_kernel_interface.h index dbb8a7e99f0..6519ae88a7d 100644 --- a/mindspore/lite/include/registry/register_kernel_interface.h +++ b/mindspore/lite/include/registry/register_kernel_interface.h @@ -25,6 +25,9 @@ #include "schema/model_generated.h" namespace mindspore { +namespace kernel { +class Kernel; +} namespace registry { /// \brief KernelInterfaceCreator defined a functor to create KernelInterface. using KernelInterfaceCreator = std::function()>; @@ -55,10 +58,12 @@ class MS_API RegisterKernelInterface { /// /// \param[in] provider Define the identification of user. /// \param[in] primitive Define the attributes of a certain op. + /// \param[in] kernel Define the kernel of a certain op. /// /// \return Boolean value to represent registration of a certain op is existing or not. static std::shared_ptr GetKernelInterface(const std::string &provider, - const schema::Primitive *primitive); + const schema::Primitive *primitive, + const kernel::Kernel *kernel = nullptr); }; /// \brief KernelInterfaceReg defined registration class of KernelInterface. diff --git a/mindspore/lite/src/cxx_api/kernel.cc b/mindspore/lite/src/cxx_api/kernel.cc new file mode 100644 index 00000000000..f3577a516c2 --- /dev/null +++ b/mindspore/lite/src/cxx_api/kernel.cc @@ -0,0 +1,30 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "include/api/kernel.h" +namespace mindspore::kernel { +void Kernel::Initialize() { + if (primitive_ == nullptr) { + return; + } + type_ = primitive_->value_type(); + if (type_ == schema::PrimitiveType_Custom) { + auto param = primitive_->value_as_Custom(); + if (param != nullptr && param->type() != nullptr) { + SetAttr("type", param->type()->str()); + } + } +} +} // namespace mindspore::kernel diff --git a/mindspore/lite/src/registry/kernel_interface_registry.cc b/mindspore/lite/src/registry/kernel_interface_registry.cc index 1c9f6ab681b..9124d85391f 100644 --- a/mindspore/lite/src/registry/kernel_interface_registry.cc +++ b/mindspore/lite/src/registry/kernel_interface_registry.cc @@ -20,6 +20,7 @@ #include "src/common/log_adapter.h" #include "src/common/version_manager.h" #include "schema/model_generated.h" +#include "include/api/kernel.h" using mindspore::registry::KernelInterfaceCreator; using mindspore::schema::PrimitiveType_MAX; @@ -32,12 +33,10 @@ static constexpr auto KMaxCustomTypeNum = 200; static const auto kMaxKernelNum = PrimitiveType_MAX - PrimitiveType_MIN + 1; std::string GetCustomType(const schema::Primitive *primitive) { auto param = primitive->value_as_Custom(); - if (param == nullptr) { - return ""; - } - if (param->type() == nullptr) { + if (param == nullptr || param->type() == nullptr) { return ""; } + return param->type()->str(); } } // namespace @@ -92,15 +91,19 @@ std::shared_ptr KernelInterfaceRegistry::GetCustomCache } std::shared_ptr KernelInterfaceRegistry::GetCustomKernelInterface( - const schema::Primitive *primitive) { - MS_ASSERT(primitive != nullptr); + const schema::Primitive *primitive, const kernel::Kernel *kernel) { std::unique_lock lock(mutex_); - auto &&type = GetCustomType(primitive); + std::string type; + if (kernel == nullptr) { + type = GetCustomType(primitive); + } else { + type = kernel->GetAttr("type"); + } for (auto &&item : custom_creators_) { auto &&provider = item.first; - auto kernel = GetCustomCacheInterface(provider, type); - if (kernel != nullptr) { - return kernel; + auto kernel_interface = GetCustomCacheInterface(provider, type); + if (kernel_interface != nullptr) { + return kernel_interface; } auto provider_iter = custom_creators_.find(provider); if (provider_iter == custom_creators_.end()) { @@ -108,32 +111,38 @@ std::shared_ptr KernelInterfaceRegistry::GetCustomKerne } auto creator_iter = provider_iter->second.find(type); if (creator_iter != provider_iter->second.end()) { - kernel = creator_iter->second(); - custom_kernels_[provider][type] = kernel; - return kernel; + kernel_interface = creator_iter->second(); + custom_kernels_[provider][type] = kernel_interface; + return kernel_interface; } } return nullptr; } -std::shared_ptr KernelInterfaceRegistry::GetKernelInterface( - const std::string &provider, const schema::Primitive *primitive) { - if (primitive == nullptr) { +std::shared_ptr KernelInterfaceRegistry::GetKernelInterface(const std::string &provider, + const schema::Primitive *primitive, + const kernel::Kernel *kernel) { + if (primitive == nullptr && kernel == nullptr) { return nullptr; } - int op_type = static_cast(primitive->value_type()); + int op_type; + if (kernel == nullptr) { + op_type = static_cast(primitive->value_type()); + } else { + op_type = static_cast(kernel->type()); + } if (op_type > PrimitiveType_MAX || op_type <= PrimitiveType_MIN) { return nullptr; } if (op_type == schema::PrimitiveType_Custom) { - return GetCustomKernelInterface(primitive); + return GetCustomKernelInterface(primitive, kernel); } std::unique_lock lock(mutex_); - auto kernel = GetCacheInterface(provider, op_type); - if (kernel != nullptr) { - return kernel; + auto kernel_interface = GetCacheInterface(provider, op_type); + if (kernel_interface != nullptr) { + return kernel_interface; } auto iter = kernel_creators_.find(provider); if (iter == kernel_creators_.end()) { @@ -142,9 +151,9 @@ std::shared_ptr KernelInterfaceRegistry::GetKernelInter auto creator = iter->second[op_type]; if (creator != nullptr) { - kernel = creator(); - kernel_interfaces_[provider][op_type] = kernel; - return kernel; + kernel_interface = creator(); + kernel_interfaces_[provider][op_type] = kernel_interface; + return kernel_interface; } return nullptr; } diff --git a/mindspore/lite/src/registry/kernel_interface_registry.h b/mindspore/lite/src/registry/kernel_interface_registry.h index 715b983a945..17bf00dc04a 100644 --- a/mindspore/lite/src/registry/kernel_interface_registry.h +++ b/mindspore/lite/src/registry/kernel_interface_registry.h @@ -35,7 +35,8 @@ class KernelInterfaceRegistry { } std::shared_ptr GetKernelInterface(const std::string &provider, - const schema::Primitive *primitive); + const schema::Primitive *primitive, + const kernel::Kernel *kernel); Status CustomReg(const std::string &provider, const std::string &op_type, const registry::KernelInterfaceCreator creator); Status Reg(const std::string &provider, int op_type, const registry::KernelInterfaceCreator creator); @@ -46,7 +47,8 @@ class KernelInterfaceRegistry { std::shared_ptr GetCacheInterface(const std::string &provider, int op_type); std::shared_ptr GetCustomCacheInterface(const std::string &provider, const std::string &type); - std::shared_ptr GetCustomKernelInterface(const schema::Primitive *primitive); + std::shared_ptr GetCustomKernelInterface(const schema::Primitive *primitive, + const kernel::Kernel *kernel); std::mutex mutex_; // key: provider diff --git a/mindspore/lite/src/registry/register_kernel_interface.cc b/mindspore/lite/src/registry/register_kernel_interface.cc index 72a06e46658..a40879f2fac 100644 --- a/mindspore/lite/src/registry/register_kernel_interface.cc +++ b/mindspore/lite/src/registry/register_kernel_interface.cc @@ -41,10 +41,11 @@ Status RegisterKernelInterface::CustomReg(const std::string &provider, const std #endif } -std::shared_ptr RegisterKernelInterface::GetKernelInterface( - const std::string &provider, const schema::Primitive *primitive) { +std::shared_ptr RegisterKernelInterface::GetKernelInterface(const std::string &provider, + const schema::Primitive *primitive, + const kernel::Kernel *kernel) { #ifndef CUSTOM_KERNEL_REGISTRY_CLIP - return KernelInterfaceRegistry::Instance()->GetKernelInterface(provider, primitive); + return KernelInterfaceRegistry::Instance()->GetKernelInterface(provider, primitive, kernel); #else MS_LOG(ERROR) << unsupport_custom_kernel_register_log; return nullptr; diff --git a/mindspore/lite/src/runtime/infer_manager.cc b/mindspore/lite/src/runtime/infer_manager.cc index 3fd0f10a727..d5b348c73be 100644 --- a/mindspore/lite/src/runtime/infer_manager.cc +++ b/mindspore/lite/src/runtime/infer_manager.cc @@ -34,23 +34,33 @@ namespace mindspore { namespace lite { #ifndef CUSTOM_KERNEL_REGISTRY_CLIP int KernelInferShape(const std::vector &inputs, const std::vector &outputs, - const void *primitive, std::set &&providers, int schema_version) { - if (primitive == nullptr) { + const void *primitive, std::set &&providers, int schema_version, + const kernel::Kernel *kernel) { + if (primitive == nullptr && kernel == nullptr) { return RET_NOT_SUPPORT; } std::shared_ptr kernel_interface = nullptr; - if (IsCustomNode(primitive, schema_version)) { - kernel_interface = - registry::RegisterKernelInterface::GetKernelInterface("", static_cast(primitive)); + bool is_custom_node = false; + if (kernel == nullptr) { + if (IsCustomNode(primitive, schema_version)) { + is_custom_node = true; + } + } else if (kernel->type() == schema::PrimitiveType_Custom) { + is_custom_node = true; + } + if (is_custom_node) { + kernel_interface = registry::RegisterKernelInterface::GetKernelInterface( + "", static_cast(primitive), kernel); } else { for (auto &&provider : providers) { kernel_interface = registry::RegisterKernelInterface::GetKernelInterface( - provider, static_cast(primitive)); + provider, static_cast(primitive), kernel); if (kernel_interface != nullptr) { break; } } } + if (kernel_interface == nullptr) { return RET_NOT_SUPPORT; } diff --git a/mindspore/lite/src/runtime/infer_manager.h b/mindspore/lite/src/runtime/infer_manager.h index e5eb98a68b8..fccb2ed63ec 100644 --- a/mindspore/lite/src/runtime/infer_manager.h +++ b/mindspore/lite/src/runtime/infer_manager.h @@ -26,13 +26,15 @@ #include "src/tensor.h" #include "nnacl/tensor_c.h" #include "nnacl/infer/infer.h" +#include "include/api/kernel.h" namespace mindspore::lite { int KernelInferShape(const std::vector &tensors_in, const std::vector &outputs, OpParameter *parameter); #ifndef CUSTOM_KERNEL_REGISTRY_CLIP int KernelInferShape(const std::vector &inputs, const std::vector &outputs, - const void *primitive, std::set &&providers, int schema_version); + const void *primitive, std::set &&providers, int schema_version, + const kernel::Kernel *kernel = nullptr); #endif class InferManager { public: diff --git a/mindspore/lite/src/runtime/kernel/opencl/opencl_fusion.cc b/mindspore/lite/src/runtime/kernel/opencl/opencl_fusion.cc index 43001569217..c6cd26e0c93 100644 --- a/mindspore/lite/src/runtime/kernel/opencl/opencl_fusion.cc +++ b/mindspore/lite/src/runtime/kernel/opencl/opencl_fusion.cc @@ -428,7 +428,7 @@ int TryFusionConvScaleWeight(LiteKernel *conv_kernel, LiteKernel *scale_kernel) MS_ASSERT(conv_kernel); MS_ASSERT(scale_kernel); auto *scale_param = - reinterpret_cast(reinterpret_cast(scale_kernel)->GetParameter()); + reinterpret_cast(reinterpret_cast(scale_kernel->kernel())->GetParameter()); MS_ASSERT(scale_param); MS_ASSERT(conv_kernel->in_tensors().size() >= INPUT_TENSOR_SIZE_2); auto *filter = conv_kernel->in_tensors().at(1); diff --git a/mindspore/lite/src/sub_graph_kernel.cc b/mindspore/lite/src/sub_graph_kernel.cc index fb7aab29f38..6b4e99cce26 100644 --- a/mindspore/lite/src/sub_graph_kernel.cc +++ b/mindspore/lite/src/sub_graph_kernel.cc @@ -107,7 +107,7 @@ int SubGraphKernel::ReSize() { int ret; #ifndef CUSTOM_KERNEL_REGISTRY_CLIP ret = lite::KernelInferShape(inputs, outputs, kernel->kernel()->primitive(), kernel->Context()->GetProviders(), - schema_version_); + schema_version_, kernel->kernel()); if (ret == lite::RET_NOT_SUPPORT) { #endif auto parameter = kernel->op_parameter();