1/**
2 * Copyright (c) Glow Contributors. See CONTRIBUTORS file.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16#ifndef GLOW_CODEGEN_MEMORYALLOCATOR_H
17#define GLOW_CODEGEN_MEMORYALLOCATOR_H
18#include <cstddef>
19#include <cstdint>
20#include <list>
21#include <set>
22#include <string>
23#include <unordered_map>
24#include <vector>
25
26#include "glow/Support/Memory.h"
27
28namespace glow {
29
30/// Utility function to verify if two contiguous integer intervals overlap. Each
31/// interval is described by a half-open interval [begin, end).
32template <class T>
33inline bool intervalsOverlap(T begin1, T end1, T begin2, T end2) {
34 return (std::max(begin1, begin2) < std::min(end1, end2));
35}
36
37/// Allocation structure which represents a request to allocate or free
38/// a memory region (buffer) within a given MemoryAllocator instance.
39struct Allocation {
40
41 /// Type that should be used as a handle.
42 using Handle = const void *;
43
44 /// Allocation handle which uniquely identifies the buffer to be allocated.
45 /// The provided handle must be unique for each buffer.
46 Handle handle;
47
48 /// Allocation buffer size in bytes. This field is mandatory for ALLOC
49 /// and is ignored for FREE.
50 uint64_t size;
51
52 /// Allocation request type flag: true for ALLOC, false for FREE.
53 bool alloc;
54
55 Allocation(Handle handle, bool alloc, uint64_t size)
56 : handle(handle), size(size), alloc(alloc) {}
57
58 Allocation(size_t id, bool alloc, uint64_t size)
59 : handle((Handle)(id)), size(size), alloc(alloc) {}
60};
61
62/// A POD struct that represents a single half-open allocation [start .. end).
63struct Segment {
64
65 /// The allocation starts at this address.
66 uint64_t begin;
67
68 /// The allocation ends before this address (half-open interval).
69 uint64_t end;
70
71 Segment(uint64_t begin, uint64_t end) : begin(begin), end(end) {}
72
73 /// \returns the size of the interval.
74 uint64_t size() const { return end - begin; }
75
76 /// \returns True if the value \p idx falls within this segment.
77 bool contains(uint64_t idx) const { return idx >= begin && idx < end; }
78};
79
80/// Allocates segments of memory.
81/// Each allocation is associated with a user-defined handle, typically
82/// representing a client-specific object, e.g. a handle can be a `Value *` and
83/// represent a value whose payload is going to be stored in the allocated
84/// memory block. This simplifies the clients of MemoryAllocator and allows them
85/// to use higher-level client-side objects instead of raw allocated addresses
86/// to refer to the allocated memory blocks.
87class MemoryAllocator {
88public:
89 /// Type that should be used as a handle.
90 using Handle = const void *;
91
92 /// A reserved value to mark invalid allocation.
93 static const uint64_t npos;
94
95 explicit MemoryAllocator(const std::string &name, uint64_t memorySize,
96 size_t alignment = TensorAlignment)
97 : name_(name), memorySize_(memorySize), alignment_{alignment} {}
98
99 void reset() {
100 liveSize_ = 0;
101 maxUsedSize_ = 0;
102 maxLiveSize_ = 0;
103 segments_.clear();
104 handleToSegmentMap_.clear();
105 addrToHandleMap_.clear();
106 }
107
108 /// \returns True if the value \p idx is within the currently allocated range.
109 bool contains(uint64_t idx) const {
110 for (auto &s : segments_) {
111 if (s.contains(idx)) {
112 return true;
113 }
114 }
115 return false;
116 }
117
118 /// Allocate a segment of size \p size and associate a \p handle with it.
119 /// \returns the allocated pointer, or MemoryAllocator::npos, if the
120 /// allocation failed.
121 uint64_t allocate(uint64_t size, Handle handle);
122
123 /// Allocate a segment of size \p size and associate a handle \p Handle with
124 /// it. If the allocation is not possible, the allocator should try to evict
125 /// some entries that are not needed at the moment, but it is not allowed to
126 /// evict any entries from \p mustNotEvict set. All evicted entries are stored
127 /// in the \p evicted set.
128 /// \returns the allocated pointer, or MemoryAllocator::npos, if the
129 /// allocation failed.
130 uint64_t allocate(uint64_t size, Handle handle,
131 const std::set<Handle> &mustNotEvict,
132 std::vector<Handle> &evicted);
133
134 /// Allocate all the segments associated with the allocations \p allocList.
135 /// This method has an improved memory allocation efficiency because all
136 /// the allocations are requested at once and the algorithm can improve the
137 /// allocation efficiency by allocating first the larger segments and so
138 /// avoiding early fragmentation. The given allocations must be consistent,
139 /// each ALLOC request must be associated with a FREE request following it
140 /// with the same handle and all the handles must be unique for each
141 /// ALLOC/FREE pair. This function can be used multiple times for multiple IR
142 /// functions and is intended to be called once for each IR function. The
143 /// parameter \p reuseMemory specifies whether the allocation for the current
144 /// function can reuse or not the memory allocated for the previous functions.
145 /// Upon function return information about the allocated segments can be
146 /// retrieved with \ref getSegment(). \returns the total memory usage or
147 /// MemoryAllocator::npos if the allocation failed.
148 /// NOTE: Do not use the functions allocateAll() and allocate() for the same
149 /// allocator object since these functions are not compatible with each other.
150 uint64_t allocateAll(const std::list<Allocation> &allocList,
151 bool reuseMemory = true);
152
153 /// \returns the handle currently associated with the allocation at \p
154 /// address.
155 Handle getHandle(uint64_t ptr) const;
156
157 /// \returns true if there is a handle currently associated with the
158 /// allocation at \p address.
159 bool hasHandle(uint64_t ptr) const;
160
161 /// \returns the segment currently associated with the \p handle.
162 const Segment &getSegment(Handle handle) const;
163
164 /// \returns the address currently associated with the \p handle.
165 uint64_t getAddress(Handle handle) const;
166
167 /// \returns the size of the allocated block currently associated with the \p
168 /// handle.
169 uint64_t getSize(Handle handle) const;
170
171 /// \returns true if there is an address currently associated with the \p
172 /// handle.
173 bool hasAddress(Handle handle) const;
174
175 /// Frees the allocation associated with \p handle.
176 void deallocate(Handle handle);
177
178 /// \returns the total size (in bytes) of the memory.
179 uint64_t getMemorySize() const { return memorySize_; }
180
181 /// \returns the maximum memory usage (in bytes).
182 uint64_t getMaxMemoryUsage() const { return maxUsedSize_; }
183
184 /// \returns the alignment boundary used to align segments.
185 size_t getAlignment() const { return alignment_; }
186
187 /// \returns the allocation efficiency as a float between 0.0 and 1.0
188 /// which quantifies the efficiency of the allocation algorithm. An
189 /// efficiency value of 1.0 means the best theoretically possible. The
190 /// efficiency is not always 1.0 due to memory fragmentation.
191 float getAllocationEfficiency() const;
192
193 /// \returns the name of the memory region.
194 const std::string &getName() const { return name_; }
195
196private:
197 /// The name of the memory region.
198 std::string name_;
199
200 /// A list of allocated segments.
201 std::list<Segment> segments_;
202
203 /// The total size (in bytes) of the memory region. A value of 0 means
204 /// unlimited size (infinite).
205 uint64_t memorySize_;
206
207 /// The maximum size (in bytes) used for segment allocation (memory usage).
208 uint64_t maxUsedSize_{0};
209
210 /// The maximum size (in bytes) for all simultaneously alive segments during
211 /// allocation. This is a theoretical size and is the best (minimum) memory
212 /// usage we can get with any allocation algorithm since it ignores memory
213 /// fragmentation. Since maxLiveSize_ <= maxUsedSize_ we can use this to
214 /// compute an allocation efficiency as maxLiveSize_ / maxUsedSize_.
215 uint64_t maxLiveSize_{0};
216
217 /// Current size (in bytes) for all the live segments currently allocated.
218 /// This is automatically updated for each call to \ref allocate or
219 /// \ref deallocate.
220 uint64_t liveSize_{0};
221
222 /// The alignment boundary for each segment allocation.
223 size_t alignment_;
224
225 /// Maps allocated addresses to the currently associated handles.
226 std::unordered_map<uint64_t, Handle> addrToHandleMap_;
227
228 /// Maps handles to the allocation information about the memory block
229 /// currently associated with them.
230 std::unordered_map<Handle, Segment> handleToSegmentMap_;
231
232 /// \returns the effective size used for allocating a segment which depends
233 /// on the requested allocation size \p size and the alignment used by this
234 /// allocator instance.
235 uint64_t getEffectiveSize(uint64_t size) const;
236
237 /// Tries to evict some entries that are not needed at the moment to free
238 /// enough memory for the allocation of \p size bytes, but it is not allowed
239 /// to evict any entries from \p mustNotEvict set. All evicted entries are
240 /// stored in the \p evicted set. Uses first-fit approach for finding eviction
241 /// candidates.
242 void evictFirstFit(uint64_t size, const std::set<Handle> &mustNotEvict,
243 std::vector<Handle> &evicted);
244
245 /// Associates a \p handle with an allocated address \p ptr and size \p size.
246 void setHandle(uint64_t ptr, uint64_t size, Handle handle);
247
248 /// Function to verify the list of allocations \p allocList before allocating
249 /// the segments with \ref allocateAll. \returns true if allocations are valid
250 /// and false otherwise.
251 bool verifyAllocations(const std::list<Allocation> &allocList) const;
252
253 /// Function to verify the segments allocated with \ref allocateAll after
254 /// using the allocation list \p allocList. \returns true if segments are
255 /// valid and false otherwise.
256 bool verifySegments(const std::list<Allocation> &allocList) const;
257
258 /// Utility function to map the handles used in the allocations \p allocList
259 /// to consecutive unique IDs between 0 and numSegments - 1. The function
260 /// returns the handle to ID map \p handleToIdMap and the ID to handle
261 /// map as a vector \p idToHandleMap. This mapping makes the algorithm used
262 /// in \ref allocateAll easier/faster.
263 void mapHandlesToIds(const std::list<Allocation> &allocList,
264 std::unordered_map<Handle, size_t> &handleToIdMap,
265 std::vector<Handle> &idToHandleMap);
266};
267
268} // namespace glow
269
270#endif // GLOW_CODEGEN_MEMORYALLOCATOR_H
271