1// Copyright 2019 The Marl Authors.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#ifndef marl_memory_h
16#define marl_memory_h
17
18#include "debug.h"
19#include "export.h"
20
21#include <stdint.h>
22
23#include <array>
24#include <cstdlib>
25#include <memory>
26#include <mutex>
27#include <utility> // std::forward
28
29namespace marl {
30
31template <typename T>
32struct StlAllocator;
33
34// pageSize() returns the size in bytes of a virtual memory page for the host
35// system.
36MARL_EXPORT
37size_t pageSize();
38
39template <typename T>
40MARL_NO_EXPORT inline T alignUp(T val, T alignment) {
41 return alignment * ((val + alignment - 1) / alignment);
42}
43
44// aligned_storage() is a replacement for std::aligned_storage that isn't busted
45// on older versions of MSVC.
46template <size_t SIZE, size_t ALIGNMENT>
47struct aligned_storage {
48 struct alignas(ALIGNMENT) type {
49 unsigned char data[SIZE];
50 };
51};
52
53///////////////////////////////////////////////////////////////////////////////
54// Allocation
55///////////////////////////////////////////////////////////////////////////////
56
57// Allocation holds the result of a memory allocation from an Allocator.
58struct Allocation {
59 // Intended usage of the allocation. Used for allocation trackers.
60 enum class Usage : uint8_t {
61 Undefined = 0,
62 Stack, // Fiber stack
63 Create, // Allocator::create(), make_unique(), make_shared()
64 Vector, // marl::containers::vector<T>
65 List, // marl::containers::list<T>
66 Stl, // marl::StlAllocator
67 Count, // Not intended to be used as a usage type - used for upper bound.
68 };
69
70 // Request holds all the information required to make an allocation.
71 struct Request {
72 size_t size = 0; // The size of the allocation in bytes.
73 size_t alignment = 0; // The minimum alignment of the allocation.
74 bool useGuards = false; // Whether the allocation is guarded.
75 Usage usage = Usage::Undefined; // Intended usage of the allocation.
76 };
77
78 void* ptr = nullptr; // The pointer to the allocated memory.
79 Request request; // Request used for the allocation.
80};
81
82///////////////////////////////////////////////////////////////////////////////
83// Allocator
84///////////////////////////////////////////////////////////////////////////////
85
86// Allocator is an interface to a memory allocator.
87// Marl provides a default implementation with Allocator::Default.
88class Allocator {
89 public:
90 // The default allocator. Initialized with an implementation that allocates
91 // from the OS. Can be assigned a custom implementation.
92 MARL_EXPORT static Allocator* Default;
93
94 // Deleter is a smart-pointer compatible deleter that can be used to delete
95 // objects created by Allocator::create(). Deleter is used by the smart
96 // pointers returned by make_shared() and make_unique().
97 struct MARL_EXPORT Deleter {
98 MARL_NO_EXPORT inline Deleter();
99 MARL_NO_EXPORT inline Deleter(Allocator* allocator, size_t count);
100
101 template <typename T>
102 MARL_NO_EXPORT inline void operator()(T* object);
103
104 Allocator* allocator = nullptr;
105 size_t count = 0;
106 };
107
108 // unique_ptr<T> is an alias to std::unique_ptr<T, Deleter>.
109 template <typename T>
110 using unique_ptr = std::unique_ptr<T, Deleter>;
111
112 virtual ~Allocator() = default;
113
114 // allocate() allocates memory from the allocator.
115 // The returned Allocation::request field must be equal to the Request
116 // parameter.
117 virtual Allocation allocate(const Allocation::Request&) = 0;
118
119 // free() frees the memory returned by allocate().
120 // The Allocation must have all fields equal to those returned by allocate().
121 virtual void free(const Allocation&) = 0;
122
123 // create() allocates and constructs an object of type T, respecting the
124 // alignment of the type.
125 // The pointer returned by create() must be deleted with destroy().
126 template <typename T, typename... ARGS>
127 inline T* create(ARGS&&... args);
128
129 // destroy() destructs and frees the object allocated with create().
130 template <typename T>
131 inline void destroy(T* object);
132
133 // make_unique() returns a new object allocated from the allocator wrapped
134 // in a unique_ptr that respects the alignment of the type.
135 template <typename T, typename... ARGS>
136 inline unique_ptr<T> make_unique(ARGS&&... args);
137
138 // make_unique_n() returns an array of n new objects allocated from the
139 // allocator wrapped in a unique_ptr that respects the alignment of the
140 // type.
141 template <typename T, typename... ARGS>
142 inline unique_ptr<T> make_unique_n(size_t n, ARGS&&... args);
143
144 // make_shared() returns a new object allocated from the allocator
145 // wrapped in a std::shared_ptr that respects the alignment of the type.
146 template <typename T, typename... ARGS>
147 inline std::shared_ptr<T> make_shared(ARGS&&... args);
148
149 protected:
150 Allocator() = default;
151};
152
153///////////////////////////////////////////////////////////////////////////////
154// Allocator::Deleter
155///////////////////////////////////////////////////////////////////////////////
156Allocator::Deleter::Deleter() : allocator(nullptr) {}
157Allocator::Deleter::Deleter(Allocator* allocator, size_t count)
158 : allocator(allocator), count(count) {}
159
160template <typename T>
161void Allocator::Deleter::operator()(T* object) {
162 object->~T();
163
164 Allocation allocation;
165 allocation.ptr = object;
166 allocation.request.size = sizeof(T) * count;
167 allocation.request.alignment = alignof(T);
168 allocation.request.usage = Allocation::Usage::Create;
169 allocator->free(allocation);
170}
171
172///////////////////////////////////////////////////////////////////////////////
173// Allocator
174///////////////////////////////////////////////////////////////////////////////
175template <typename T, typename... ARGS>
176T* Allocator::create(ARGS&&... args) {
177 Allocation::Request request;
178 request.size = sizeof(T);
179 request.alignment = alignof(T);
180 request.usage = Allocation::Usage::Create;
181
182 auto alloc = allocate(request);
183 new (alloc.ptr) T(std::forward<ARGS>(args)...);
184 return reinterpret_cast<T*>(alloc.ptr);
185}
186
187template <typename T>
188void Allocator::destroy(T* object) {
189 object->~T();
190
191 Allocation alloc;
192 alloc.ptr = object;
193 alloc.request.size = sizeof(T);
194 alloc.request.alignment = alignof(T);
195 alloc.request.usage = Allocation::Usage::Create;
196 free(alloc);
197}
198
199template <typename T, typename... ARGS>
200Allocator::unique_ptr<T> Allocator::make_unique(ARGS&&... args) {
201 return make_unique_n<T>(1, std::forward<ARGS>(args)...);
202}
203
204template <typename T, typename... ARGS>
205Allocator::unique_ptr<T> Allocator::make_unique_n(size_t n, ARGS&&... args) {
206 if (n == 0) {
207 return nullptr;
208 }
209
210 Allocation::Request request;
211 request.size = sizeof(T) * n;
212 request.alignment = alignof(T);
213 request.usage = Allocation::Usage::Create;
214
215 auto alloc = allocate(request);
216 new (alloc.ptr) T(std::forward<ARGS>(args)...);
217 return unique_ptr<T>(reinterpret_cast<T*>(alloc.ptr), Deleter{this, n});
218}
219
220template <typename T, typename... ARGS>
221std::shared_ptr<T> Allocator::make_shared(ARGS&&... args) {
222 Allocation::Request request;
223 request.size = sizeof(T);
224 request.alignment = alignof(T);
225 request.usage = Allocation::Usage::Create;
226
227 auto alloc = allocate(request);
228 new (alloc.ptr) T(std::forward<ARGS>(args)...);
229 return std::shared_ptr<T>(reinterpret_cast<T*>(alloc.ptr), Deleter{this, 1});
230}
231
232///////////////////////////////////////////////////////////////////////////////
233// TrackedAllocator
234///////////////////////////////////////////////////////////////////////////////
235
236// TrackedAllocator wraps an Allocator to track the allocations made.
237class TrackedAllocator : public Allocator {
238 public:
239 struct UsageStats {
240 // Total number of allocations.
241 size_t count = 0;
242 // total allocation size in bytes (as requested, may be higher due to
243 // alignment or guards).
244 size_t bytes = 0;
245 };
246
247 struct Stats {
248 // numAllocations() returns the total number of allocations across all
249 // usages for the allocator.
250 inline size_t numAllocations() const;
251
252 // bytesAllocated() returns the total number of bytes allocated across all
253 // usages for the allocator.
254 inline size_t bytesAllocated() const;
255
256 // Statistics per usage.
257 std::array<UsageStats, size_t(Allocation::Usage::Count)> byUsage;
258 };
259
260 // Constructor that wraps an existing allocator.
261 inline TrackedAllocator(Allocator* allocator);
262
263 // stats() returns the current allocator statistics.
264 inline Stats stats();
265
266 // Allocator compliance
267 inline Allocation allocate(const Allocation::Request&) override;
268 inline void free(const Allocation&) override;
269
270 private:
271 Allocator* const allocator;
272 std::mutex mutex;
273 Stats stats_;
274};
275
276size_t TrackedAllocator::Stats::numAllocations() const {
277 size_t out = 0;
278 for (auto& stats : byUsage) {
279 out += stats.count;
280 }
281 return out;
282}
283
284size_t TrackedAllocator::Stats::bytesAllocated() const {
285 size_t out = 0;
286 for (auto& stats : byUsage) {
287 out += stats.bytes;
288 }
289 return out;
290}
291
292TrackedAllocator::TrackedAllocator(Allocator* allocator)
293 : allocator(allocator) {}
294
295TrackedAllocator::Stats TrackedAllocator::stats() {
296 std::unique_lock<std::mutex> lock(mutex);
297 return stats_;
298}
299
300Allocation TrackedAllocator::allocate(const Allocation::Request& request) {
301 {
302 std::unique_lock<std::mutex> lock(mutex);
303 auto& usageStats = stats_.byUsage[int(request.usage)];
304 ++usageStats.count;
305 usageStats.bytes += request.size;
306 }
307 return allocator->allocate(request);
308}
309
310void TrackedAllocator::free(const Allocation& allocation) {
311 {
312 std::unique_lock<std::mutex> lock(mutex);
313 auto& usageStats = stats_.byUsage[int(allocation.request.usage)];
314 MARL_ASSERT(usageStats.count > 0,
315 "TrackedAllocator detected abnormal free()");
316 MARL_ASSERT(usageStats.bytes >= allocation.request.size,
317 "TrackedAllocator detected abnormal free()");
318 --usageStats.count;
319 usageStats.bytes -= allocation.request.size;
320 }
321 return allocator->free(allocation);
322}
323
324///////////////////////////////////////////////////////////////////////////////
325// StlAllocator
326///////////////////////////////////////////////////////////////////////////////
327
328// StlAllocator exposes an STL-compatible allocator wrapping a marl::Allocator.
329template <typename T>
330struct StlAllocator {
331 using value_type = T;
332 using pointer = T*;
333 using const_pointer = const T*;
334 using reference = T&;
335 using const_reference = const T&;
336 using size_type = size_t;
337 using difference_type = size_t;
338
339 // An equivalent STL allocator for a different type.
340 template <class U>
341 struct rebind {
342 typedef StlAllocator<U> other;
343 };
344
345 // Constructs an StlAllocator that will allocate using allocator.
346 // allocator must remain valid until this StlAllocator has been destroyed.
347 inline StlAllocator(Allocator* allocator);
348
349 template <typename U>
350 inline StlAllocator(const StlAllocator<U>& other);
351
352 // Returns the actual address of x even in presence of overloaded operator&.
353 inline pointer address(reference x) const;
354 inline const_pointer address(const_reference x) const;
355
356 // Allocates the memory for n objects of type T.
357 // Does not actually construct the objects.
358 inline T* allocate(std::size_t n);
359
360 // Deallocates the memory for n objects of type T.
361 inline void deallocate(T* p, std::size_t n);
362
363 // Returns the maximum theoretically possible number of T stored in this
364 // allocator.
365 inline size_type max_size() const;
366
367 // Copy constructs an object of type T at the address p.
368 inline void construct(pointer p, const_reference val);
369
370 // Constructs an object of type U at the address P forwarning all other
371 // arguments to the constructor.
372 template <typename U, typename... Args>
373 inline void construct(U* p, Args&&... args);
374
375 // Deconstructs the object at p. It does not free the memory.
376 inline void destroy(pointer p);
377
378 // Deconstructs the object at p. It does not free the memory.
379 template <typename U>
380 inline void destroy(U* p);
381
382 private:
383 inline Allocation::Request request(size_t n) const;
384
385 template <typename U>
386 friend struct StlAllocator;
387 Allocator* allocator;
388};
389
390template <typename T>
391StlAllocator<T>::StlAllocator(Allocator* allocator) : allocator(allocator) {}
392
393template <typename T>
394template <typename U>
395StlAllocator<T>::StlAllocator(const StlAllocator<U>& other) {
396 allocator = other.allocator;
397}
398
399template <typename T>
400typename StlAllocator<T>::pointer StlAllocator<T>::address(reference x) const {
401 return &x;
402}
403template <typename T>
404typename StlAllocator<T>::const_pointer StlAllocator<T>::address(
405 const_reference x) const {
406 return &x;
407}
408
409template <typename T>
410T* StlAllocator<T>::allocate(std::size_t n) {
411 auto alloc = allocator->allocate(request(n));
412 return reinterpret_cast<T*>(alloc.ptr);
413}
414
415template <typename T>
416void StlAllocator<T>::deallocate(T* p, std::size_t n) {
417 Allocation alloc;
418 alloc.ptr = p;
419 alloc.request = request(n);
420 allocator->free(alloc);
421}
422
423template <typename T>
424typename StlAllocator<T>::size_type StlAllocator<T>::max_size() const {
425 return std::numeric_limits<size_type>::max() / sizeof(value_type);
426}
427
428template <typename T>
429void StlAllocator<T>::construct(pointer p, const_reference val) {
430 new (p) T(val);
431}
432
433template <typename T>
434template <typename U, typename... Args>
435void StlAllocator<T>::construct(U* p, Args&&... args) {
436 ::new ((void*)p) U(std::forward<Args>(args)...);
437}
438
439template <typename T>
440void StlAllocator<T>::destroy(pointer p) {
441 ((T*)p)->~T();
442}
443
444template <typename T>
445template <typename U>
446void StlAllocator<T>::destroy(U* p) {
447 p->~U();
448}
449
450template <typename T>
451Allocation::Request StlAllocator<T>::request(size_t n) const {
452 Allocation::Request req = {};
453 req.size = sizeof(T) * n;
454 req.alignment = alignof(T);
455 req.usage = Allocation::Usage::Stl;
456 return req;
457}
458
459} // namespace marl
460
461#endif // marl_memory_h
462