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 "include/registry/node_parser_registry.h"
18 #include <map>
19 #include <mutex>
20 #include <string>
21
22 namespace mindspore {
23 namespace registry {
24 namespace {
25 std::map<converter::FmkType, std::map<std::string, converter::NodeParserPtr>> node_parser_room;
26 std::mutex node_mutex;
27 } // namespace
NodeParserRegistry(converter::FmkType fmk_type,const std::string & node_type,const converter::NodeParserPtr & node_parser)28 NodeParserRegistry::NodeParserRegistry(converter::FmkType fmk_type, const std::string &node_type,
29 const converter::NodeParserPtr &node_parser) {
30 std::unique_lock<std::mutex> lock(node_mutex);
31 node_parser_room[fmk_type][node_type] = node_parser;
32 }
33
GetNodeParser(converter::FmkType fmk_type,const std::string & node_type)34 converter::NodeParserPtr NodeParserRegistry::GetNodeParser(converter::FmkType fmk_type, const std::string &node_type) {
35 auto iter_level1 = node_parser_room.find(fmk_type);
36 if (iter_level1 == node_parser_room.end()) {
37 return nullptr;
38 }
39 if (node_type.empty()) {
40 return nullptr;
41 }
42 auto iter_level2 = iter_level1->second.find(node_type);
43 if (iter_level2 == iter_level1->second.end()) {
44 return nullptr;
45 }
46 return iter_level2->second;
47 }
48 } // namespace registry
49 } // namespace mindspore
50