1 /** 2 * Copyright 2021 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef MINDSPORE_LITE_SRC_OPS_MS_OPS_UTILS_H_ 18 #define MINDSPORE_LITE_SRC_OPS_MS_OPS_UTILS_H_ 19 20 #include <map> 21 #include <string> 22 #include <memory> 23 #include "src/ops/ops_func_declare.h" 24 25 #ifdef PRIMITIVE_WRITEABLE 26 namespace mindspore { 27 namespace lite { 28 typedef std::unique_ptr<schema::PrimitiveT> (*PrimitiveTCreator)(const AnfNodePtr &node); 29 30 class MSOpsRegistry { 31 public: GetInstance()32 static MSOpsRegistry *GetInstance() { 33 static MSOpsRegistry registry; 34 return ®istry; 35 } InsertPrimitiveTMap(const std::string & name,PrimitiveTCreator creator)36 void InsertPrimitiveTMap(const std::string &name, PrimitiveTCreator creator) { primitive_creators[name] = creator; } GetPrimitiveCreator(const std::string & name)37 PrimitiveTCreator GetPrimitiveCreator(const std::string &name) { 38 if (primitive_creators.find(name) != primitive_creators.end()) { 39 return primitive_creators[name]; 40 } else { 41 MS_LOG(ERROR) << "Unsupported primitive type in Create: " << name; 42 return nullptr; 43 } 44 } 45 46 protected: 47 std::map<std::string, PrimitiveTCreator> primitive_creators; 48 }; 49 50 class RegistryMSOps { 51 public: RegistryMSOps(const std::string & name,PrimitiveTCreator creator)52 RegistryMSOps(const std::string &name, PrimitiveTCreator creator) { 53 MSOpsRegistry::GetInstance()->InsertPrimitiveTMap(name, creator); 54 } 55 ~RegistryMSOps() = default; 56 }; 57 58 std::unique_ptr<schema::PrimitiveT> GetPrimitiveT(const mindspore::AnfNodePtr &node); 59 } // namespace lite 60 } // namespace mindspore 61 #endif 62 63 #endif // MINDSPORE_LITE_SRC_OPS_MS_OPS_UTILS_H_ 64