op.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#ifndef NNVM_OP_H_
#define NNVM_OP_H_
 
#include <dmlc/parameter.h>
 
#include <functional>
#include <limits>
#include <string>
#include <typeinfo>
#include <utility>
#include <vector>
 
#include "base.h"
#include "c_api.h"
 
namespace nnvm {
 
// forward declarations
class Node;
struct NodeAttrs;
template <typename ValueType>
class OpMap;
class OpGroup;
class OpRegistryEntry;
using dmlc::ParamFieldInfo;
 
static const uint32_t kVarg = std::numeric_limits<uint32_t>::max();
 
class NNVM_DLL Op {
 public:
  std::string name;
  std::string description;
  /* \brief description of inputs and keyword arguments*/
  std::vector<ParamFieldInfo> arguments;
  uint32_t num_inputs = 1;
  uint32_t num_outputs = 1;
  uint32_t support_level = 10;
  std::function<uint32_t(const NodeAttrs& attrs)> get_num_outputs = nullptr;
  std::function<uint32_t(const NodeAttrs& attrs)> get_num_inputs = nullptr;
  std::function<void(NodeAttrs* attrs)> attr_parser = nullptr;
  // function fields.
  inline Op& describe(const std::string& descr);  // NOLINT(*)
  inline Op& add_argument(const std::string& name, const std::string& type,
                          const std::string& description);
  inline Op& add_arguments(const std::vector<ParamFieldInfo>& args);
  inline Op& set_num_inputs(uint32_t n);  // NOLINT(*)
  inline Op& set_support_level(uint32_t level);  // NOLINT(*)
  inline Op& set_num_inputs(std::function<uint32_t(const NodeAttrs& attr)> fn);  // NOLINT(*)
  inline Op& set_num_outputs(uint32_t n);  // NOLINT(*)
  inline Op& set_num_outputs(std::function<uint32_t(const NodeAttrs& attr)> fn);  // NOLINT(*)
  inline Op& set_attr_parser(std::function<void(NodeAttrs* attrs)> fn);  // NOLINT(*)
  template <typename ValueType>
  inline Op& set_attr(const std::string& attr_name,  // NOLINT(*)
                      const ValueType& value, int plevel = 10);
  Op& add_alias(const std::string& alias);  // NOLINT(*)
  Op& include(const std::string& group_name);
  static const Op* Get(const std::string& op_name);
  template <typename ValueType>
  static const OpMap<ValueType>& GetAttr(const std::string& attr_name);
 
 private:
  template <typename ValueType>
  friend class OpMap;
  friend class OpGroup;
  friend class dmlc::Registry<Op>;
  // Program internal unique index of operator.
  // Used to help index the program.
  uint32_t index_{0};
  // internal constructor
  Op();
  // get const reference to certain attribute
  static const any* GetAttrMap(const std::string& key);
  // update the attribute OpMap
  static void UpdateAttrMap(const std::string& key, std::function<void(any*)> updater);
  // add a trigger based on tag matching on certain tag attribute
  // This will apply trigger on all the op such that
  // include the corresponding group.
  // The trigger will also be applied to all future registrations
  // that calls include
  static void AddGroupTrigger(const std::string& group_name, std::function<void(Op*)> trigger);
};
 
template <typename ValueType>
class OpMap {
 public:
  inline const ValueType& operator[](const Op* op) const;
  inline const ValueType& get(const Op* op, const ValueType& def_value) const;
  inline int count(const Op* op) const;
 
  inline bool contains(const Op* op) const;
 
