diff --git a/include/api/kernel.h b/include/api/kernel.h index 6ec62dec020..c5461f3edd6 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; @@ -102,6 +101,44 @@ 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 ""; + } + + /// \brief set kernel's config. + /// + /// \param[in] config define the kernel's config. + void SetConfig(const std::map> *config) { + config_ = config; + } + /// \brief set kernel's config. + /// + /// \param[in] config define the kernel's config. + std::map GetConfig(const std::string §ion) const { + if (config_ == nullptr) { + return std::map(); + } + auto iter = config_->find(section); + if (iter != config_->end()) { + return iter->second; + } + return std::map(); + } + + 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_; const mindspore::Context *context_ = nullptr; @@ -109,6 +146,11 @@ class Kernel { std::vector outputs_; schema::PrimitiveType type_ = schema::PrimitiveType_NONE; const schema::Primitive *primitive_ = nullptr; + std::map attrs_; + const std::map> *config_; + + private: + void Initialize(); }; } // namespace mindspore::kernel diff --git a/include/api/model.h b/include/api/model.h index 8d34871527b..8de66fa0d80 100644 --- a/include/api/model.h +++ b/include/api/model.h @@ -114,6 +114,14 @@ class MS_API Model { /// \return Status. inline Status LoadConfig(const std::string &config_path); + /// \brief Update config. + /// + /// \param[in] section define the config section. + /// \param[in] config define the config will be updated. + /// + /// \return Status. + inline Status UpdateConfig(const std::string §ion, const std::pair &config); + /// \brief Obtains all input tensors of the model. /// /// \return The vector that includes all input tensors. @@ -223,6 +231,7 @@ class MS_API Model { MSTensor GetOutputByTensorName(const std::vector &tensor_name); std::vector GetOutputsByNodeName(const std::vector &node_name); Status LoadConfig(const std::vector &config_path); + Status UpdateConfig(const std::vector §ion, const std::pair, std::vector> &config); Status Build(const void *model_data, size_t data_size, ModelType model_type, const std::shared_ptr &model_context, const Key &dec_key, const std::vector &dec_mode); Status Build(const std::vector &model_path, ModelType model_type, const std::shared_ptr &model_context, @@ -249,6 +258,12 @@ Status Model::LoadConfig(const std::string &config_path) { return LoadConfig(StringToChar(config_path)); } +Status Model::UpdateConfig(const std::string §ion, const std::pair &config) { + std::pair, std::vector> config_pair = {StringToChar(config.first), + StringToChar(config.second)}; + return UpdateConfig(StringToChar(section), config_pair); +} + Status Model::Build(const void *model_data, size_t data_size, ModelType model_type, const std::shared_ptr &model_context, const Key &dec_key, const std::string &dec_mode) { return Build(model_data, data_size, model_type, model_context, dec_key, StringToChar(dec_mode)); diff --git a/mindspore/lite/include/registry/register_kernel.h b/mindspore/lite/include/registry/register_kernel.h index 8a9cf476b64..38084cbc61a 100644 --- a/mindspore/lite/include/registry/register_kernel.h +++ b/mindspore/lite/include/registry/register_kernel.h @@ -71,7 +71,7 @@ class MS_API RegisterKernel { /// /// \return Status as a status identification of registering. inline static Status RegKernel(const std::string &arch, const std::string &provider, DataType data_type, int type, - CreateKernel creator); + const CreateKernel creator); /// \brief Static method to register kernel which is corresponding to custom op. /// @@ -83,7 +83,7 @@ class MS_API RegisterKernel { /// /// \return Status as a status identification of registering. inline static Status RegCustomKernel(const std::string &arch, const std::string &provider, DataType data_type, - const std::string &type, CreateKernel creator); + const std::string &type, const CreateKernel creator); /// \brief Static methon to get a kernel's create function. /// @@ -95,9 +95,9 @@ class MS_API RegisterKernel { private: static Status RegKernel(const std::vector &arch, const std::vector &provider, DataType data_type, - int type, CreateKernel creator); + int type, const CreateKernel creator); static Status RegCustomKernel(const std::vector &arch, const std::vector &provider, DataType data_type, - const std::vector &type, CreateKernel creator); + const std::vector &type, const CreateKernel creator); static CreateKernel GetCreator(const schema::Primitive *primitive, KernelDescHelper *desc); }; @@ -115,7 +115,7 @@ class MS_API KernelReg { /// \param[in] op_type Define the ordinary op type. /// \param[in] creator Define a function pointer to create a kernel. KernelReg(const std::string &arch, const std::string &provider, DataType data_type, int op_type, - CreateKernel creator) { + const CreateKernel creator) { RegisterKernel::RegKernel(arch, provider, data_type, op_type, creator); } @@ -127,18 +127,18 @@ class MS_API KernelReg { /// \param[in] op_type Define the concrete type of a custom op. /// \param[in] creator Define a function pointer to create a kernel. KernelReg(const std::string &arch, const std::string &provider, DataType data_type, const std::string &op_type, - CreateKernel creator) { + const CreateKernel creator) { RegisterKernel::RegCustomKernel(arch, provider, data_type, op_type, creator); } }; Status RegisterKernel::RegKernel(const std::string &arch, const std::string &provider, DataType data_type, int type, - CreateKernel creator) { + const CreateKernel creator) { return RegKernel(StringToChar(arch), StringToChar(provider), data_type, type, creator); } Status RegisterKernel::RegCustomKernel(const std::string &arch, const std::string &provider, DataType data_type, - const std::string &type, CreateKernel creator) { + const std::string &type, const CreateKernel creator) { return RegCustomKernel(StringToChar(arch), StringToChar(provider), data_type, StringToChar(type), creator); } diff --git a/mindspore/lite/include/registry/register_kernel_interface.h b/mindspore/lite/include/registry/register_kernel_interface.h index d8ddd260225..93e02a45741 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()>; @@ -40,7 +43,7 @@ class MS_API RegisterKernelInterface { /// /// \return Status as a status identification of registering. inline static Status CustomReg(const std::string &provider, const std::string &op_type, - KernelInterfaceCreator creator); + const KernelInterfaceCreator creator); /// \brief Static method to register op whose primitive type is ordinary. /// @@ -49,23 +52,26 @@ class MS_API RegisterKernelInterface { /// \param[in] creator Define the KernelInterface create function. /// /// \return Status as a status identification of registering. - inline static Status Reg(const std::string &provider, int op_type, KernelInterfaceCreator creator); + inline static Status Reg(const std::string &provider, int op_type, const KernelInterfaceCreator creator); /// \brief Static method to get registration of a certain op. /// /// \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. inline static std::shared_ptr GetKernelInterface(const std::string &provider, - const schema::Primitive *primitive); + const schema::Primitive *primitive, + const kernel::Kernel *kernel = nullptr); private: static Status CustomReg(const std::vector &provider, const std::vector &op_type, - KernelInterfaceCreator creator); - static Status Reg(const std::vector &provider, int op_type, KernelInterfaceCreator creator); + const KernelInterfaceCreator creator); + static Status Reg(const std::vector &provider, int op_type, const KernelInterfaceCreator creator); static std::shared_ptr GetKernelInterface(const std::vector &provider, - const schema::Primitive *primitive); + const schema::Primitive *primitive, + const kernel::Kernel *kernel = nullptr); }; /// \brief KernelInterfaceReg defined registration class of KernelInterface. @@ -76,7 +82,7 @@ class MS_API KernelInterfaceReg { /// \param[in] provider Define the identification of user. /// \param[in] op_type Define the ordinary op type. /// \param[in] creator Define the KernelInterface create function. - KernelInterfaceReg(const std::string &provider, int op_type, KernelInterfaceCreator creator) { + KernelInterfaceReg(const std::string &provider, int op_type, const KernelInterfaceCreator creator) { RegisterKernelInterface::Reg(provider, op_type, creator); } @@ -85,23 +91,26 @@ class MS_API KernelInterfaceReg { /// \param[in] provider Define the identification of user. /// \param[in] op_type Define the concrete type of a custom op. /// \param[in] creator Define the KernelInterface create function. - KernelInterfaceReg(const std::string &provider, const std::string &op_type, KernelInterfaceCreator creator) { + KernelInterfaceReg(const std::string &provider, const std::string &op_type, const KernelInterfaceCreator creator) { RegisterKernelInterface::CustomReg(provider, op_type, creator); } + + virtual ~KernelInterfaceReg() = default; }; Status RegisterKernelInterface::CustomReg(const std::string &provider, const std::string &op_type, - KernelInterfaceCreator creator) { + const KernelInterfaceCreator creator) { return CustomReg(StringToChar(provider), StringToChar(op_type), creator); } -Status RegisterKernelInterface::Reg(const std::string &provider, int op_type, KernelInterfaceCreator creator) { +Status RegisterKernelInterface::Reg(const std::string &provider, int op_type, const KernelInterfaceCreator creator) { return Reg(StringToChar(provider), op_type, creator); } -std::shared_ptr RegisterKernelInterface::GetKernelInterface( - const std::string &provider, const schema::Primitive *primitive) { - return GetKernelInterface(StringToChar(provider), primitive); +std::shared_ptr RegisterKernelInterface::GetKernelInterface(const std::string &provider, + const schema::Primitive *primitive, + const kernel::Kernel *kernel) { + return GetKernelInterface(StringToChar(provider), primitive, kernel); } /// \brief Defined registering macro to register ordinary op, which called by user directly. diff --git a/mindspore/lite/src/common/config_file.cc b/mindspore/lite/src/common/config_file.cc index 833a617797e..71a26a37e2d 100644 --- a/mindspore/lite/src/common/config_file.cc +++ b/mindspore/lite/src/common/config_file.cc @@ -21,19 +21,55 @@ #endif namespace { constexpr size_t kLengthOfParentheses = 2; -} +constexpr size_t kMinSectionLineLength = 2; +constexpr size_t kMaxValidLineCount = 100000; +constexpr size_t kMaxLineCount = 100100; + +} // namespace namespace mindspore { namespace lite { -int GetSectionInfoFromConfigFile(const std::string &file, const std::string §ion_name, - std::map *section_info) { - if (file.empty()) { - MS_LOG(ERROR) << "file is nullptr"; +namespace { +void ParseLine(const std::string &line, std::map *section_config, std::string *section, + size_t *valid_line_count, std::map> *config) { + // eg: [section] + // key=value + if (line[0] == '[' && line[line.length() - 1] == ']') { + if (!section->empty() && !section_config->empty()) { + config->insert(std::make_pair(*section, *section_config)); + } + section_config->clear(); + *section = line.substr(1, line.length() - kLengthOfParentheses); + *valid_line_count = *valid_line_count + 1; + } + + if (!section->empty()) { + auto index = line.find('='); + if (index == std::string::npos) { + return; + } + auto key = line.substr(0, index); + if (index + 1 > line.size()) { + return; + } + auto value = line.substr(index + 1); + lite::Trim(&key); + lite::Trim(&value); + section_config->insert(std::make_pair(key, value)); + *valid_line_count = *valid_line_count + 1; + } +} +} // namespace + +int GetAllSectionInfoFromConfigFile(const std::string &file, + std::map> *config) { + if (file.empty() || config == nullptr) { + MS_LOG(ERROR) << "input Invalid!check file and config."; return RET_ERROR; } auto resolved_path = std::make_unique(PATH_MAX); if (resolved_path == nullptr) { - MS_LOG(ERROR) << "new resolved_path failed"; + MS_LOG(ERROR) << "new resolved_path fail!"; return RET_ERROR; } @@ -56,44 +92,25 @@ int GetSectionInfoFromConfigFile(const std::string &file, const std::string &sec return RET_ERROR; } std::string line; - - bool find_section = false; + std::string section; + std::map section_config; + size_t line_count = 0; + size_t valid_line_count = 0; while (std::getline(ifs, line)) { + line_count++; + if (line_count >= kMaxLineCount || valid_line_count >= kMaxValidLineCount) { + MS_LOG(ERROR) << "config too many lines!"; + return RET_ERROR; + } lite::Trim(&line); - if (line.empty()) { + if (line.length() <= kMinSectionLineLength || line[0] == '#') { continue; } - if (line[0] == '#') { - continue; - } - - if (line[0] == '[') { - if (find_section == true) { - break; - } - std::string section = line.substr(1, line.length() - kLengthOfParentheses); - if (section != section_name) { - continue; - } - find_section = true; - } - - if (find_section == true) { - auto index = line.find('='); - if (index == std::string::npos) { - continue; - } - auto key = line.substr(0, index); - if (index + 1 > line.size()) { - return RET_ERROR; - } - auto value = line.substr(index + 1); - lite::Trim(&key); - lite::Trim(&value); - section_info->insert(std::make_pair(key, value)); - } + ParseLine(line, §ion_config, §ion, &valid_line_count, config); + } + if (!section.empty() && !section_config.empty()) { + config->insert(std::make_pair(section, section_config)); } - ifs.close(); return RET_OK; } diff --git a/mindspore/lite/src/common/config_file.h b/mindspore/lite/src/common/config_file.h index 70aac9dcb05..47a14c60138 100644 --- a/mindspore/lite/src/common/config_file.h +++ b/mindspore/lite/src/common/config_file.h @@ -35,10 +35,8 @@ namespace mindspore { namespace lite { constexpr int MAX_CONFIG_FILE_LENGTH = 1024; -#define CONFIG_FILE_EXECUTION_PLAN "execution_plan" - -int GetSectionInfoFromConfigFile(const std::string &file, const std::string §ion_name, - std::map *section_info); +int GetAllSectionInfoFromConfigFile(const std::string &file, + std::map> *config); void ParserExecutionPlan(const std::map *config_infos, std::map *data_type_plan); 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/cxx_api/model/model.cc b/mindspore/lite/src/cxx_api/model/model.cc index 1cd1f2c3972..fbf32afde28 100644 --- a/mindspore/lite/src/cxx_api/model/model.cc +++ b/mindspore/lite/src/cxx_api/model/model.cc @@ -217,6 +217,19 @@ Status Model::LoadConfig(const std::vector &config_path) { return kSuccess; } +Status Model::UpdateConfig(const std::vector §ion, + const std::pair, std::vector> &config) { + std::unique_lock impl_lock(g_impl_init_lock); + if (impl_ == nullptr) { + impl_ = std::shared_ptr(new (std::nothrow) ModelImpl()); + } + if (impl_ != nullptr) { + return impl_->UpdateConfig(CharToString(section), {CharToString(config.first), CharToString(config.second)}); + } + MS_LOG(ERROR) << "Model implement is null!"; + return kLiteFileError; +} + Status Model::SetTrainMode(bool train) { if ((impl_ == nullptr) || (impl_->session_ == nullptr)) { MS_LOG(ERROR) << "Model is null."; diff --git a/mindspore/lite/src/cxx_api/model/model_impl.cc b/mindspore/lite/src/cxx_api/model/model_impl.cc index ccc3bff55f4..8939b401108 100644 --- a/mindspore/lite/src/cxx_api/model/model_impl.cc +++ b/mindspore/lite/src/cxx_api/model/model_impl.cc @@ -17,6 +17,10 @@ #include "src/cxx_api/model/model_impl.h" #include #include +#include +#include +#include +#include #include "include/api/types.h" #include "include/api/context.h" #include "include/lite_session.h" @@ -32,6 +36,11 @@ #include "src/common/config_file.h" namespace mindspore { +namespace { +static const char *kExecutionPlan = "execution_plan"; +static constexpr size_t kMaxSectionNum = 100; +static constexpr size_t kMaxConfigNumPerSection = 1000; +} // namespace using mindspore::lite::RET_ERROR; using mindspore::lite::RET_OK; @@ -195,15 +204,16 @@ Status ModelImpl::RunGraph(const MSKernelCallBack &before, const MSKernelCallBac bool ModelImpl::IsTrainModel() { return (graph_ && graph_->graph_data_ && graph_->graph_data_->IsTrainModel()); } Status ModelImpl::LoadConfig(const std::string &config_path) { - std::map config_info; - int ret = lite::GetSectionInfoFromConfigFile(config_path, CONFIG_FILE_EXECUTION_PLAN, &config_info); + std::map> all_config_info; + int ret = lite::GetAllSectionInfoFromConfigFile(config_path, &all_config_info); if (ret != RET_OK) { - MS_LOG(ERROR) << "GetSectionInfoFromConfigFile failed."; + MS_LOG(ERROR) << "GetAllSectionInfoFromConfigFile fail!ret: " << ret; return kLiteFileError; } - + config_info_ = all_config_info; + std::map config_info = all_config_info[kExecutionPlan]; if (config_info.empty()) { - MS_LOG(WARNING) << "No valid info in config file."; + MS_LOG(WARNING) << "No valid execution plan info in config file."; return kSuccess; } @@ -211,6 +221,24 @@ Status ModelImpl::LoadConfig(const std::string &config_path) { return kSuccess; } +Status ModelImpl::UpdateConfig(const std::string §ion, const std::pair &config) { + auto iter = config_info_.find(section); + if (iter == config_info_.end()) { + if (config_info_.size() >= kMaxSectionNum) { + MS_LOG(ERROR) << "config too many sections!"; + return kLiteError; + } + config_info_[section][config.first] = config.second; + return kSuccess; + } + if (iter->second.size() >= kMaxConfigNumPerSection) { + MS_LOG(ERROR) << "config too many items!"; + return kLiteError; + } + iter->second[config.first] = config.second; + return kSuccess; +} + Status ModelImpl::Predict(const std::vector &inputs, std::vector *outputs, const MSKernelCallBack &before, const MSKernelCallBack &after) { if (outputs == nullptr) { @@ -590,6 +618,7 @@ session::LiteSession *ModelImpl::CreateLiteSession(lite::InnerContext *context) } session->InitExecutionConfig(&execution_plan_); + session->SetConfigInfo(&config_info_); auto ret = session->Init(context); if (ret != mindspore::lite::RET_OK) { diff --git a/mindspore/lite/src/cxx_api/model/model_impl.h b/mindspore/lite/src/cxx_api/model/model_impl.h index 90a6533cb90..3fb45439c0b 100644 --- a/mindspore/lite/src/cxx_api/model/model_impl.h +++ b/mindspore/lite/src/cxx_api/model/model_impl.h @@ -71,6 +71,7 @@ class ModelImpl { session::LiteSession *CreateLiteSession(lite::InnerContext *context); Status LoadConfig(const std::string &config_path); + Status UpdateConfig(const std::string §ion, const std::pair &config); std::vector GetInputs(); std::vector GetOutputs(); std::vector GetGradients() const; @@ -113,6 +114,7 @@ class ModelImpl { void SetConfig(const std::shared_ptr cfg) { cfg_ = cfg; } Status RunGraph(const MSKernelCallBack &before, const MSKernelCallBack &after); std::map execution_plan_; + std::map> config_info_; }; } // namespace mindspore diff --git a/mindspore/lite/src/lite_session.cc b/mindspore/lite/src/lite_session.cc index 382e3a91653..d2aa8f42a58 100644 --- a/mindspore/lite/src/lite_session.cc +++ b/mindspore/lite/src/lite_session.cc @@ -529,6 +529,7 @@ int LiteSession::CompileGraph(Model *model) { Scheduler scheduler(context_, ms_context_, model, &tensors_, inputs_, outputs_, is_train_session_, &is_infershape_, &is_control_flow_, execution_plan_, delegate_, delegate_device_type_); scheduler.SetupSchedulerCb(std::move(sched_cb_)); + scheduler.SetConfig(config_info_); ret = scheduler.Schedule(&kernels_); if (ret != RET_OK) { MS_LOG(ERROR) << "Schedule kernels failed: " << ret; diff --git a/mindspore/lite/src/lite_session.h b/mindspore/lite/src/lite_session.h index 7cdbe6e3cee..68e6fa1799b 100644 --- a/mindspore/lite/src/lite_session.h +++ b/mindspore/lite/src/lite_session.h @@ -87,6 +87,10 @@ class LiteSession : public session::LiteSession { const Delegate *get_delegate() const { return this->delegate_.get(); } + void SetConfigInfo(const std::map> *config_info) { + config_info_ = config_info; + } + protected: static void ConvertTensorsQuantParam(const schema::Tensor *src_tensor, lite::Tensor *dst_tensor); @@ -181,6 +185,7 @@ class LiteSession : public session::LiteSession { std::shared_ptr delegate_ = nullptr; int delegate_device_type_ = -1; // -1: not specified; 0: CPU; 1: GPU; 2: NPU std::map *execution_plan_ = nullptr; + const std::map> *config_info_ = nullptr; }; } // namespace lite } // namespace mindspore diff --git a/mindspore/lite/src/ops/populate/affine_populate.cc b/mindspore/lite/src/ops/populate/affine_populate.cc index 999849aa1e6..d84df44f90f 100644 --- a/mindspore/lite/src/ops/populate/affine_populate.cc +++ b/mindspore/lite/src/ops/populate/affine_populate.cc @@ -70,14 +70,14 @@ OpParameter *PopulateAffineParameter(const void *prim) { affine_param->context_size_ = static_cast(context.size()); // malloc && memset for context - affine_param->context_ = reinterpret_cast(malloc(affine_param->context_size_ * sizeof(int))); + affine_param->context_ = reinterpret_cast(malloc(context.size() * sizeof(int))); if (affine_param->context_ == nullptr) { MS_LOG(ERROR) << "malloc param context_ for affine layer failed!"; ReleaseParam(affine_param, matmul_param); return nullptr; } - memset(affine_param->context_, 0, affine_param->context_size_ * sizeof(int)); - for (int i = 0; i < affine_param->context_size_; ++i) { + (void)memset(affine_param->context_, 0, context.size() * sizeof(int)); + for (size_t i = 0; i < context.size(); ++i) { affine_param->context_[i] = context.at(i); } affine_param->output_dim_ = value->output_dim(); diff --git a/mindspore/lite/src/ops/populate/control/tensor_array_populate.cc b/mindspore/lite/src/ops/populate/control/tensor_array_populate.cc index 23b552c54e4..4cee72fa93f 100644 --- a/mindspore/lite/src/ops/populate/control/tensor_array_populate.cc +++ b/mindspore/lite/src/ops/populate/control/tensor_array_populate.cc @@ -43,8 +43,8 @@ OpParameter *PopulateTensorArrayParameter(const void *prim) { bool identical_element_shapes = value->identical_element_shapes(); param->identical_element_shapes_ = identical_element_shapes; std::vector primitive_element_shape(value->element_shape()->begin(), value->element_shape()->end()); - param->element_shape_size_ = primitive_element_shape.size(); - int size = sizeof(int) * param->element_shape_size_; + param->element_shape_size_ = static_cast(primitive_element_shape.size()); + auto size = sizeof(int) * param->element_shape_size_; param->element_shape_ = static_cast(malloc(size)); if (param->element_shape_ == nullptr) { MS_LOG(ERROR) << "malloc element_shape failed!"; diff --git a/mindspore/lite/src/ops/populate/splice_populate.cc b/mindspore/lite/src/ops/populate/splice_populate.cc index fbbab170710..635fb1fdb16 100644 --- a/mindspore/lite/src/ops/populate/splice_populate.cc +++ b/mindspore/lite/src/ops/populate/splice_populate.cc @@ -52,7 +52,7 @@ OpParameter *PopulateSpliceParameter(const void *prim) { param->context_dim_ = static_cast(primitive_context.size()); // malloc && memset for context - param->context_ = reinterpret_cast(malloc(param->context_dim_ * sizeof(int))); + param->context_ = reinterpret_cast(malloc(primitive_context.size() * sizeof(int))); if (param->context_ == nullptr) { MS_LOG(ERROR) << "malloc param context_ error"; free(param); @@ -60,8 +60,8 @@ OpParameter *PopulateSpliceParameter(const void *prim) { } // src_to_dst_row_offset int src_to_dst_row_offset = INT32_MIN; - memset(param->context_, 0, param->context_dim_ * sizeof(int)); - for (int i = 0; i < param->context_dim_; ++i) { + (void)memset(param->context_, 0, primitive_context.size() * sizeof(int)); + for (size_t i = 0; i < primitive_context.size(); ++i) { param->context_[i] = primitive_context[i]; src_to_dst_row_offset = std::max(src_to_dst_row_offset, std::abs(primitive_context.at(i))); } @@ -83,15 +83,15 @@ OpParameter *PopulateSpliceParameter(const void *prim) { param->forward_indexes_dim_ = static_cast(primitive_forward_indexes.size()); // malloc && memset for forward_indexes - param->forward_indexes_ = reinterpret_cast(malloc(param->forward_indexes_dim_ * sizeof(int))); + param->forward_indexes_ = reinterpret_cast(malloc(primitive_context.size() * sizeof(int))); if (param->forward_indexes_ == nullptr) { MS_LOG(ERROR) << "malloc param forward_indexes_ error"; free(param->context_); free(param); return nullptr; } - memset(param->forward_indexes_, 0, param->forward_indexes_dim_ * sizeof(int)); - memcpy(param->forward_indexes_, primitive_forward_indexes.data(), param->forward_indexes_dim_ * sizeof(int)); + (void)memset(param->forward_indexes_, 0, primitive_context.size() * sizeof(int)); + (void)memcpy(param->forward_indexes_, primitive_forward_indexes.data(), primitive_context.size() * sizeof(int)); param->output_dim_ = value->output_dim(); return reinterpret_cast(param); } diff --git a/mindspore/lite/src/registry/kernel_interface_registry.cc b/mindspore/lite/src/registry/kernel_interface_registry.cc index d90352c071c..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; @@ -27,16 +28,33 @@ 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 || param->type() == nullptr) { + return ""; + } + return param->type()->str(); } } // namespace Status KernelInterfaceRegistry::CustomReg(const std::string &provider, const std::string &type, - KernelInterfaceCreator creator) { + 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; } @@ -73,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()) { @@ -89,47 +111,54 @@ 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; + 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; } - int op_type = primitive->value_type(); 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()) { return nullptr; } - if (op_type > PrimitiveType_MAX || op_type <= PrimitiveType_MIN) { - return nullptr; - } + 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; } -Status KernelInterfaceRegistry::Reg(const std::string &provider, int op_type, KernelInterfaceCreator creator) { +Status KernelInterfaceRegistry::Reg(const std::string &provider, int op_type, const KernelInterfaceCreator creator) { if (op_type <= PrimitiveType_MIN || op_type > PrimitiveType_MAX) { MS_LOG(ERROR) << "reg op_type invalid!op_type: " << op_type << ", max value: " << PrimitiveType_MAX; return kLiteParamInvalid; @@ -142,6 +171,10 @@ Status KernelInterfaceRegistry::Reg(const std::string &provider, int op_type, Ke 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/kernel_interface_registry.h b/mindspore/lite/src/registry/kernel_interface_registry.h index 0739eb64ab7..17bf00dc04a 100644 --- a/mindspore/lite/src/registry/kernel_interface_registry.h +++ b/mindspore/lite/src/registry/kernel_interface_registry.h @@ -35,9 +35,11 @@ class KernelInterfaceRegistry { } std::shared_ptr GetKernelInterface(const std::string &provider, - const schema::Primitive *primitive); - Status CustomReg(const std::string &provider, const std::string &op_type, registry::KernelInterfaceCreator creator); - Status Reg(const std::string &provider, int op_type, registry::KernelInterfaceCreator creator); + 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); virtual ~KernelInterfaceRegistry(); private: @@ -45,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.cc b/mindspore/lite/src/registry/register_kernel.cc index 325c5b24b89..bc5f9d883b7 100644 --- a/mindspore/lite/src/registry/register_kernel.cc +++ b/mindspore/lite/src/registry/register_kernel.cc @@ -23,7 +23,7 @@ namespace mindspore { namespace registry { Status RegisterKernel::RegCustomKernel(const std::vector &arch, const std::vector &provider, - DataType data_type, const std::vector &type, CreateKernel creator) { + DataType data_type, const std::vector &type, const CreateKernel creator) { #ifndef CUSTOM_KERNEL_REGISTRY_CLIP return RegistryKernelImpl::GetInstance()->RegCustomKernel(CharToString(arch), CharToString(provider), data_type, CharToString(type), creator); @@ -34,7 +34,7 @@ Status RegisterKernel::RegCustomKernel(const std::vector &arch, const std: } Status RegisterKernel::RegKernel(const std::vector &arch, const std::vector &provider, DataType data_type, - int op_type, CreateKernel creator) { + int op_type, const CreateKernel creator) { #ifndef CUSTOM_KERNEL_REGISTRY_CLIP return RegistryKernelImpl::GetInstance()->RegKernel(CharToString(arch), CharToString(provider), data_type, op_type, creator); diff --git a/mindspore/lite/src/registry/register_kernel_impl.cc b/mindspore/lite/src/registry/register_kernel_impl.cc index 8ae95577a46..68917ea585f 100644 --- a/mindspore/lite/src/registry/register_kernel_impl.cc +++ b/mindspore/lite/src/registry/register_kernel_impl.cc @@ -25,15 +25,14 @@ 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; -} // namespace - -int RegistryKernelImpl::GetFuncIndex(const KernelDesc &desc) { +static const auto kKernelMaxNum = kOpTypeLen * kDataTypeLen; +static constexpr auto kMaxProviderNum = 10; +static constexpr auto kMaxArchPerProviderNum = 10; +static constexpr auto kMaxCustomTypeNum = 200; +int GetFuncIndex(const KernelDesc &desc) { if (desc.data_type >= DataType::kNumberTypeEnd) { return -1; } @@ -47,14 +46,36 @@ int RegistryKernelImpl::GetFuncIndex(const KernelDesc &desc) { } return index; } +} // namespace Status RegistryKernelImpl::RegCustomKernel(const std::string &arch, const std::string &provider, DataType data_type, - const std::string &type, CreateKernel creator) { - if (data_type >= DataType::kNumberTypeEnd) { + const std::string &type, const CreateKernel creator) { + 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 +85,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, - registry::CreateKernel creator) { + 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 +117,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 +129,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 +136,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 || data_type_index >= kDataTypeLen) { + 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(); diff --git a/mindspore/lite/src/registry/register_kernel_impl.h b/mindspore/lite/src/registry/register_kernel_impl.h index 844f62d9528..666df603c1d 100644 --- a/mindspore/lite/src/registry/register_kernel_impl.h +++ b/mindspore/lite/src/registry/register_kernel_impl.h @@ -37,10 +37,10 @@ class RegistryKernelImpl { } Status RegCustomKernel(const std::string &arch, const std::string &provider, DataType data_type, - const std::string &type, registry::CreateKernel creator); + const std::string &type, const registry::CreateKernel creator); Status RegKernel(const std::string &arch, const std::string &provider, DataType data_type, int type, - registry::CreateKernel creator); + const registry::CreateKernel creator); virtual registry::CreateKernel GetProviderCreator(const schema::Primitive *primitive, registry::KernelDesc *desc); @@ -60,7 +60,6 @@ class RegistryKernelImpl { std::mutex lock_; registry::CreateKernel GetCustomKernelCreator(const schema::Primitive *primitive, registry::KernelDesc *desc); - int GetFuncIndex(const registry::KernelDesc &desc); }; } // namespace mindspore::registry diff --git a/mindspore/lite/src/registry/register_kernel_interface.cc b/mindspore/lite/src/registry/register_kernel_interface.cc index d24fb5b391a..676805fe80b 100644 --- a/mindspore/lite/src/registry/register_kernel_interface.cc +++ b/mindspore/lite/src/registry/register_kernel_interface.cc @@ -22,7 +22,8 @@ namespace mindspore { namespace registry { -Status RegisterKernelInterface::Reg(const std::vector &provider, int op_type, KernelInterfaceCreator creator) { +Status RegisterKernelInterface::Reg(const std::vector &provider, int op_type, + const KernelInterfaceCreator creator) { #ifndef CUSTOM_KERNEL_REGISTRY_CLIP return KernelInterfaceRegistry::Instance()->Reg(CharToString(provider), op_type, creator); #else @@ -32,7 +33,7 @@ Status RegisterKernelInterface::Reg(const std::vector &provider, int op_ty } Status RegisterKernelInterface::CustomReg(const std::vector &provider, const std::vector &op_type, - KernelInterfaceCreator creator) { + const KernelInterfaceCreator creator) { #ifndef CUSTOM_KERNEL_REGISTRY_CLIP return KernelInterfaceRegistry::Instance()->CustomReg(CharToString(provider), CharToString(op_type), creator); #else @@ -41,10 +42,11 @@ Status RegisterKernelInterface::CustomReg(const std::vector &provider, con #endif } -std::shared_ptr RegisterKernelInterface::GetKernelInterface( - const std::vector &provider, const schema::Primitive *primitive) { +std::shared_ptr RegisterKernelInterface::GetKernelInterface(const std::vector &provider, + const schema::Primitive *primitive, + const kernel::Kernel *kernel) { #ifndef CUSTOM_KERNEL_REGISTRY_CLIP - return KernelInterfaceRegistry::Instance()->GetKernelInterface(CharToString(provider), primitive); + return KernelInterfaceRegistry::Instance()->GetKernelInterface(CharToString(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 90f85374d3a..c188f202039 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 8cde2a63a4a..ae6657daf45 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/scheduler.cc b/mindspore/lite/src/scheduler.cc index a87c42b0af3..07aac24690a 100644 --- a/mindspore/lite/src/scheduler.cc +++ b/mindspore/lite/src/scheduler.cc @@ -1376,6 +1376,9 @@ kernel::LiteKernel *Scheduler::ScheduleNodeToKernel(const lite::Model::Node *src SetKernelTensorDataType(kernel); kernel->set_name(src_node->name_); + if (kernel->kernel() != nullptr) { + kernel->kernel()->SetConfig(config_info_); + } return kernel; } diff --git a/mindspore/lite/src/scheduler.h b/mindspore/lite/src/scheduler.h index 69be1ab474f..d1e0d73c0aa 100644 --- a/mindspore/lite/src/scheduler.h +++ b/mindspore/lite/src/scheduler.h @@ -59,6 +59,9 @@ class Scheduler { ~Scheduler() = default; int Schedule(std::vector *dst_kernels); void SetupSchedulerCb(std::unique_ptr cb) { sched_cb_ = std::move(cb); } + void SetConfig(const std::map> *config_info) { + config_info_ = config_info; + } private: int SchedulePreProcess(); @@ -165,6 +168,7 @@ class Scheduler { #endif int schema_version_ = SCHEMA_VERSION::SCHEMA_CUR; std::map *execution_plan_ = nullptr; + const std::map> *config_info_ = nullptr; }; } // namespace mindspore::lite diff --git a/mindspore/lite/src/sub_graph_kernel.cc b/mindspore/lite/src/sub_graph_kernel.cc index 36ea65d37b1..b03d82d8aaa 100644 --- a/mindspore/lite/src/sub_graph_kernel.cc +++ b/mindspore/lite/src/sub_graph_kernel.cc @@ -92,7 +92,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(); diff --git a/mindspore/lite/test/st/mix_data_type_test.cc b/mindspore/lite/test/st/mix_data_type_test.cc index 79af9ac5215..cb96cb3f408 100644 --- a/mindspore/lite/test/st/mix_data_type_test.cc +++ b/mindspore/lite/test/st/mix_data_type_test.cc @@ -51,10 +51,10 @@ TEST_F(MixDataTypeTest, Config1) { std::string filename = "MixDataTypeTestConfig"; std::string sectionname = "execution_plan"; - std::map config_info; - ret = lite::GetSectionInfoFromConfigFile(filename, sectionname, &config_info); + std::map> configs; + ret = lite::GetAllSectionInfoFromConfigFile(filename, &configs); ASSERT_EQ(ret, 0); - + std::map config_info = configs[sectionname]; ASSERT_EQ(config_info.size(), 2); auto info0 = config_info.at("op1");