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 | |
39 | namespace tvm { |
40 | namespace support { |
41 | |
42 | namespace { |
43 | template <typename T> // For lvalues (T is T&), |
44 | T&& 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 | */ |
52 | struct { |
53 | /*! \brief points to the next page. */ |
54 | ArenaPageHeader* ; |
55 | /*! |
56 | * \brief Total size of the page. |
57 | */ |
58 | size_t ; |
59 | /*! \brief memory allocator offset inside page. */ |
60 | size_t ; |
61 | }; |
62 | |
63 | /*! |
64 | * \brief Arena allocator that allocates memory from continuous |
65 | * chunk and frees them all only during destruction. |
66 | */ |
67 | template <typename PageAllocator> |
68 | class 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 (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 | |