 private:
  friend class Op;
  // internal attribute name
  std::string attr_name_;
  // internal data
  std::vector<std::pair<ValueType, int>> data_;
  OpMap() = default;
};
 
class OpGroup {
 public:
  std::string group_name;
  template <typename ValueType>
  inline OpGroup& set_attr(const std::string& attr_name,  // NOLINT(*)
                           const ValueType& value, int plevel = 1);
};
 
// internal macros to make
#define NNVM_REGISTER_VAR_DEF(OpName) \
  static DMLC_ATTRIBUTE_UNUSED ::nnvm::Op& __make_##NnvmOp##_##OpName
 
#define NNVM_REGISTER_GVAR_DEF(TagName) \
  static DMLC_ATTRIBUTE_UNUSED ::nnvm::OpGroup __make_##NnvmOpGroup##_##TagName
 
#define NNVM_REGISTER_OP(OpName)                                \
  DMLC_STR_CONCAT(NNVM_REGISTER_VAR_DEF(OpName), __COUNTER__) = \
      ::dmlc::Registry<::nnvm::Op>::Get()->__REGISTER_OR_GET__(#OpName)
 
#define NNVM_REGISTER_OP_GROUP(GroupName) \
  DMLC_STR_CONCAT(NNVM_REGISTER_GVAR_DEF(GroupName), __COUNTER__) = ::nnvm::OpGroup { #GroupName }
 
// implementations of template functions after this.
// member function of Op
template <typename ValueType>
inline const OpMap<ValueType>& Op::GetAttr(const std::string& key) {
  const any* ref = GetAttrMap(key);
  if (ref == nullptr) {
    // update the attribute map of the key by creating new empty OpMap
    UpdateAttrMap(key, [key](any* pmap) {
      // use callback so it is in lockscope
      if (pmap->empty()) {
        OpMap<ValueType> pm;
        pm.attr_name_ = key;
        *pmap = std::move(pm);
      }
    });
    ref = GetAttrMap(key);
  }
  return nnvm::get<OpMap<ValueType>>(*ref);
}
 
template <typename ValueType>
inline Op& Op::set_attr(  // NOLINT(*)
    const std::string& attr_name, const ValueType& value, int plevel) {
  CHECK_GT(plevel, 0) << "plevel in set_attr must be greater than 0";
  // update the attribute map of the key by creating new empty if needed.
  UpdateAttrMap(attr_name, [this, attr_name, value, plevel](any* pmap) {
    // the callback is in lockscope so is threadsafe.
    if (pmap->empty()) {
      OpMap<ValueType> pm;
      pm.attr_name_ = attr_name;
      *pmap = std::move(pm);
    }
    CHECK(pmap->type() == typeid(OpMap<ValueType>))
        << "Attribute " << attr_name << " of operator " << this->name
        << " is registered as inconsistent types"
        << " previously " << pmap->type().name() << " current " << typeid(OpMap<ValueType>).name();
    std::vector<std::pair<ValueType, int>>& vec = nnvm::get<OpMap<ValueType>>(*pmap).data_;
    // resize the value type.
    if (vec.size() <= index_) {
      vec.resize(index_ + 1, std::make_pair(ValueType(), 0));
    }
    std::pair<ValueType, int>& p = vec[index_];
    CHECK(p.second != plevel) << "Attribute " << attr_name << " of operator " << this->name
                              << " is already registered with same plevel=" << plevel;
    if (p.second < plevel) {
      vec[index_] = std::make_pair(value, plevel);
    }
  });
  return *this;
}
 
inline Op& Op::describe(const std::string& descr) {  // NOLINT(*)
  this->description = descr;
  return *this;
}
 
inline Op& Op::add_argument(const std::string& name, const std::string& type,
                            const std::string& description) {
  arguments.push_back({name, type, type, description});
  return *this;
}
 
inline Op& Op::add_arguments(const std::vector<ParamFieldInfo>& args) {
  this->arguments.insert(arguments.end(), args.begin(), args.end());
  return *this;
}
 
inline Op& Op::set_num_inputs(uint32_t n) {  // NOLINT(*)
  this->num_inputs = n;
  return *this;
}
 
inline Op& Op::set_support_level(uint32_t n) {  // NOLINT(*)
  this->support_level = n;
  return *this;
}
 
inline Op& Op::set_num_inputs(std::function<uint32_t(const NodeAttrs& attr)> fn) {  // NOLINT(*)
  this->get_num_inputs = fn;
  return *this;
}
 
inline Op& Op::set_num_outputs(uint32_t n) {  // NOLINT(*)
  this->num_outputs = n;
  return *this;
}
 
inline Op& Op::set_num_outputs(std::function<uint32_t(const NodeAttrs& attr)> fn) {  // NOLINT(*)
  this->get_num_outputs = fn;
  return *this;
}
 
inline Op& Op::set_attr_parser(std::function<void(NodeAttrs* attrs)> fn) {  // NOLINT(*)
  this->attr_parser = fn;
  return *this;
}
 
// member functions of OpMap
template <typename ValueType>
inline int OpMap<ValueType>::count(const Op* op) const {
  if (contains(op)) {
    return 1;
  } else {
    return 0;
  }
}
 
template <typename ValueType>
inline bool OpMap<ValueType>::contains(const Op* op) const {
  if (op == nullptr) {
    return false;
  }
  const uint32_t idx = op->index_;
  return idx < data_.size() ? (data_[idx].second != 0) : false;
}
 
template <typename ValueType>
inline const ValueType& OpMap<ValueType>::operator[](const Op* op) const {
  CHECK(op != nullptr);
  const uint32_t idx = op->index_;
  CHECK(idx < data_.size() && data_[idx].second)
      << "Attribute " << attr_name_ << " has not been registered for Operator " << op->name;
  return data_[idx].first;
}
 
template <typename ValueType>
inline const ValueType& OpMap<ValueType>::get(const Op* op, const ValueType& def_value) const {
  if (op == nullptr) return def_value;
  const uint32_t idx = op->index_;
  if (idx < data_.size() && data_[idx].second) {
    return data_[idx].first;
  } else {
    return def_value;
  }
}
 
template <typename ValueType>
inline OpGroup& OpGroup::set_attr(const std::string& attr_name, const ValueType& value,
                                  int plevel) {
  auto trigger = [attr_name, value, plevel](Op* op) {
    op->set_attr<ValueType>(attr_name, value, plevel);
  };
  Op::AddGroupTrigger(group_name, trigger);
  return *this;
}
 
}  // namespace nnvm
 
#endif  // NNVM_OP_H_