43 template<
typename ValueType>
46 class OpRegistryEntry;
47 using dmlc::ParamFieldInfo;
50 static const uint32_t kVarg = std::numeric_limits<uint32_t>::max();
122 uint32_t num_inputs = 1;
130 uint32_t num_outputs = 1;
136 uint32_t support_level = 10;
142 std::function<uint32_t(const NodeAttrs& attrs)> get_num_outputs =
nullptr;
148 std::function<uint32_t(const NodeAttrs& attrs)> get_num_inputs =
nullptr;
181 std::function<void(NodeAttrs* attrs)> attr_parser =
nullptr;
189 inline Op& describe(
const std::string& descr);
197 inline Op& add_argument(
const std::string &name,
198 const std::string &type,
199 const std::string &description);
205 inline Op& add_arguments(
const std::vector<ParamFieldInfo> &args);
211 inline Op& set_num_inputs(uint32_t n);
217 inline Op& set_support_level(uint32_t level);
223 inline Op& set_num_inputs(std::function<uint32_t (
const NodeAttrs& attr)> fn);
229 inline Op& set_num_outputs(uint32_t n);
235 inline Op& set_num_outputs(std::function<uint32_t (
const NodeAttrs& attr)> fn);
241 inline Op& set_attr_parser(std::function<
void (
NodeAttrs* attrs)> fn);
255 template<
typename ValueType>
256 inline Op& set_attr(
const std::string& attr_name,
257 const ValueType& value,
265 Op& add_alias(
const std::string& alias);
273 Op& include(
const std::string& group_name);
280 static const Op* Get(
const std::string& op_name);
288 template<
typename ValueType>
292 template<
typename ValueType>
302 static const any* GetAttrMap(
const std::string& key);
304 static void UpdateAttrMap(
const std::string& key,
305 std::function<
void(any*)> updater);
311 static void AddGroupTrigger(
const std::string& group_name,
312 std::function<
void(
Op*)> trigger);
320 template<
typename ValueType>
328 inline const ValueType& operator[](
const Op* op)
const;
335 inline const ValueType&
get(
const Op* op,
const ValueType& def_value)
const;
341 inline int count(
const Op* op)
const;
348 inline bool contains(
const Op* op)
const;
353 std::string attr_name_;
355 std::vector<std::pair<ValueType, int> > data_;
380 template<
typename ValueType>
381 inline OpGroup& set_attr(
const std::string& attr_name,
382 const ValueType& value,
387 #define NNVM_REGISTER_VAR_DEF(OpName) \ 388 static DMLC_ATTRIBUTE_UNUSED ::nnvm::Op & __make_ ## NnvmOp ## _ ## OpName 390 #define NNVM_REGISTER_GVAR_DEF(TagName) \ 391 static DMLC_ATTRIBUTE_UNUSED ::nnvm::OpGroup __make_ ## NnvmOpGroup ## _ ## TagName 408 #define NNVM_REGISTER_OP(OpName) \ 409 DMLC_STR_CONCAT(NNVM_REGISTER_VAR_DEF(OpName), __COUNTER__) = \ 410 ::dmlc::Registry<::nnvm::Op>::Get()->__REGISTER_OR_GET__(#OpName) 433 #define NNVM_REGISTER_OP_GROUP(GroupName) \ 434 DMLC_STR_CONCAT(NNVM_REGISTER_GVAR_DEF(GroupName), __COUNTER__) = \ 435 ::nnvm::OpGroup {#GroupName} 439 template<
typename ValueType>
441 const any* ref = GetAttrMap(key);
442 if (ref ==
nullptr) {
444 UpdateAttrMap(key, [key](any* pmap) {
449 *pmap = std::move(pm);
452 ref = GetAttrMap(key);
454 return nnvm::get<OpMap<ValueType> >(*ref);
457 template<
typename ValueType>
459 const std::string& attr_name,
460 const ValueType& value,
463 <<
"plevel in set_attr must be greater than 0";
465 UpdateAttrMap(attr_name,
466 [
this, attr_name, value, plevel](any* pmap) {
470 pm.attr_name_ = attr_name;
471 *pmap = std::move(pm);
474 <<
"Attribute " << attr_name
475 <<
" of operator " << this->name
476 <<
" is registered as inconsistent types" 477 <<
" previously " << pmap->type().name()
479 std::vector<std::pair<ValueType, int> >& vec =
480 nnvm::get<OpMap<ValueType> >(*pmap).data_;
482 if (vec.size() <= index_) {
483 vec.resize(index_ + 1,
484 std::make_pair(ValueType(), 0));
486 std::pair<ValueType, int>& p = vec[index_];
487 CHECK(p.second != plevel)
488 <<
"Attribute " << attr_name
489 <<
" of operator " << this->name
490 <<
" is already registered with same plevel=" << plevel;
491 if (p.second < plevel) {
492 vec[index_] = std::make_pair(value, plevel);
500 this->description = descr;
505 const std::string &type,
506 const std::string &description) {
507 arguments.push_back({name, type, type, description});
512 this->arguments.insert(arguments.end(), args.begin(), args.end());
517 this->num_inputs = n;
522 this->support_level = n;
527 this->get_num_inputs = fn;
532 this->num_outputs = n;
537 this->get_num_outputs = fn;
542 this->attr_parser = fn;
547 template<
typename ValueType>
556 template<
typename ValueType>
561 const uint32_t idx = op->index_;
562 return idx < data_.size() ? (data_[idx].second != 0) :
false;
565 template<
typename ValueType>
567 CHECK(op !=
nullptr);
568 const uint32_t idx = op->index_;
569 CHECK(idx < data_.size() && data_[idx].second)
570 <<
"Attribute " << attr_name_
571 <<
" has not been registered for Operator " << op->
name;
572 return data_[idx].first;
575 template<
typename ValueType>
577 if (op ==
nullptr)
return def_value;
578 const uint32_t idx = op->index_;
579 if (idx < data_.size() && data_[idx].second) {
580 return data_[idx].first;
586 template<
typename ValueType>
588 const ValueType& value,
590 auto trigger = [attr_name, value, plevel](
Op* op) {
591 op->
set_attr<ValueType>(attr_name, value, plevel);
593 Op::AddGroupTrigger(group_name, trigger);
std::vector< ParamFieldInfo > arguments
Definition: op.h:114
Registry class. Registry can be used to register global singletons. The most commonly use case are fa...
Definition: registry.h:27
Op & add_argument(const std::string &name, const std::string &type, const std::string &description)
Add argument information to the function.
Definition: op.h:504
std::string description
detailed description of the operator This can be used to generate docstring automatically for the ope...
Definition: op.h:112
Op & set_attr_parser(std::function< void(NodeAttrs *attrs)> fn)
Set the attr_parser function.
Definition: op.h:541
The attributes of the current operation node. Usually are additional parameters like axis...
Definition: node.h:120
const ValueType & operator[](const Op *op) const
get the corresponding value element at op
Definition: op.h:566
static const OpMap< ValueType > & GetAttr(const std::string &attr_name)
Get additional registered attribute about operators. If nothing has been registered, an empty OpMap will be returned.
Definition: op.h:440
Op & describe(const std::string &descr)
setter function during registration Set the description of operator
Definition: op.h:499
Op & set_num_outputs(uint32_t n)
Set the num_outputs.
Definition: op.h:531
OpGroup & set_attr(const std::string &attr_name, const ValueType &value, int plevel=1)
Register additional attributes to operator group.
Definition: op.h:587
Op & set_support_level(uint32_t level)
Set the support level of op.
Definition: op.h:521
#define NNVM_DLL
NNVM_DLL prefix for windows.
Definition: c_api.h:38
Op & set_num_inputs(uint32_t n)
Set the num_inputs.
Definition: op.h:516
bool contains(const Op *op) const
Check if the map has op as key.
Definition: op.h:557
auxiliary data structure used to set attributes to a group of operators
Definition: op.h:363
Op & add_arguments(const std::vector< ParamFieldInfo > &args)
Append list if arguments to the end.
Definition: op.h:511
std::string name
name of the operator
Definition: op.h:107
nnvm::Op Op
operator structure from NNVM
Definition: base.h:99
int count(const Op *op) const
Check if the map has op as key.
Definition: op.h:548
A map data structure that takes Op* as key and returns ValueType.
Definition: op.h:44
const ValueType & get(const Op *op, const ValueType &def_value) const
get the corresponding value element at op with default value.
Definition: op.h:576
std::string group_name
the tag key to be matched
Definition: op.h:366
C API of NNVM symbolic construction and pass. Enables construction and transformation of Graph in any...
Op & set_attr(const std::string &attr_name, const ValueType &value, int plevel=10)
Register additional attributes to operator.
Definition: op.h:458
Provide lightweight util to do parameter setup and checking.
Configuration of nnvm as well as basic data structure.
Operator structure.
Definition: op.h:104