42 template<
typename ValueType>
45 class OpRegistryEntry;
46 using dmlc::ParamFieldInfo;
49 static const uint32_t kVarg = std::numeric_limits<uint32_t>::max();
121 uint32_t num_inputs = 1;
129 uint32_t num_outputs = 1;
135 uint32_t support_level = 10;
141 std::function<uint32_t(const NodeAttrs& attrs)> get_num_outputs =
nullptr;
147 std::function<uint32_t(const NodeAttrs& attrs)> get_num_inputs =
nullptr;
180 std::function<void(NodeAttrs* attrs)> attr_parser =
nullptr;
188 inline Op& describe(
const std::string& descr);
196 inline Op& add_argument(
const std::string &name,
197 const std::string &type,
198 const std::string &description);
204 inline Op& add_arguments(
const std::vector<ParamFieldInfo> &args);
210 inline Op& set_num_inputs(uint32_t n);
216 inline Op& set_support_level(uint32_t level);
222 inline Op& set_num_inputs(std::function<uint32_t (
const NodeAttrs& attr)> fn);
228 inline Op& set_num_outputs(uint32_t n);
234 inline Op& set_num_outputs(std::function<uint32_t (
const NodeAttrs& attr)> fn);
240 inline Op& set_attr_parser(std::function<
void (
NodeAttrs* attrs)> fn);
254 template<
typename ValueType>
255 inline Op& set_attr(
const std::string& attr_name,
256 const ValueType& value,
264 Op& add_alias(
const std::string& alias);
272 Op& include(
const std::string& group_name);
279 static const Op* Get(
const std::string& op_name);
287 template<
typename ValueType>
291 template<
typename ValueType>
301 static const any* GetAttrMap(
const std::string& key);
303 static void UpdateAttrMap(
const std::string& key,
304 std::function<
void(any*)> updater);
310 static void AddGroupTrigger(
const std::string& group_name,
311 std::function<
void(
Op*)> trigger);
319 template<
typename ValueType>
327 inline const ValueType& operator[](
const Op* op)
const;
334 inline const ValueType&
get(
const Op* op,
const ValueType& def_value)
const;
340 inline int count(
const Op* op)
const;
347 inline bool contains(
const Op* op)
const;
352 std::string attr_name_;
354 std::vector<std::pair<ValueType, int> > data_;
379 template<
typename ValueType>
380 inline OpGroup& set_attr(
const std::string& attr_name,
381 const ValueType& value,
386 #define NNVM_REGISTER_VAR_DEF(OpName) \ 387 static DMLC_ATTRIBUTE_UNUSED ::nnvm::Op & __make_ ## NnvmOp ## _ ## OpName 389 #define NNVM_REGISTER_GVAR_DEF(TagName) \ 390 static DMLC_ATTRIBUTE_UNUSED ::nnvm::OpGroup __make_ ## NnvmOpGroup ## _ ## TagName 407 #define NNVM_REGISTER_OP(OpName) \ 408 DMLC_STR_CONCAT(NNVM_REGISTER_VAR_DEF(OpName), __COUNTER__) = \ 409 ::dmlc::Registry<::nnvm::Op>::Get()->__REGISTER_OR_GET__(#OpName) 432 #define NNVM_REGISTER_OP_GROUP(GroupName) \ 433 DMLC_STR_CONCAT(NNVM_REGISTER_GVAR_DEF(GroupName), __COUNTER__) = \ 434 ::nnvm::OpGroup {#GroupName} 438 template<
typename ValueType>
440 const any* ref = GetAttrMap(key);
441 if (ref ==
nullptr) {
443 UpdateAttrMap(key, [key](any* pmap) {
448 *pmap = std::move(pm);
451 ref = GetAttrMap(key);
453 return nnvm::get<OpMap<ValueType> >(*ref);
456 template<
typename ValueType>
458 const std::string& attr_name,
459 const ValueType& value,
462 <<
"plevel in set_attr must be greater than 0";
464 UpdateAttrMap(attr_name,
465 [
this, attr_name, value, plevel](any* pmap) {
469 pm.attr_name_ = attr_name;
470 *pmap = std::move(pm);
473 <<
"Attribute " << attr_name
474 <<
" of operator " << this->name
475 <<
" is registered as inconsistent types" 476 <<
" previously " << pmap->type().name()
478 std::vector<std::pair<ValueType, int> >& vec =
479 nnvm::get<OpMap<ValueType> >(*pmap).data_;
481 if (vec.size() <= index_) {
482 vec.resize(index_ + 1,
483 std::make_pair(ValueType(), 0));
485 std::pair<ValueType, int>& p = vec[index_];
486 CHECK(p.second != plevel)
487 <<
"Attribute " << attr_name
488 <<
" of operator " << this->name
489 <<
" is already registered with same plevel=" << plevel;
490 if (p.second < plevel) {
491 vec[index_] = std::make_pair(value, plevel);
499 this->description = descr;
504 const std::string &type,
505 const std::string &description) {
506 arguments.push_back({name, type, type, description});
511 this->arguments.insert(arguments.end(), args.begin(), args.end());
516 this->num_inputs = n;
521 this->support_level = n;
526 this->get_num_inputs = fn;
531 this->num_outputs = n;
536 this->get_num_outputs = fn;
541 this->attr_parser = fn;
546 template<
typename ValueType>
555 template<
typename ValueType>
560 const uint32_t idx = op->index_;
561 return idx < data_.size() ? (data_[idx].second != 0) :
false;
564 template<
typename ValueType>
566 CHECK(op !=
nullptr);
567 const uint32_t idx = op->index_;
568 CHECK(idx < data_.size() && data_[idx].second)
569 <<
"Attribute " << attr_name_
570 <<
" has not been registered for Operator " << op->
name;
571 return data_[idx].first;
574 template<
typename ValueType>
576 if (op ==
nullptr)
return def_value;
577 const uint32_t idx = op->index_;
578 if (idx < data_.size() && data_[idx].second) {
579 return data_[idx].first;
585 template<
typename ValueType>
587 const ValueType& value,
589 auto trigger = [attr_name, value, plevel](
Op* op) {
590 op->
set_attr<ValueType>(attr_name, value, plevel);
592 Op::AddGroupTrigger(group_name, trigger);
std::vector< ParamFieldInfo > arguments
Definition: op.h:113
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:503
std::string description
detailed description of the operator This can be used to generate docstring automatically for the ope...
Definition: op.h:111
Op & set_attr_parser(std::function< void(NodeAttrs *attrs)> fn)
Set the attr_parser function.
Definition: op.h:540
The attributes of the current operation node. Usually are additional parameters like axis...
Definition: node.h:119
const ValueType & operator[](const Op *op) const
get the corresponding value element at op
Definition: op.h:565
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:439
Op & describe(const std::string &descr)
setter function during registration Set the description of operator
Definition: op.h:498
Op & set_num_outputs(uint32_t n)
Set the num_outputs.
Definition: op.h:530
OpGroup & set_attr(const std::string &attr_name, const ValueType &value, int plevel=1)
Register additional attributes to operator group.
Definition: op.h:586
Op & set_support_level(uint32_t level)
Set the support level of op.
Definition: op.h:520
#define NNVM_DLL
NNVM_DLL prefix for windows.
Definition: c_api.h:37
Op & set_num_inputs(uint32_t n)
Set the num_inputs.
Definition: op.h:515
bool contains(const Op *op) const
Check if the map has op as key.
Definition: op.h:556
auxiliary data structure used to set attributes to a group of operators
Definition: op.h:362
Op & add_arguments(const std::vector< ParamFieldInfo > &args)
Append list if arguments to the end.
Definition: op.h:510
std::string name
name of the operator
Definition: op.h:106
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:547
A map data structure that takes Op* as key and returns ValueType.
Definition: op.h:43
const ValueType & get(const Op *op, const ValueType &def_value) const
get the corresponding value element at op with default value.
Definition: op.h:575
std::string group_name
the tag key to be matched
Definition: op.h:365
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:457
Provide lightweight util to do parameter setup and checking.
Configuration of nnvm as well as basic data structure.
Operator structure.
Definition: op.h:103