1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20/*!
21 *
22 * \file arena.h
23 * \brief Arena allocator that allocates memory chunks and frees them all during destruction time.
24 *
25 * NOTE: This file is portable to bare-metal embedded devices. Don't use operator new (without
26 * placement parameters) or malloc.
27 */
28#ifndef TVM_SUPPORT_GENERIC_ARENA_H_
29#define TVM_SUPPORT_GENERIC_ARENA_H_
30
31#ifndef TVM_ARENA_HAS_DESTRUCTOR
32#define TVM_ARENA_HAS_DESTRUCTOR 1
33#endif
34
35#include <stddef.h>
36
37#include <utility>
38
39namespace tvm {
40namespace support {
41
42namespace {
43template <typename T> // For lvalues (T is T&),
44T&& forward(T&& param) { // take/return lvalue refs.
45 return static_cast<T&&>(param); // For rvalues (T is T),
46} // take/return rvalue refs.
47} // namespace
48
49/*!
50 * \brief An arena page header.
51 */
52struct ArenaPageHeader {
53 /*! \brief points to the next page. */
54 ArenaPageHeader* next;
55 /*!
56 * \brief Total size of the page.
57 */
58 size_t size;
59 /*! \brief memory allocator offset inside page. */
60 size_t offset;
61};
62
63/*!
64 * \brief Arena allocator that allocates memory from continuous
65 * chunk and frees them all only during destruction.
66 */
67template <typename PageAllocator>
68class GenericArena {
69 public:
70 explicit GenericArena(PageAllocator alloc = PageAllocator()) : alloc_(alloc) {
71 // eagerly allocate the first page.
72 head_ = tail_ = alloc_.allocate(1);
73 head_->next = nullptr;
74 }
75
76#if TVM_ARENA_HAS_DESTRUCTOR
77 ~GenericArena() { this->FreeAll(); }
78#endif
79
80 /*! \brief Free all pages. */
81 void FreeAll() {
82 FreePageList(&head_);
83 FreePageList(&free_list_);
84 }
85 /*! \brief Recycle all the pages in the arena */
86 void RecycleAll() {
87 // put all the current list to the free list.
88 tail_->next = free_list_;
89 // allocate the first in the free list to head
90 free_list_ = head_->next;
91 head_->next = nullptr;
92 // Reset the head.
93 head_->offset = sizeof(ArenaPageHeader);
94 tail_ = head_;
95 }
96 /*!
97 * \brief Allocate a space from Arena for type T
98 * \param T the data type to be allocated
99 * \param count Numberof elements
100 * \note The space of T is not initialized.
101 */
102 template <typename T>
103 T* allocate_(int count = 1) {
104 static_assert(PageAllocator::kPageAlign % alignof(T) == 0, "To large alignment");
105 return static_cast<T*>(Alloc(sizeof(T) * count, alignof(T)));
106 }
107 /*!
108 * \brief Create a new instance of type T.
109 * \param args The constructor argument.
110 * \tparam T the type to be created.
111 * \tparam Args Arguments to the constructor.
112 *
113 * \return The allocated object.
114 * \note The type T must be simple type, or only contain
115 * memory allocated from the same arena.
116 * Otherwise the destructor needs to be called explicitly.
117 */
118 template <typename T, typename... Args>
119 T* make(Args&&... args) {
120 T* ptr = allocate_<T>();
121 new (ptr) T(forward<Args>(args)...);
122 return ptr;
123 }
124
125 private:
126 /*! \brief internal page allocator. */
127 PageAllocator alloc_;
128 /* \brief The head of the allocated list. */
129 ArenaPageHeader* head_{nullptr};
130 /*! \brief The tail of the allocated list. */
131 ArenaPageHeader* tail_{nullptr};
132 /* \brief List of free pages. */
133 ArenaPageHeader* free_list_{nullptr};
134 /*!
135 * \brief Align ptr by upper bound.
136 * \param offset The offset value.
137 * \param align The alignment requirement.
138 */
139 size_t UpperAlign(size_t offset, size_t align) {
140 return offset + (align - (offset % align)) % align;
141 }
142 /*!
143 * \brief Internal aligned alloc function.
144 * \param size The size of the memory.
145 * \param align The alignment requirement.
146 */
147 void* Alloc(size_t size, size_t align) {
148 size_t offset = UpperAlign(head_->offset, align);
149 if (offset + size <= head_->size) {
150 head_->offset = offset + size;
151 return reinterpret_cast<char*>(head_) + offset;
152 } else {
153 ArenaPageHeader* new_head;
154 offset = UpperAlign(sizeof(ArenaPageHeader), align);
155 if (free_list_ != nullptr && offset + size <= free_list_->size) {
156 new_head = free_list_;
157 free_list_ = free_list_->next;
158 } else {
159 new_head = alloc_.allocate(offset + size);
160 }
161 new_head->next = head_;
162 new_head->offset = offset + size;
163 head_ = new_head;
164 return reinterpret_cast<char*>(head_) + offset;
165 }
166 }
167 /*!
168 * \brief Free all the pages in the list.
169 * \param ptr The head ptr.
170 */
171 void FreePageList(ArenaPageHeader** ptr) {
172 // delete all the allocated pages.
173 while (ptr[0] != nullptr) {
174 ArenaPageHeader* temp = ptr[0];
175 ptr[0] = ptr[0]->next;
176 alloc_.deallocate(temp);
177 }
178 }
179};
180
181} // namespace support
182} // namespace tvm
183#endif // TVM_SUPPORT_GENERIC_ARENA_H_
184