1 /** 2 * Copyright 2020 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 #ifndef MINDSPORE_CCSRC_BACKEND_OPTIMIZER_COMMON_CONST_INPUT_TO_ATTR_REGISTRY_H_ 17 #define MINDSPORE_CCSRC_BACKEND_OPTIMIZER_COMMON_CONST_INPUT_TO_ATTR_REGISTRY_H_ 18 #include <string> 19 #include <unordered_map> 20 #include <unordered_set> 21 22 #include "utils/ms_utils.h" 23 24 namespace mindspore { 25 namespace opt { 26 class ConstInputToAttrInfoRegister { 27 public: op_name_(op_name)28 explicit ConstInputToAttrInfoRegister(const std::string &op_name = "") : op_name_(op_name) {} 29 virtual ~ConstInputToAttrInfoRegister() = default; 30 SetConstInputToAttr(size_t input_index)31 ConstInputToAttrInfoRegister &SetConstInputToAttr(size_t input_index) { 32 (void)input_attr_set_.insert(input_index); 33 return *this; 34 } 35 SetConstInputToAttr(const std::unordered_set<size_t> & input_attr_set)36 ConstInputToAttrInfoRegister &SetConstInputToAttr(const std::unordered_set<size_t> &input_attr_set) { 37 (void)input_attr_set_.insert(input_attr_set.begin(), input_attr_set.end()); 38 return *this; 39 } 40 GetConstInputAttrInfo()41 const std::unordered_set<size_t> &GetConstInputAttrInfo() const { return input_attr_set_; } GetOpName()42 const std::string &GetOpName() const { return op_name_; } 43 44 private: 45 std::string op_name_; 46 std::unordered_set<size_t> input_attr_set_; 47 }; 48 49 class ConstInputToAttrInfoRegistry { 50 public: 51 static ConstInputToAttrInfoRegistry &Instance(); 52 void Register(const ConstInputToAttrInfoRegister ®); 53 void Register(const std::string &op_name, const std::unordered_set<size_t> &input_attr_set); 54 bool GetRegisterByOpName(const std::string &op_name, ConstInputToAttrInfoRegister *reg) const; 55 56 private: 57 ConstInputToAttrInfoRegistry(); 58 ~ConstInputToAttrInfoRegistry() = default; 59 DISABLE_COPY_AND_ASSIGN(ConstInputToAttrInfoRegistry) 60 std::unordered_map<std::string, ConstInputToAttrInfoRegister> op_input_to_attr_map_; 61 }; 62 63 struct ConstInputToAttrInfoReceiver { 64 // Note: This is implicitly converting constructor ConstInputToAttrInfoReceiverConstInputToAttrInfoReceiver65 ConstInputToAttrInfoReceiver(const ConstInputToAttrInfoRegister ®) { // NOLINT(runtime/explicit) 66 ConstInputToAttrInfoRegistry::Instance().Register(reg); 67 } 68 }; 69 } // namespace opt 70 71 #define REG_CONST_INPUT_TO_ATTR(op_name) REG_CONST_INPUT_TO_ATTR_UNIQ_HELPER(__COUNTER__, op_name) 72 #define REG_CONST_INPUT_TO_ATTR_UNIQ_HELPER(ctr, op_name) REG_CONST_INPUT_TO_ATTR_UNIQ(ctr, op_name) 73 #define REG_CONST_INPUT_TO_ATTR_UNIQ(ctr, op_name) \ 74 static ::mindspore::opt::ConstInputToAttrInfoReceiver g_const_input_to_attr_register_##ctr __attribute__((unused)) = \ 75 ::mindspore::opt::ConstInputToAttrInfoRegister(op_name) 76 } // namespace mindspore 77 78 #endif // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_COMMON_CONST_INPUT_TO_ATTR_REGISTRY_H_ 79