object.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.
 */
// Acknowledgement: This file originates from incubator-tvm
#ifndef MXNET_RUNTIME_OBJECT_H_
#define MXNET_RUNTIME_OBJECT_H_
 
#include <dmlc/logging.h>
#include <type_traits>
#include <string>
#include <utility>
#include "c_runtime_api.h"
 
#ifndef MXNET_OBJECT_ATOMIC_REF_COUNTER
#define MXNET_OBJECT_ATOMIC_REF_COUNTER 1
#endif
 
#if MXNET_OBJECT_ATOMIC_REF_COUNTER
#include <atomic>
#endif  // MXNET_OBJECT_ATOMIC_REF_COUNTER
 
namespace mxnet {
namespace runtime {
 
enum TypeIndex {
  kRoot         = 0,
  kMXNetTensor  = 1,
  kMXNetClosure = 2,
  kMXNetADT     = 3,
  kMXNetMap     = 4,
  kMXNetString  = 5,
  kEllipsis     = 6,
  kSlice        = 7,
  kInteger      = 8,
  kFloat        = 9,
  kStaticIndexEnd,
  kDynamic = kStaticIndexEnd
};
 
class Object {
 public:
  typedef void (*FDeleter)(Object* self);
  uint32_t type_index() const {
    return type_index_;
  }
  std::string GetTypeKey() const {
    return TypeIndex2Key(type_index_);
  }
  size_t GetTypeKeyHash() const {
    return TypeIndex2KeyHash(type_index_);
  }
  template <typename TargetType>
  inline bool IsInstance() const;
 
  MXNET_DLL static std::string TypeIndex2Key(uint32_t tindex);
  MXNET_DLL static size_t TypeIndex2KeyHash(uint32_t tindex);
  MXNET_DLL static uint32_t TypeKey2Index(const std::string& key);
 
#if MXNET_OBJECT_ATOMIC_REF_COUNTER
  using RefCounterType = std::atomic<int32_t>;
#else
  using RefCounterType = int32_t;
#endif
 
  static constexpr const char* _type_key = "Object";
 
  static uint32_t _GetOrAllocRuntimeTypeIndex() {
    return TypeIndex::kRoot;
  }
  static uint32_t RuntimeTypeIndex() {
    return TypeIndex::kRoot;
  }
 
  // Default object type properties for sub-classes
  static constexpr bool _type_final                    = false;
  static constexpr uint32_t _type_child_slots          = 0;
  static constexpr bool _type_child_slots_can_overflow = true;
  // NOTE: the following field is not type index of Object
  // but was intended to be used by sub-classes as default value.
  // The type index of Object is TypeIndex::kRoot
  static constexpr uint32_t _type_index = TypeIndex::kDynamic;
 
  // Default constructor and copy constructor
  Object() {}
  // 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.
  // This will enable us to use make_object<ObjectClass>(*obj_ptr)
  // to copy an existing object.
  Object(const Object& other) {  // NOLINT(*)
  }
  Object(Object&& other) {  // NOLINT(*)
  }
  Object& operator=(const Object& other) {  // NOLINT(*)
    return *this;
  }
  Object& operator=(Object&& other) {  // NOLINT(*)
    return *this;
  }
 
 protected:
  // The fields of the base object cell.
  uint32_t type_index_{0};
  RefCounterType ref_counter_{0};
  FDeleter deleter_ = nullptr;
  // Invariant checks.
  static_assert(sizeof(int32_t) == sizeof(RefCounterType) &&
                    alignof(int32_t) == sizeof(RefCounterType),
                "RefCounter ABI check.");
 
  MXNET_DLL static uint32_t GetOrAllocRuntimeTypeIndex(const std::string& key,
                                                       uint32_t static_tindex,
                                                       uint32_t parent_tindex,
                                                       uint32_t type_child_slots,
                                                       bool type_child_slots_can_overflow);
 
  // reference counter related operations
  inline void IncRef();
  inline void DecRef();
 
