qsdy
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

!25193 [MS][LITE]limit provider max number 

Merge pull request !25193 from chenjianping/r1.5_dev2
This commit is contained in:
i-robot 2021-10-23 06:46:25 +00:00 committed by Gitee
parent 6a81dad0f6 beff650f8f
commit 586d62bfa1
2 changed files with 75 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
mindspore/lite/src/registry/kernel_interface_registry.cc
View File

@ -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<kernel::KernelInterface> KernelInterfaceRegistry::GetKernelInter
return nullptr;
}
int op_type = static_cast<int>(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<kernel::KernelInterface> 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<std::mutex> 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<KernelInterfaceCreator *>(calloc(kMaxKernelNum, sizeof(KernelInterfaceCreator)));
if (kernel_creators_[provider] == nullptr) {

68
mindspore/lite/src/registry/register_kernel_impl.cc
View File

@ -25,12 +25,13 @@ using mindspore::schema::PrimitiveType_MAX;
using mindspore::schema::PrimitiveType_MIN;
namespace mindspore::registry {
namespace {
static const auto kKernelMaxNum =
(static_cast<int>(DataType::kNumberTypeEnd) - static_cast<int>(DataType::kNumberTypeBegin) - 1) *
(PrimitiveType_MAX - PrimitiveType_MIN);
static const auto kOpTypeLen = PrimitiveType_MAX - PrimitiveType_MIN + 1;
static const auto kDataTypeLen =
static_cast<int>(DataType::kNumberTypeEnd) - static_cast<int>(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<int>(data_type) - static_cast<int>(DataType::kNumberTypeBegin) - 1;
if (data_type_index < 0 || data_type_index >= kDataTypeLen) {
MS_LOG(ERROR) << "invalid data_type: " << static_cast<int>(data_type) << "!provider: " << provider;
return kLiteError;
}
std::unique_lock<std::mutex> 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<CreateKernel *>(calloc(kDataTypeLen, sizeof(CreateKernel)));
@ -64,20 +86,30 @@ Status RegistryKernelImpl::RegCustomKernel(const std::string &arch, const std::s
}
}
int data_type_index = static_cast<int>(data_type) - static_cast<int>(DataType::kNumberTypeBegin) - 1;
if (data_type_index < 0 || data_type_index >= kDataTypeLen) {
MS_LOG(ERROR) << "invalid data_type: " << static_cast<int>(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<int>(PrimitiveType_MIN) || type > static_cast<int>(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<int>(data_type) << ",op type "
<< type;
return kLiteError;
}
std::unique_lock<std::mutex> 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<CreateKernel *>(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<CreateKernel *>(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<int>(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<int>(desc->data_type) - static_cast<int>(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();