node_base.h

/*!
 *  Copyright (c) 2018 by Contributors
 * \file tvm/runtime/node_base.h
 * \brief Base data structure for Node.
 *
 * \note Node is not a runtime feature.
 *  This file only exposes the signature of NodePtr for PackedFunc.
 */
#ifndef TVM_RUNTIME_NODE_BASE_H_
#define TVM_RUNTIME_NODE_BASE_H_

#include <utility>
#include <atomic>

namespace tvm {

// forward declarations
template<typename T>
class NodePtr;
class Node;
class NodeRef;

/*!
 * \brief Base class of Node for runtime destructor purposes.
 *
 * Node is a reference counted object which is used to construct AST.
 * Each node is backed by a custom deleter, which deletes the object.
 * Do not call create raw Node pointer, always use tvm::make_node.
 *
 * \note In most cases, please inheritate tvm::Node.
 * \sa Node, NodePtr, make_node
 */
class NodeBase {
 public:
  /*!
   * \brief type of NodeBase deleter
   * \param self pointer to the NodeBase.
   */
  typedef void (*FDeleter)(NodeBase* self);

 protected:
  // default constructor and copy constructor
  NodeBase() {}
  // override the copy and assign constructors to do nothing.
  // This is to make sure only contents, but not deleter and ref_counter
  // are copied when a child class copies itself.
  NodeBase(const NodeBase& other) {  // NOLINT(*)
  }
  NodeBase(NodeBase&& other) {  // NOLINT(*)
  }
  NodeBase& operator=(const NodeBase& other) {  //NOLINT(*)
    return *this;
  }
  NodeBase& operator=(NodeBase&& other) {  //NOLINT(*)
    return *this;
  }

 private:
  /*! \brief Internal reference counter */
  std::atomic<int> ref_counter_{0};
  /*!
   * \brief deleter of this object to enable customized allocation.
   * If the deleter is nullptr, no deletion will be performed.
   * The creator of the Node must always set the deleter field properly.
   */
  FDeleter deleter_ = nullptr;
  // reference counting functions
  void IncRef() {
    ref_counter_.fetch_add(1, std::memory_order_relaxed);
  }
  void DecRef() {
    if (ref_counter_.fetch_sub(1, std::memory_order_release) == 1) {
      std::atomic_thread_fence(std::memory_order_acquire);
      if (this->deleter_ != nullptr) {
        (*this->deleter_)(this);
      }
    }
  }
  int use_count() const {
    return ref_counter_.load(std::memory_order_relaxed);
  }
  // friend declaration
  template<typename>
  friend class NodePtr;
  template<typename Y, typename... Args>
  friend NodePtr<Y> make_node(Args&&...);
};

/*!
 * \brief Smart pointer for Node containers,
 *  must be subclass of NodeBase
 * \tparam T the content data type.
 */
template<typename T>
class NodePtr {
 public:
  /*! \brief default constructor */
  NodePtr() {}
  /*! \brief default constructor */
  NodePtr(std::nullptr_t) {}  // NOLINT(*)
  /*!
   * \brief copy constructor
   * \param other The value to be moved
   */
  NodePtr(const NodePtr<T>& other)  // NOLINT(*)
      : NodePtr(other.data_) {
  }
  /*!
   * \brief copy constructor
   * \param other The value to be moved
   */
  template<typename Y>
  NodePtr(const NodePtr<Y>& other)  // NOLINT(*)
      : NodePtr(other.data_) {
    static_assert(std::is_base_of<T, Y>::value,
                  "can only assign of child class NodePtr to parent");
  }
  /*!
   * \brief move constructor
   * \param other The value to be moved
   */
  NodePtr(NodePtr<T>&& other) // NOLINT(*)
      : data_(other.data_) {
    other.data_ = nullptr;
  }
  /*!
   * \brief move constructor
   * \param other The value to be moved
   */
  template<typename Y>
  NodePtr(NodePtr<Y>&& other)  // NOLINT(*)
      : data_(other.data_) {
    static_assert(std::is_base_of<T, Y>::value,
                  "can only assign of child class NodePtr to parent");
    other.data_ = nullptr;
  }
  /*! \brief destructor */
  ~NodePtr() {
    this->reset();
  }
  /*!
   * \brief Swap this array with another NDArray
   * \param other The other NDArray
   */
  void swap(NodePtr<T>& other) {  // NOLINT(*)
    std::swap(data_, other.data_);
  }
  /*!
   * \return Get the content of the pointer
   */
  T* get() const {
    return static_cast<T*>(data_);
  }
  /*!
   * \return The pointer
   */
  T* operator->() const {
    return get();
  }
  /*!
   * \return The reference
   */
  T& operator*() const { // NOLINT(*)
    return *get();
  }
  /*!
   * \brief copy assignmemt
   * \param other The value to be assigned.
   * \return reference to self.
   */
  NodePtr<T>& operator=(const NodePtr<T>& other) {  // NOLINT(*)
    // takes in plane operator to enable copy elison.
    // copy-and-swap idiom
    NodePtr(other).swap(*this);  // NOLINT(*)
    return *this;
  }
  /*!
   * \brief move assignmemt
   * \param other The value to be assigned.
   * \return reference to self.
   */
  NodePtr<T>& operator=(NodePtr<T>&& other) {  // NOLINT(*)
    // copy-and-swap idiom
    NodePtr(std::move(other)).swap(*this); // NOLINT(*)
    return *this;
  }
  /*! \brief reset the content of ptr to be nullptr */
  void reset() {
    if (data_ != nullptr) {
      data_->DecRef();
      data_ = nullptr;
    }
  }
  /*! \return The use count of the ptr, for debug purposes */
  int use_count() const {
    return data_ != nullptr ? data_->use_count() : 0;
  }
  /*! \return whether the reference is unique */
  bool unique() const {
    return data_ != nullptr && data_->use_count() == 1;
  }
  /*! \return Whether two NodePtr do not equals each other */
  bool operator==(const NodePtr<T>& other) const {
    return data_ == other.data_;
  }
  /*! \return Whether two NodePtr equals each other */
  bool operator!=(const NodePtr<T>& other) const {
    return data_ != other.data_;
  }
  /*! \return Whether the pointer is nullptr */
  bool operator==(std::nullptr_t null) const {
    return data_ == nullptr;
  }
  /*! \return Whether the pointer is not nullptr */
  bool operator!=(std::nullptr_t null) const {
    return data_ != nullptr;
  }

 private:
  /*! \brief internal pointer field */
  NodeBase* data_{nullptr};
  /*!
   * \brief constructor from NodeBase
   * \param data The node base pointer
   */
  explicit NodePtr(NodeBase* data)
      : data_(data) {
    if (data != nullptr) {
      data_->IncRef();
    }
  }
  // friend declaration
  friend class Node;
  template<typename>
  friend class NodePtr;
  template<typename Y, typename... Args>
  friend NodePtr<Y> make_node(Args&&...);
};
}  // namespace tvm

#endif  // TVM_RUNTIME_NODE_BASE_H_