blob: c5da093a70b881d45ff07ea33d6f7f37033ac572 [file] [log] [blame]
/*!
* Copyright 2018 by Contributors
*
* \file arena.h
* \brief Arena allocator that allocates
* memory chunks and frees them all during destruction time.
*/
#ifndef TVM_COMMON_ARENA_H_
#define TVM_COMMON_ARENA_H_
#include <type_traits>
namespace tvm {
namespace common {
const constexpr int kArenaPageSize = 16 << 10;
/*!
* \brief Arena allocator that allocates memory from continuous
* chunk and frees them all only during destruction.
*/
class Arena {
public:
Arena() {
// eagerly allocate the first page.
head_ = reinterpret_cast<PageHeader*>(new Page());
head_->next = nullptr;
head_->ptr = sizeof(PageHeader);
}
~Arena() {
// delete all the allocated pages.
while (head_ != nullptr) {
Page* page = reinterpret_cast<Page*>(head_);
head_ = head_->next;
delete page;
}
}
/*!
* \brief Allocate a space from Arena for type T
* \param T the data type to be allocated
* \note The space of T is not initialized.
*/
template<typename T>
T* allocate_() {
return static_cast<T*>(Alloc(sizeof(T), alignof(T)));
}
/*!
* \brief Create a new instance of type T.
* \param args The constructor argument.
* \tparam T the type to be created.
* \tparam Args Arguments to the constructor.
*
* \return The allocated object.
* \note The type T must be simple type, or only contain
* memory allocated from the same arena.
* Otherwise the destructor needs to be called explicitly.
*/
template<typename T, typename... Args>
T* make(Args&&... args) {
T* ptr = allocate_<T>();
new (ptr) T(std::forward<Args>(args)...);
return ptr;
}
private:
// page size 16 KB
// The page data type;
using Page = std::aligned_storage<kArenaPageSize, 1024>::type;
/*! \brief Page header */
struct PageHeader {
/*! \brief points to the next page */
PageHeader* next;
/*! \brief memory allocator ptr inside page */
size_t ptr;
};
/* \brief The page header */
PageHeader* head_{nullptr};
/*!
* \brief Align ptr by upper bound.
* \param ptr The pointer value.
* \param align The alignment requirement.
*/
size_t UpperAlign(size_t ptr, size_t align) {
return ptr + (align - (ptr % align)) % align;
}
/*!
* \brief Internal aligned alloc function.
* \param size The size of the memory.
* \param align The alignment requirement.
*/
void* Alloc(size_t size, size_t align) {
size_t ptr = UpperAlign(head_->ptr, align);
if (ptr + size <= kArenaPageSize) {
head_->ptr = ptr + size;
return reinterpret_cast<char*>(head_) + ptr;
} else {
PageHeader* new_head = reinterpret_cast<PageHeader*>(new Page());
new_head->next = head_;
ptr = UpperAlign(sizeof(PageHeader), align);
CHECK_LE(ptr + size, kArenaPageSize);
new_head->ptr = ptr + size;
head_ = new_head;
return reinterpret_cast<char*>(head_) + ptr;
}
}
};
/*!
* \brief Link list node
* \tparam T the content data type
*/
template<typename T>
struct LinkNode {
/*! \brief The content value */
T value;
/*! \brief pointer to the next location */
LinkNode<T>* next{nullptr};
};
/*!
* \brief LinkedList structure
* \tparam T the content data type
* \note This is a simple data structure that can be used together with the arena.
* \sa LinkNode
*/
template<typename T>
struct LinkedList {
/*! \brief Head pointer */
LinkNode<T>* head{nullptr};
/*! \brief Tail pointer */
LinkNode<T>* tail{nullptr};
/*!
* \brief Push a new node to the end of the linked list.
* \param node The node to be pushed.
*/
void Push(LinkNode<T>* node) {
node->next = nullptr;
if (this->tail != nullptr) {
this->tail->next = node;
this->tail = node;
} else {
head = tail = node;
}
}
};
} // namespace common
} // namespace tvm
#endif // TVM_COMMON_ARENA_H_