1 /*! 2 * Copyright (c) 2015 by Contributors 3 * \file memory.h 4 * \brief Additional memory hanlding utilities. 5 */ 6 #ifndef DMLC_MEMORY_H_ 7 #define DMLC_MEMORY_H_ 8 9 #include <vector> 10 #include <memory> 11 #include <utility> 12 #include "./base.h" 13 #include "./logging.h" 14 #include "./thread_local.h" 15 16 namespace dmlc { 17 18 /*! 19 * \brief A memory pool that allocate memory of fixed size and alignment. 20 * \tparam size The size of each piece. 21 * \tparam align The alignment requirement of the memory. 22 */ 23 template<size_t size, size_t align> 24 class MemoryPool { 25 public: 26 /*! \brief constructor */ MemoryPool()27 MemoryPool() { 28 static_assert(align % alignof(LinkedList) == 0, 29 "alignment requirement failed."); 30 curr_page_.reset(new Page()); 31 } 32 /*! \brief allocate a new memory of size */ allocate()33 inline void* allocate() { 34 if (head_ != nullptr) { 35 LinkedList* ret = head_; 36 head_ = head_->next; 37 return ret; 38 } else { 39 if (page_ptr_ < kPageSize) { 40 return &(curr_page_->data[page_ptr_++]); 41 } else { 42 allocated_.push_back(std::move(curr_page_)); 43 curr_page_.reset(new Page()); 44 page_ptr_ = 1; 45 return &(curr_page_->data[0]); 46 } 47 } 48 } 49 /*! 50 * \brief deallocate a piece of memory 51 * \param p The pointer to the memory to be de-allocated. 52 */ deallocate(void * p)53 inline void deallocate(void* p) { 54 LinkedList* ptr = static_cast<LinkedList*>(p); 55 ptr->next = head_; 56 head_ = ptr; 57 } 58 59 private: 60 // page size of each member 61 static const int kPageSize = ((1 << 22) / size); 62 // page to be requested. 63 struct Page { 64 typename std::aligned_storage<size, align>::type data[kPageSize]; 65 }; 66 // internal linked list structure. 67 struct LinkedList { 68 LinkedList* next{nullptr}; 69 }; 70 // head of free list 71 LinkedList* head_{nullptr}; 72 // current free page 73 std::unique_ptr<Page> curr_page_; 74 // pointer to the current free page position. 75 size_t page_ptr_{0}; 76 // allocated pages. 77 std::vector<std::unique_ptr<Page> > allocated_; 78 }; 79 80 81 /*! 82 * \brief A thread local allocator that get memory from a threadlocal memory pool. 83 * This is suitable to allocate objects that do not cross thread. 84 * \tparam T the type of the data to be allocated. 85 */ 86 template<typename T> 87 class ThreadlocalAllocator { 88 public: 89 /*! \brief pointer type */ 90 typedef T* pointer; 91 /*! \brief const pointer type */ 92 typedef const T* const_ptr; 93 /*! \brief value type */ 94 typedef T value_type; 95 /*! \brief default constructor */ ThreadlocalAllocator()96 ThreadlocalAllocator() {} 97 /*! 98 * \brief constructor from another allocator 99 * \param other another allocator 100 * \tparam U another type 101 */ 102 template<typename U> ThreadlocalAllocator(const ThreadlocalAllocator<U> & other)103 ThreadlocalAllocator(const ThreadlocalAllocator<U>& other) {} 104 /*! 105 * \brief allocate memory 106 * \param n number of blocks 107 * \return an uninitialized memory of type T. 108 */ allocate(size_t n)109 inline T* allocate(size_t n) { 110 CHECK_EQ(n, 1); 111 typedef ThreadLocalStore<MemoryPool<sizeof(T), alignof(T)> > Store; 112 return static_cast<T*>(Store::Get()->allocate()); 113 } 114 /*! 115 * \brief deallocate memory 116 * \param p a memory to be returned. 117 * \param n number of blocks 118 */ deallocate(T * p,size_t n)119 inline void deallocate(T* p, size_t n) { 120 CHECK_EQ(n, 1); 121 typedef ThreadLocalStore<MemoryPool<sizeof(T), alignof(T)> > Store; 122 Store::Get()->deallocate(p); 123 } 124 }; 125 126 127 /*! 128 * \brief a shared pointer like type that allocate object 129 * from a threadlocal object pool. This object is not thread-safe 130 * but can be faster than shared_ptr in certain usecases. 131 * \tparam T the data type. 132 */ 133 template<typename T> 134 struct ThreadlocalSharedPtr { 135 public: 136 /*! \brief default constructor */ ThreadlocalSharedPtrThreadlocalSharedPtr137 ThreadlocalSharedPtr() : block_(nullptr) {} 138 /*! 139 * \brief constructor from nullptr 140 * \param other the nullptr type 141 */ ThreadlocalSharedPtrThreadlocalSharedPtr142 ThreadlocalSharedPtr(std::nullptr_t other) : block_(nullptr) {} // NOLINT(*) 143 /*! 144 * \brief copy constructor 145 * \param other another pointer. 146 */ ThreadlocalSharedPtrThreadlocalSharedPtr147 ThreadlocalSharedPtr(const ThreadlocalSharedPtr<T>& other) 148 : block_(other.block_) { 149 IncRef(block_); 150 } 151 /*! 152 * \brief move constructor 153 * \param other another pointer. 154 */ ThreadlocalSharedPtrThreadlocalSharedPtr155 ThreadlocalSharedPtr(ThreadlocalSharedPtr<T>&& other) 156 : block_(other.block_) { 157 other.block_ = nullptr; 158 } 159 /*! 160 * \brief destructor 161 */ ~ThreadlocalSharedPtrThreadlocalSharedPtr162 ~ThreadlocalSharedPtr() { 163 DecRef(block_); 164 } 165 /*! 166 * \brief move assignment 167 * \param other another object to be assigned. 168 * \return self. 169 */ 170 inline ThreadlocalSharedPtr<T>& operator=(ThreadlocalSharedPtr<T>&& other) { 171 DecRef(block_); 172 block_ = other.block_; 173 other.block_ = nullptr; 174 return *this; 175 } 176 /*! 177 * \brief copy assignment 178 * \param other another object to be assigned. 179 * \return self. 180 */ 181 inline ThreadlocalSharedPtr<T> &operator=(const ThreadlocalSharedPtr<T>& other) { 182 DecRef(block_); 183 block_ = other.block_; 184 IncRef(block_); 185 return *this; 186 } 187 /*! \brief check if nullptr */ 188 inline bool operator==(std::nullptr_t other) const { 189 return block_ == nullptr; 190 } 191 /*! 192 * \return get the pointer content. 193 */ getThreadlocalSharedPtr194 inline T* get() const { 195 if (block_ == nullptr) return nullptr; 196 return reinterpret_cast<T*>(&(block_->data)); 197 } 198 /*! 199 * \brief reset the pointer to nullptr. 200 */ resetThreadlocalSharedPtr201 inline void reset() { 202 DecRef(block_); 203 block_ = nullptr; 204 } 205 /*! \return if use_count == 1*/ uniqueThreadlocalSharedPtr206 inline bool unique() const { 207 if (block_ == nullptr) return false; 208 return block_->use_count_ == 1; 209 } 210 /*! \return dereference pointer */ 211 inline T* operator*() const { 212 return reinterpret_cast<T*>(&(block_->data)); 213 } 214 /*! \return dereference pointer */ 215 inline T* operator->() const { 216 return reinterpret_cast<T*>(&(block_->data)); 217 } 218 /*! 219 * \brief create a new space from threadlocal storage and return it. 220 * \tparam Args the arguments. 221 * \param args The input argument 222 * \return the allocated pointer. 223 */ 224 template <typename... Args> CreateThreadlocalSharedPtr225 inline static ThreadlocalSharedPtr<T> Create(Args&&... args) { 226 ThreadlocalAllocator<RefBlock> arena; 227 ThreadlocalSharedPtr<T> p; 228 p.block_ = arena.allocate(1); 229 p.block_->use_count_ = 1; 230 new (&(p.block_->data)) T(std::forward<Args>(args)...); 231 return p; 232 } 233 234 private: 235 // internal reference block 236 struct RefBlock { 237 typename std::aligned_storage<sizeof(T), alignof(T)>::type data; 238 unsigned use_count_; 239 }; 240 // decrease ref counter DecRefThreadlocalSharedPtr241 inline static void DecRef(RefBlock* block) { 242 if (block != nullptr) { 243 if (--block->use_count_ == 0) { 244 ThreadlocalAllocator<RefBlock> arena; 245 T* dptr = reinterpret_cast<T*>(&(block->data)); 246 dptr->~T(); 247 arena.deallocate(block, 1); 248 } 249 } 250 } 251 // increase ref counter IncRefThreadlocalSharedPtr252 inline static void IncRef(RefBlock* block) { 253 if (block != nullptr) { 254 ++block->use_count_; 255 } 256 } 257 // internal block 258 RefBlock *block_; 259 }; 260 261 } // namespace dmlc 262 263 #endif // DMLC_MEMORY_H_ 264