mxnet
memory.h
Go to the documentation of this file.
1 
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 
23 template<size_t size, size_t align>
24 class MemoryPool {
25  public:
28  static_assert(align % alignof(LinkedList) == 0,
29  "alignment requirement failed.");
30  curr_page_.reset(new Page());
31  }
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  }
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 
86 template<typename T>
88  public:
90  typedef T* pointer;
92  typedef const T* const_ptr;
94  typedef T value_type;
102  template<typename U>
109  inline T* allocate(size_t n) {
110  CHECK_EQ(n, 1);
112  return static_cast<T*>(Store::Get()->allocate());
113  }
119  inline void deallocate(T* p, size_t n) {
120  CHECK_EQ(n, 1);
122  Store::Get()->deallocate(p);
123  }
124 };
125 
126 
133 template<typename T>
135  public:
137  ThreadlocalSharedPtr() : block_(nullptr) {}
142  ThreadlocalSharedPtr(std::nullptr_t other) : block_(nullptr) {} // NOLINT(*)
148  : block_(other.block_) {
149  IncRef(block_);
150  }
156  : block_(other.block_) {
157  other.block_ = nullptr;
158  }
163  DecRef(block_);
164  }
171  DecRef(block_);
172  block_ = other.block_;
173  other.block_ = nullptr;
174  return *this;
175  }
182  DecRef(block_);
183  block_ = other.block_;
184  IncRef(block_);
185  return *this;
186  }
188  inline bool operator==(std::nullptr_t other) const {
189  return block_ == nullptr;
190  }
194  inline T* get() const {
195  if (block_ == nullptr) return nullptr;
196  return reinterpret_cast<T*>(&(block_->data));
197  }
201  inline void reset() {
202  DecRef(block_);
203  block_ = nullptr;
204  }
206  inline bool unique() const {
207  if (block_ == nullptr) return false;
208  return block_->use_count_ == 1;
209  }
211  inline T* operator*() const {
212  return reinterpret_cast<T*>(&(block_->data));
213  }
215  inline T* operator->() const {
216  return reinterpret_cast<T*>(&(block_->data));
217  }
224  template <typename... Args>
225  inline static ThreadlocalSharedPtr<T> Create(Args&&... args) {
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
241  inline static void DecRef(RefBlock* block) {
242  if (block != nullptr) {
243  if (--block->use_count_ == 0) {
245  T* dptr = reinterpret_cast<T*>(&(block->data));
246  dptr->~T();
247  arena.deallocate(block, 1);
248  }
249  }
250  }
251  // increase ref counter
252  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_
void deallocate(void *p)
deallocate a piece of memory
Definition: memory.h:53
static ThreadlocalSharedPtr< T > Create(Args &&...args)
create a new space from threadlocal storage and return it.
Definition: memory.h:225
A memory pool that allocate memory of fixed size and alignment.
Definition: memory.h:24
ThreadlocalSharedPtr()
default constructor
Definition: memory.h:137
T * operator*() const
Definition: memory.h:211
bool operator==(std::nullptr_t other) const
check if nullptr
Definition: memory.h:188
T * pointer
pointer type
Definition: memory.h:90
ThreadlocalAllocator()
default constructor
Definition: memory.h:96
void * allocate()
allocate a new memory of size
Definition: memory.h:33
ThreadlocalSharedPtr(ThreadlocalSharedPtr< T > &&other)
move constructor
Definition: memory.h:155
A threadlocal store to store threadlocal variables. Will return a thread local singleton of type T...
Definition: thread_local.h:35
T * operator->() const
Definition: memory.h:215
MemoryPool()
constructor
Definition: memory.h:27
ThreadlocalSharedPtr(std::nullptr_t other)
constructor from nullptr
Definition: memory.h:142
ThreadlocalAllocator(const ThreadlocalAllocator< U > &other)
constructor from another allocator
Definition: memory.h:103
namespace for dmlc
Definition: array_view.h:12
Portable thread local storage.
a shared pointer like type that allocate object from a threadlocal object pool. This object is not th...
Definition: memory.h:134
void reset()
reset the pointer to nullptr.
Definition: memory.h:201
A thread local allocator that get memory from a threadlocal memory pool. This is suitable to allocate...
Definition: memory.h:87
ThreadlocalSharedPtr< T > & operator=(const ThreadlocalSharedPtr< T > &other)
copy assignment
Definition: memory.h:181
bool unique() const
Definition: memory.h:206
void deallocate(T *p, size_t n)
deallocate memory
Definition: memory.h:119
~ThreadlocalSharedPtr()
destructor
Definition: memory.h:162
const T * const_ptr
const pointer type
Definition: memory.h:92
ThreadlocalSharedPtr< T > & operator=(ThreadlocalSharedPtr< T > &&other)
move assignment
Definition: memory.h:170
T * allocate(size_t n)
allocate memory
Definition: memory.h:109
ThreadlocalSharedPtr(const ThreadlocalSharedPtr< T > &other)
copy constructor
Definition: memory.h:147
T value_type
value type
Definition: memory.h:94