mxnet
op_map.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
27 #ifndef MXNET_CPP_OP_MAP_H_
28 #define MXNET_CPP_OP_MAP_H_
29 
30 #include <map>
31 #include <string>
32 #include "mxnet-cpp/base.h"
33 #include "dmlc/logging.h"
34 
35 namespace mxnet {
36 namespace cpp {
37 
43 class OpMap {
44  public:
48  inline OpMap() {
49  mx_uint num_symbol_creators = 0;
50  AtomicSymbolCreator *symbol_creators = nullptr;
51  int r =
52  MXSymbolListAtomicSymbolCreators(&num_symbol_creators, &symbol_creators);
53  CHECK_EQ(r, 0);
54  for (mx_uint i = 0; i < num_symbol_creators; i++) {
55  const char *name;
56  const char *description;
57  mx_uint num_args;
58  const char **arg_names;
59  const char **arg_type_infos;
60  const char **arg_descriptions;
61  const char *key_var_num_args;
62  r = MXSymbolGetAtomicSymbolInfo(symbol_creators[i], &name, &description,
63  &num_args, &arg_names, &arg_type_infos,
64  &arg_descriptions, &key_var_num_args);
65  CHECK_EQ(r, 0);
66  symbol_creators_[name] = symbol_creators[i];
67  }
68 
69  nn_uint num_ops;
70  const char **op_names;
71  r = NNListAllOpNames(&num_ops, &op_names);
72  CHECK_EQ(r, 0);
73  for (nn_uint i = 0; i < num_ops; i++) {
74  OpHandle handle;
75  r = NNGetOpHandle(op_names[i], &handle);
76  CHECK_EQ(r, 0);
77  op_handles_[op_names[i]] = handle;
78  }
79  }
80 
87  inline AtomicSymbolCreator GetSymbolCreator(const std::string &name) {
88  if (symbol_creators_.count(name) == 0)
89  return GetOpHandle(name);
90  return symbol_creators_[name];
91  }
92 
99  inline OpHandle GetOpHandle(const std::string &name) {
100  return op_handles_[name];
101  }
102 
103  private:
104  std::map<std::string, AtomicSymbolCreator> symbol_creators_;
105  std::map<std::string, OpHandle> op_handles_;
106 };
107 
108 } // namespace cpp
109 } // namespace mxnet
110 
111 #endif // MXNET_CPP_OP_MAP_H_
OpHandle GetOpHandle(const std::string &name)
Get an op handle with its name.
Definition: op_map.h:99
MXNET_DLL int MXSymbolListAtomicSymbolCreators(mx_uint *out_size, AtomicSymbolCreator **out_array)
list all the available AtomicSymbolEntry
OpMap instance holds a map of all the symbol creators so we can get symbol creators by name...
Definition: op_map.h:43
namespace of mxnet
Definition: base.h:118
MXNET_DLL int MXSymbolGetAtomicSymbolInfo(AtomicSymbolCreator creator, const char **name, const char **description, mx_uint *num_args, const char ***arg_names, const char ***arg_type_infos, const char ***arg_descriptions, const char **key_var_num_args, const char **return_type DEFAULT(NULL))
Get the detailed information about atomic symbol.
OpMap()
Create an Mxnet instance.
Definition: op_map.h:48
unsigned int mx_uint
manually define unsigned int
Definition: c_api.h:58
void * AtomicSymbolCreator
handle to a function that takes param and creates symbol
Definition: c_api.h:71
AtomicSymbolCreator GetSymbolCreator(const std::string &name)
Get a symbol creator with its name.
Definition: op_map.h:87