 private:
  inline int use_count() const;
  MXNET_DLL bool DerivedFrom(uint32_t parent_tindex) const;
  // friend classes
  template <typename>
  friend class ObjAllocatorBase;
  template <typename>
  friend class ObjectPtr;
  friend class MXNetRetValue;
  friend class ObjectInternal;
};
 
template <typename RefType, typename ObjectType>
inline RefType GetRef(const ObjectType* ptr);
 
template <typename SubRef, typename BaseRef>
inline SubRef Downcast(BaseRef ref);
 
template <typename T>
class ObjectPtr {
 public:
  ObjectPtr() {}
  ObjectPtr(std::nullptr_t) {}  // NOLINT(*)
  ObjectPtr(const ObjectPtr<T>& other)  // NOLINT(*)
      : ObjectPtr(other.data_) {}
  template <typename U>
  ObjectPtr(const ObjectPtr<U>& other)  // NOLINT(*)
      : ObjectPtr(other.data_) {
    static_assert(std::is_base_of<T, U>::value,
                  "can only assign of child class ObjectPtr to parent");
  }
  ObjectPtr(ObjectPtr<T>&& other)  // NOLINT(*)
      : data_(other.data_) {
    other.data_ = nullptr;
  }
  template <typename Y>
  ObjectPtr(ObjectPtr<Y>&& other)  // NOLINT(*)
      : data_(other.data_) {
    static_assert(std::is_base_of<T, Y>::value,
                  "can only assign of child class ObjectPtr to parent");
    other.data_ = nullptr;
  }
  ~ObjectPtr() {
    this->reset();
  }
  void swap(ObjectPtr<T>& other) {  // NOLINT(*)
    std::swap(data_, other.data_);
  }
  T* get() const {
    return static_cast<T*>(data_);
  }
  T* operator->() const {
    return get();
  }
  T& operator*() const {  // NOLINT(*)
    return *get();
  }
  ObjectPtr<T>& operator=(const ObjectPtr<T>& other) {  // NOLINT(*)
    // takes in plane operator to enable copy elison.
    // copy-and-swap idiom
    ObjectPtr(other).swap(*this);  // NOLINT(*)
    return *this;
  }
  ObjectPtr<T>& operator=(ObjectPtr<T>&& other) {  // NOLINT(*)
    // copy-and-swap idiom
    ObjectPtr(std::move(other)).swap(*this);  // NOLINT(*)
    return *this;
  }
  void reset() {
    if (data_ != nullptr) {
      data_->DecRef();
      data_ = nullptr;
    }
  }
  int use_count() const {
    return data_ != nullptr ? data_->use_count() : 0;
  }
  bool unique() const {
    return data_ != nullptr && data_->use_count() == 1;
  }
  bool operator==(const ObjectPtr<T>& other) const {
    return data_ == other.data_;
  }
  bool operator!=(const ObjectPtr<T>& other) const {
    return data_ != other.data_;
  }
  bool operator==(std::nullptr_t null) const {
    return data_ == nullptr;
  }
  bool operator!=(std::nullptr_t null) const {
    return data_ != nullptr;
  }
 
 private:
  Object* data_{nullptr};
  explicit ObjectPtr(Object* data) : data_(data) {
    if (data != nullptr) {
      data_->IncRef();
    }
  }
  // friend classes
  friend class Object;
  friend class ObjectRef;
  friend struct ObjectHash;
  template <typename>
  friend class ObjectPtr;
  template <typename>
  friend class ObjAllocatorBase;
  friend class MXNetPODValue_;
  friend class MXNetArgsSetter;
  friend class MXNetRetValue;
  friend class MXNetArgValue;
  template <typename RefType, typename ObjType>
  friend RefType GetRef(const ObjType* ptr);
  template <typename BaseType, typename ObjType>
  friend ObjectPtr<BaseType> GetObjectPtr(ObjType* ptr);
};
 
class ObjectRef {
 public:
  ObjectRef() = default;
  explicit ObjectRef(ObjectPtr<Object> data) : data_(data) {}
  bool same_as(const ObjectRef& other) const {
    return data_ == other.data_;
  }
  bool operator==(const ObjectRef& other) const {
    return data_ == other.data_;
  }
  bool operator!=(const ObjectRef& other) const {
    return data_ != other.data_;
  }
  bool operator<(const ObjectRef& other) const {
    return data_.get() < other.data_.get();
  }
  bool defined() const {
    return data_ != nullptr;
  }
  const Object* get() const {
    return data_.get();
  }
  const Object* operator->() const {
    return get();
  }
  bool unique() const {
    return data_.unique();
  }
  template <typename ObjectType>
  inline const ObjectType* as() const;
 
  using ContainerType = Object;
  // Default type properties for the reference class.
  static constexpr bool _type_is_nullable = true;
 
