memory.h

Go to the documentation of this file.

 
#ifndef DMLC_MEMORY_H_
#define DMLC_MEMORY_H_
 
#include <vector>
#include <memory>
#include <utility>
#include "./base.h"
#include "./logging.h"
#include "./thread_local.h"
 
namespace dmlc {
 
template<size_t size, size_t align>
class MemoryPool {
 public:
  MemoryPool() {
    static_assert(align % alignof(LinkedList) == 0,
                  "alignment requirement failed.");
    curr_page_.reset(new Page());
  }
  inline void* allocate() {
    if (head_ != nullptr) {
      LinkedList* ret = head_;
      head_ = head_->next;
      return ret;
    } else {
      if (page_ptr_ < kPageSize) {
        return &(curr_page_->data[page_ptr_++]);
      } else {
        allocated_.push_back(std::move(curr_page_));
        curr_page_.reset(new Page());
        page_ptr_ = 1;
        return &(curr_page_->data[0]);
      }
    }
  }
  inline void deallocate(void* p) {
    LinkedList* ptr = static_cast<LinkedList*>(p);
    ptr->next = head_;
    head_ = ptr;
  }
 
 private:
  // page size of each member
  static const int kPageSize = ((1 << 22) / size);
  // page to be requested.
  struct Page {
    typename std::aligned_storage<size, align>::type data[kPageSize];
  };
  // internal linked list structure.
  struct LinkedList {
    LinkedList* next{nullptr};
  };
  // head of free list
  LinkedList* head_{nullptr};
  // current free page
  std::unique_ptr<Page> curr_page_;
  // pointer to the current free page position.
  size_t page_ptr_{0};
  // allocated pages.
  std::vector<std::unique_ptr<Page> > allocated_;
};
 
 
template<typename T>
class ThreadlocalAllocator {
 public:
  typedef T* pointer;
  typedef const T* const_ptr;
  typedef T value_type;
  ThreadlocalAllocator() {}
  template<typename U>
  ThreadlocalAllocator(const ThreadlocalAllocator<U>& other) {}
  inline T* allocate(size_t n) {
    CHECK_EQ(n, 1);
    typedef ThreadLocalStore<MemoryPool<sizeof(T), alignof(T)> > Store;
    return static_cast<T*>(Store::Get()->allocate());
  }
  inline void deallocate(T* p, size_t n) {
    CHECK_EQ(n, 1);
    typedef ThreadLocalStore<MemoryPool<sizeof(T), alignof(T)> > Store;
    Store::Get()->deallocate(p);
  }
};
 
 
template<typename T>
struct ThreadlocalSharedPtr {
 public:
  ThreadlocalSharedPtr() : block_(nullptr) {}
  ThreadlocalSharedPtr(std::nullptr_t other) : block_(nullptr) {}  // NOLINT(*)
  ThreadlocalSharedPtr(const ThreadlocalSharedPtr<T>& other)
      : block_(other.block_) {
    IncRef(block_);
  }
  ThreadlocalSharedPtr(ThreadlocalSharedPtr<T>&& other)
      : block_(other.block_) {
    other.block_ = nullptr;
  }
  ~ThreadlocalSharedPtr() {
    DecRef(block_);
  }
  inline ThreadlocalSharedPtr<T>& operator=(ThreadlocalSharedPtr<T>&& other) {
    DecRef(block_);
    block_ = other.block_;
    other.block_ = nullptr;
    return *this;
  }
  inline ThreadlocalSharedPtr<T> &operator=(const ThreadlocalSharedPtr<T>& other) {
    DecRef(block_);
    block_ = other.block_;
    IncRef(block_);
    return *this;
  }
  inline bool operator==(std::nullptr_t other) const {
    return block_ == nullptr;
  }
  inline T* get() const {
    if (block_ == nullptr) return nullptr;
    return reinterpret_cast<T*>(&(block_->data));
  }
  inline void reset() {
    DecRef(block_);
    block_ = nullptr;
  }
  inline bool unique() const {
    if (block_ == nullptr) return false;
    return block_->use_count_ == 1;
  }
  inline T* operator*() const {
    return reinterpret_cast<T*>(&(block_->data));
  }
  inline T* operator->() const {
    return reinterpret_cast<T*>(&(block_->data));
  }
  template <typename... Args>
  inline static ThreadlocalSharedPtr<T> Create(Args&&... args) {
    ThreadlocalAllocator<RefBlock> arena;
    ThreadlocalSharedPtr<T> p;
    p.block_ = arena.allocate(1);
    p.block_->use_count_ = 1;
    new (&(p.block_->data)) T(std::forward<Args>(args)...);
    return p;
  }
 
 private:
  // internal reference block
  struct RefBlock {
    typename std::aligned_storage<sizeof(T), alignof(T)>::type data;
    unsigned use_count_;
  };
  // decrease ref counter
  inline static void DecRef(RefBlock* block) {
    if (block != nullptr) {
      if (--block->use_count_ == 0) {
        ThreadlocalAllocator<RefBlock> arena;
        T* dptr = reinterpret_cast<T*>(&(block->data));
        dptr->~T();
        arena.deallocate(block, 1);
      }
    }
  }
  // increase ref counter
  inline static void IncRef(RefBlock* block) {
    if (block != nullptr) {
      ++block->use_count_;
    }
  }
  // internal block
  RefBlock *block_;
};
 
}  // namespace dmlc
 
#endif  // DMLC_MEMORY_H_