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 #include "tools/optimizer/parallel/operator_info_register.h"
18 #include <utility>
19 namespace mindspore {
20 namespace opt {
21
22 // find the only key of operator_info
operator <(const SplitOpKey & rhs) const23 bool SplitOpKey::operator<(const SplitOpKey &rhs) const {
24 return std::tie(this->op_type_, this->data_type_, this->is_depth_wise_) <
25 std::tie(rhs.op_type_, rhs.data_type_, rhs.is_depth_wise_);
26 }
27
ToString() const28 std::string SplitOpKey::ToString() const {
29 std::ostringstream split_info;
30 split_info << "op_type: " << schema::EnumNamePrimitiveType(static_cast<schema::PrimitiveType>(op_type_)) << "\t"
31 << " data_type_: " << data_type_ << " is_depth_wise: " << is_depth_wise_ << "\t";
32 return split_info.str();
33 }
34
GeInstance()35 OperatorInfoFactory *OperatorInfoFactory::GeInstance() {
36 static OperatorInfoFactory factory;
37 return &factory;
38 }
39
RegisterOperatorInfo(schema::PrimitiveType operator_type,TypeId type_id,bool is_depth_wise,const OperatorInfoCreatorFunc & creator_func)40 void OperatorInfoFactory::RegisterOperatorInfo(schema::PrimitiveType operator_type, TypeId type_id, bool is_depth_wise,
41 const OperatorInfoCreatorFunc &creator_func) {
42 // create a key to find the only create function
43 SplitOpKey op_key(operator_type, type_id, is_depth_wise);
44 if (operator_info_map_.find(op_key) != operator_info_map_.end()) {
45 MS_LOG(ERROR) << " Operator already exist " << op_key.ToString();
46 return;
47 }
48 this->operator_info_map_.insert(std::pair<SplitOpKey, OperatorInfoCreatorFunc>(op_key, creator_func));
49 }
50
FindOperatorInfo(const SplitOpKey & op_key)51 OperatorInfoCreatorFunc OperatorInfoFactory::FindOperatorInfo(const SplitOpKey &op_key) {
52 auto iterator = this->operator_info_map_.find(op_key);
53 if (iterator != this->operator_info_map_.end()) {
54 return iterator->second;
55 }
56 MS_LOG(ERROR) << "operator_info do not register" << op_key.ToString();
57 return nullptr;
58 }
59
OperatorInfoRegister(schema::PrimitiveType operator_type,TypeId type_id,bool is_depth_wise,const OperatorInfoCreatorFunc & creator_func)60 OperatorInfoRegister::OperatorInfoRegister(schema::PrimitiveType operator_type, TypeId type_id, bool is_depth_wise,
61 const OperatorInfoCreatorFunc &creator_func) {
62 OperatorInfoFactory::GeInstance()->RegisterOperatorInfo(operator_type, type_id, is_depth_wise, creator_func);
63 }
64 } // namespace opt
65 } // namespace mindspore
66