 protected:
  ObjectPtr<Object> data_;
  Object* get_mutable() const {
    return data_.get();
  }
  template <typename T>
  static T DowncastNoCheck(ObjectRef ref) {
    return T(std::move(ref.data_));
  }
  template <typename ObjectType>
  static ObjectPtr<ObjectType> GetDataPtr(const ObjectRef& ref) {
    return ObjectPtr<ObjectType>(ref.data_.data_);
  }
  // friend classes.
  friend struct ObjectHash;
  friend class MXNetRetValue;
  friend class MXNetArgsSetter;
  template <typename SubRef, typename BaseRef>
  friend SubRef Downcast(BaseRef ref);
};
 
template <typename BaseType, typename ObjectType>
inline ObjectPtr<BaseType> GetObjectPtr(ObjectType* ptr);
 
struct ObjectHash {
  size_t operator()(const ObjectRef& a) const {
    return operator()(a.data_);
  }
 
  template <typename T>
  size_t operator()(const ObjectPtr<T>& a) const {
    return std::hash<Object*>()(a.get());
  }
};
 
struct ObjectEqual {
  bool operator()(const ObjectRef& a, const ObjectRef& b) const {
    return a.same_as(b);
  }
 
  template <typename T>
  size_t operator()(const ObjectPtr<T>& a, const ObjectPtr<T>& b) const {
    return a == b;
  }
};
 
#define MXNET_DECLARE_BASE_OBJECT_INFO(TypeName, ParentType)                                     \
  static uint32_t RuntimeTypeIndex() {                                                           \
    return TypeName::_type_index != ::mxnet::runtime::TypeIndex::kDynamic ?                      \
               TypeName::_type_index :                                                           \
               _GetOrAllocRuntimeTypeIndex();                                                    \
  }                                                                                              \
  static uint32_t _GetOrAllocRuntimeTypeIndex() {                                                \
    static uint32_t tidx = GetOrAllocRuntimeTypeIndex(TypeName::_type_key,                       \
                                                      TypeName::_type_index,                     \
                                                      ParentType::_GetOrAllocRuntimeTypeIndex(), \
                                                      TypeName::_type_child_slots,               \
                                                      TypeName::_type_child_slots_can_overflow); \
    return tidx;                                                                                 \
  }
 
#define MXNET_DECLARE_FINAL_OBJECT_INFO(TypeName, ParentType) \
  static const constexpr bool _type_final      = true;        \
  static const constexpr int _type_child_slots = 0;           \
  MXNET_DECLARE_BASE_OBJECT_INFO(TypeName, ParentType)
 
#define MXNET_REGISTER_OBJECT_TYPE(TypeName)                                  \
  static DMLC_ATTRIBUTE_UNUSED uint32_t __make_Object_tidx##_##TypeName##__ = \
      TypeName::_GetOrAllocRuntimeTypeIndex()
 
#define MXNET_DEFINE_DEFAULT_COPY_MOVE_AND_ASSIGN(TypeName) \
  TypeName(const TypeName& other) = default;                \
  TypeName(TypeName&& other)      = default;                \
  TypeName& operator=(const TypeName& other) = default;     \
  TypeName& operator=(TypeName&& other) = default;
 
#define MXNET_DEFINE_OBJECT_REF_METHODS(TypeName, ParentType, ObjectName)                       \
  TypeName() {}                                                                                 \
  explicit TypeName(::mxnet::runtime::ObjectPtr<::mxnet::runtime::Object> n) : ParentType(n) {} \
  const ObjectName* operator->() const {                                                        \
    return static_cast<const ObjectName*>(data_.get());                                         \
  }                                                                                             \
  operator bool() const {                                                                       \
    return data_ != nullptr;                                                                    \
  }                                                                                             \
  using ContainerType = ObjectName;
 
#define MXNET_DEFINE_OBJECT_REF_METHODS_MUT(TypeName, ParentType, ObjectName)                   \
  TypeName() {}                                                                                 \
  explicit TypeName(::mxnet::runtime::ObjectPtr<::mxnet::runtime::Object> n) : ParentType(n) {} \
  ObjectName* operator->() {                                                                    \
    return static_cast<ObjectName*>(data_.get());                                               \
  }                                                                                             \
  operator bool() const {                                                                       \
    return data_ != nullptr;                                                                    \
  }                                                                                             \
  using ContainerType = ObjectName;
 
#define MXNET_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(TypeName, ParentType, ObjectName)           \
  explicit TypeName(::mxnet::runtime::ObjectPtr<::mxnet::runtime::Object> n) : ParentType(n) {} \
  MXNET_DEFINE_DEFAULT_COPY_MOVE_AND_ASSIGN(TypeName);                                          \
  const ObjectName* operator->() const {                                                        \
    return static_cast<const ObjectName*>(data_.get());                                         \
  }                                                                                             \
  const ObjectName* get() const {                                                               \
    return operator->();                                                                        \
  }                                                                                             \
  static constexpr bool _type_is_nullable = false;                                              \
  using ContainerType                     = ObjectName;
 
// Implementations details below
// Object reference counting.
#if MXNET_OBJECT_ATOMIC_REF_COUNTER
 
inline void Object::IncRef() {
  ref_counter_.fetch_add(1, std::memory_order_relaxed);
}
 
inline void Object::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);
    }
  }
}
 
