From beff650f8f9f3542d11dd9e62f7ec5a69d348d13 Mon Sep 17 00:00:00 2001 From: chenjianping Date: Wed, 20 Oct 2021 16:48:35 +0800 Subject: [PATCH] limit max provider num --- .../src/registry/kernel_interface_registry.cc | 31 +++++++-- .../lite/src/registry/register_kernel_impl.cc | 68 +++++++++++++------ 2 files changed, 75 insertions(+), 24 deletions(-) diff --git a/mindspore/lite/src/registry/kernel_interface_registry.cc b/mindspore/lite/src/registry/kernel_interface_registry.cc index c0117fde5f9..1c9f6ab681b 100644 --- a/mindspore/lite/src/registry/kernel_interface_registry.cc +++ b/mindspore/lite/src/registry/kernel_interface_registry.cc @@ -27,16 +27,35 @@ using mindspore::schema::PrimitiveType_MIN; namespace mindspore { namespace registry { namespace { +static constexpr auto kMaxProviderNum = 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(); - MS_ASSERT(param != nullptr); + if (param == nullptr) { + return ""; + } + if (param->type() == nullptr) { + return ""; + } return param->type()->str(); } } // namespace Status KernelInterfaceRegistry::CustomReg(const std::string &provider, const std::string &type, const KernelInterfaceCreator creator) { + auto provider_iter = custom_creators_.find(provider); + if (provider_iter == custom_creators_.end() && custom_creators_.size() >= kMaxProviderNum) { + MS_LOG(ERROR) << "register too many provider!"; + return kLiteError; + } + if (provider_iter != custom_creators_.end()) { + auto type_iter = provider_iter->second.find(type); + if (type_iter == provider_iter->second.end() && provider_iter->second.size() >= KMaxCustomTypeNum) { + MS_LOG(ERROR) << "register too many custom type!"; + return kLiteError; + } + } custom_creators_[provider][type] = creator; return kSuccess; } @@ -104,6 +123,9 @@ std::shared_ptr KernelInterfaceRegistry::GetKernelInter return nullptr; } int op_type = static_cast(primitive->value_type()); + if (op_type > PrimitiveType_MAX || op_type <= PrimitiveType_MIN) { + return nullptr; + } if (op_type == schema::PrimitiveType_Custom) { return GetCustomKernelInterface(primitive); } @@ -118,9 +140,6 @@ std::shared_ptr KernelInterfaceRegistry::GetKernelInter return nullptr; } - if (op_type > PrimitiveType_MAX || op_type <= PrimitiveType_MIN) { - return nullptr; - } auto creator = iter->second[op_type]; if (creator != nullptr) { kernel = creator(); @@ -143,6 +162,10 @@ Status KernelInterfaceRegistry::Reg(const std::string &provider, int op_type, co std::unique_lock lock(mutex_); auto iter = kernel_creators_.find(provider); if (iter == kernel_creators_.end()) { + if (kernel_creators_.size() >= kMaxProviderNum) { + MS_LOG(ERROR) << "register too many provider!"; + return kLiteError; + } kernel_creators_[provider] = reinterpret_cast(calloc(kMaxKernelNum, sizeof(KernelInterfaceCreator))); if (kernel_creators_[provider] == nullptr) { diff --git a/mindspore/lite/src/registry/register_kernel_impl.cc b/mindspore/lite/src/registry/register_kernel_impl.cc index 1866baf1509..a87ae961faf 100644 --- a/mindspore/lite/src/registry/register_kernel_impl.cc +++ b/mindspore/lite/src/registry/register_kernel_impl.cc @@ -25,12 +25,13 @@ using mindspore::schema::PrimitiveType_MAX; using mindspore::schema::PrimitiveType_MIN; namespace mindspore::registry { namespace { -static const auto kKernelMaxNum = - (static_cast(DataType::kNumberTypeEnd) - static_cast(DataType::kNumberTypeBegin) - 1) * - (PrimitiveType_MAX - PrimitiveType_MIN); +static const auto kOpTypeLen = PrimitiveType_MAX - PrimitiveType_MIN + 1; static const auto kDataTypeLen = static_cast(DataType::kNumberTypeEnd) - static_cast(DataType::kNumberTypeBegin) - 1; -static const auto kOpTypeLen = PrimitiveType_MAX - PrimitiveType_MIN; +static const auto kKernelMaxNum = kOpTypeLen * kDataTypeLen; +static constexpr auto kMaxProviderNum = 10; +static constexpr auto kMaxArchPerProviderNum = 10; +static constexpr auto kMaxCustomTypeNum = 200; } // namespace int RegistryKernelImpl::GetFuncIndex(const KernelDesc &desc) const { @@ -50,11 +51,32 @@ int RegistryKernelImpl::GetFuncIndex(const KernelDesc &desc) const { Status RegistryKernelImpl::RegCustomKernel(const std::string &arch, const std::string &provider, DataType data_type, const std::string &type, const CreateKernel creator) { - if (data_type >= DataType::kNumberTypeEnd) { + int data_type_index = static_cast(data_type) - static_cast(DataType::kNumberTypeBegin) - 1; + if (data_type_index < 0 || data_type_index >= kDataTypeLen) { MS_LOG(ERROR) << "invalid data_type: " << static_cast(data_type) << "!provider: " << provider; return kLiteError; } std::unique_lock lock(lock_); + auto provider_iter = custom_kernel_creators_.find(provider); + if (provider_iter == custom_kernel_creators_.end() && custom_kernel_creators_.size() >= kMaxProviderNum) { + MS_LOG(ERROR) << "register too many provider!"; + return kLiteError; + } + if (provider_iter != custom_kernel_creators_.end()) { + auto arch_iter = provider_iter->second.find(arch); + if (arch_iter == provider_iter->second.end()) { + if (provider_iter->second.size() >= kMaxArchPerProviderNum) { + MS_LOG(ERROR) << "register too many arch!"; + return kLiteError; + } + } else { + auto type_iter = arch_iter->second.find(type); + if (type_iter == arch_iter->second.end() && arch_iter->second.size() >= kMaxCustomTypeNum) { + MS_LOG(ERROR) << "register too many type!"; + return kLiteError; + } + } + } if (custom_kernel_creators_[provider][arch][type] == nullptr) { custom_kernel_creators_[provider][arch][type] = reinterpret_cast(calloc(kDataTypeLen, sizeof(CreateKernel))); @@ -64,20 +86,30 @@ Status RegistryKernelImpl::RegCustomKernel(const std::string &arch, const std::s } } - int data_type_index = static_cast(data_type) - static_cast(DataType::kNumberTypeBegin) - 1; - if (data_type_index < 0 || data_type_index >= kDataTypeLen) { - MS_LOG(ERROR) << "invalid data_type: " << static_cast(data_type) << "!provider: " << provider; - return kLiteError; - } custom_kernel_creators_[provider][arch][type][data_type_index] = creator; return kSuccess; } Status RegistryKernelImpl::RegKernel(const std::string &arch, const std::string &provider, DataType data_type, int type, const registry::CreateKernel creator) { + if (type <= static_cast(PrimitiveType_MIN) || type > static_cast(PrimitiveType_MAX)) { + MS_LOG(ERROR) << "Invalid op type : " << type; + return kLiteParamInvalid; + } + KernelDesc desc = {data_type, type, arch, provider}; + int index = GetFuncIndex(desc); + if (index < 0) { + MS_LOG(ERROR) << "invalid kernel key, arch " << arch << ", data_type" << static_cast(data_type) << ",op type " + << type; + return kLiteError; + } std::unique_lock lock(lock_); auto iter = kernel_creators_.find(provider); if (iter == kernel_creators_.end()) { + if (kernel_creators_.size() >= kMaxProviderNum) { + MS_LOG(ERROR) << "register too many provider!"; + return kLiteError; + } kernel_creators_[provider][arch] = reinterpret_cast(calloc(kKernelMaxNum, sizeof(CreateKernel))); if (kernel_creators_[provider][arch] == nullptr) { MS_LOG(ERROR) << "malloc kernel creator buffer fail! provider: " << provider << ",arch:" << arch; @@ -86,6 +118,10 @@ Status RegistryKernelImpl::RegKernel(const std::string &arch, const std::string } else { auto iter_arch = iter->second.find(arch); if (iter_arch == iter->second.end()) { + if (iter->second.size() >= kMaxArchPerProviderNum) { + MS_LOG(ERROR) << "register too many arch!"; + return kLiteError; + } iter->second[arch] = reinterpret_cast(calloc(kKernelMaxNum, sizeof(CreateKernel))); if (iter->second[arch] == nullptr) { MS_LOG(ERROR) << "malloc kernel creator buffer fail! provider: " << provider << ",arch:" << arch; @@ -94,14 +130,6 @@ Status RegistryKernelImpl::RegKernel(const std::string &arch, const std::string } } - KernelDesc desc = {data_type, type, arch, provider}; - int index = GetFuncIndex(desc); - if (index < 0) { - MS_LOG(ERROR) << "invalid kernel key, arch " << arch << ", data_type" << static_cast(data_type) << ",op type " - << type; - return kLiteError; - } - kernel_creators_[provider][arch][index] = creator; return kSuccess; } @@ -109,11 +137,11 @@ Status RegistryKernelImpl::RegKernel(const std::string &arch, const std::string registry::CreateKernel RegistryKernelImpl::GetCustomKernelCreator(const schema::Primitive *primitive, KernelDesc *desc) { int data_type_index = static_cast(desc->data_type) - static_cast(DataType::kNumberTypeBegin) - 1; - if (data_type_index < 0) { + if (data_type_index < 0 || desc->data_type >= DataType::kNumberTypeEnd) { return nullptr; } auto param = primitive->value_as_Custom(); - if (param == nullptr) { + if (param == nullptr || param->type() == nullptr) { return nullptr; } auto custom_type = param->type()->str();