inline int Object::use_count() const {
  return ref_counter_.load(std::memory_order_relaxed);
}
 
#else
 
inline void Object::IncRef() {
  ++ref_counter_;
}
 
inline void Object::DecRef() {
  if (--ref_counter == 0) {
    if (this->deleter_ != nullptr) {
      (*this->deleter_)(this);
    }
  }
}
 
inline int Object::use_count() const {
  return ref_counter_;
}
 
#endif  // MXNET_OBJECT_ATOMIC_REF_COUNTER
 
template <typename TargetType>
inline bool Object::IsInstance() const {
  const Object* self = this;
  // NOTE: the following code can be optimized by
  // compiler dead-code elimination for already known constants.
  if (self != nullptr) {
    // Everything is a subclass of object.
    if (std::is_same<TargetType, Object>::value)
      return true;
    if (TargetType::_type_final) {
      // if the target type is a final type
      // then we only need to check the equivalence.
      return self->type_index_ == TargetType::RuntimeTypeIndex();
    } else {
      // if target type is a non-leaf type
      // Check if type index falls into the range of reserved slots.
      uint32_t begin = TargetType::RuntimeTypeIndex();
      // The condition will be optimized by constant-folding.
      if (TargetType::_type_child_slots != 0) {
        uint32_t end = begin + TargetType::_type_child_slots;
        if (self->type_index_ >= begin && self->type_index_ < end)
          return true;
      } else {
        if (self->type_index_ == begin)
          return true;
      }
      if (!TargetType::_type_child_slots_can_overflow)
        return false;
      // Invariance: parent index is always smaller than the child.
      if (self->type_index_ < TargetType::RuntimeTypeIndex())
        return false;
      // The rare slower-path, check type hierachy.
      return self->DerivedFrom(TargetType::RuntimeTypeIndex());
    }
  } else {
    return false;
  }
}
 
template <typename ObjectType>
inline const ObjectType* ObjectRef::as() const {
  if (data_ != nullptr && data_->IsInstance<ObjectType>()) {
    return static_cast<ObjectType*>(data_.get());
  } else {
    return nullptr;
  }
}
 
template <typename RefType, typename ObjType>
inline RefType GetRef(const ObjType* ptr) {
  static_assert(std::is_base_of<typename RefType::ContainerType, ObjType>::value,
                "Can only cast to the ref of same container type");
  if (!RefType::_type_is_nullable) {
    CHECK(ptr != nullptr);
  }
  return RefType(ObjectPtr<Object>(const_cast<Object*>(static_cast<const Object*>(ptr))));
}
 
template <typename BaseType, typename ObjType>
inline ObjectPtr<BaseType> GetObjectPtr(ObjType* ptr) {
  static_assert(std::is_base_of<BaseType, ObjType>::value,
                "Can only cast to the ref of same container type");
  return ObjectPtr<BaseType>(static_cast<Object*>(ptr));
}
 
template <typename SubRef, typename BaseRef>
inline SubRef Downcast(BaseRef ref) {
  if (ref.defined()) {
    CHECK(ref->template IsInstance<typename SubRef::ContainerType>())
        << "Downcast from " << ref->GetTypeKey() << " to " << SubRef::ContainerType::_type_key
        << " failed.";
  } else {
    CHECK(SubRef::_type_is_nullable) << "Downcast from nullptr to not nullable reference of "
                                     << SubRef::ContainerType::_type_key;
  }
  return SubRef(std::move(ref.data_));
}
 
}  // namespace runtime
 
template <typename T>
using NodePtr = runtime::ObjectPtr<T>;
 
}  // namespace mxnet
 
#endif  // MXNET_RUNTIME_OBJECT_H_