1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
/**
* \file src/core/include/megbrain/utils/mempool.h
* MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
*
* Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
*/
#pragma once
#include <algorithm>
#include <cstdint>
#include <memory>
#include <vector>
#include "megbrain_build_config.h"
namespace mgb {
class MemPoolStorage {
bool m_disable_freelist = false;
size_t m_cur_buf_pos = 0, m_cur_buf_size_bytes = 0;
std::vector<std::unique_ptr<uint8_t[]>> m_buf;
std::vector<void*> m_free;
public:
MGE_WIN_DECLSPEC_FUC MemPoolStorage() noexcept;
MGE_WIN_DECLSPEC_FUC MemPoolStorage(MemPoolStorage&& rhs) noexcept;
MGE_WIN_DECLSPEC_FUC ~MemPoolStorage() noexcept;
MGE_WIN_DECLSPEC_FUC MemPoolStorage& operator=(MemPoolStorage&& rhs) noexcept;
MGE_WIN_DECLSPEC_FUC void swap(MemPoolStorage& other);
/*!
* \brief allocate sotrage for an object of specified size
* \param elem_size size of the object; it must remain unchanged
* during lifespan of this MemPoolStorage
*/
MGE_WIN_DECLSPEC_FUC void* alloc(size_t elem_size);
MGE_WIN_DECLSPEC_FUC void free(void* ptr);
MGE_WIN_DECLSPEC_FUC void reorder_free();
//! clear all allocated storage
MGE_WIN_DECLSPEC_FUC void clear();
void disable_freelist() { m_disable_freelist = true; }
};
/*!
* \brief a memory pool for abundant small objects
*
* Note that the memory would not be released and returned to upstream
* allocator until the mem pool is destructed.
*
* The caller must match alloc() and free() calls; no additional check is
* performed.
*/
template <typename T>
class MemPool {
// use another template so T only needs to be complete when alloc() or
// free() is called
template <typename = void>
struct Const {
static constexpr size_t ELEM_SIZE =
((sizeof(T) - 1) / alignof(T) + 1) * alignof(T);
};
MemPoolStorage m_storage;
public:
class Deleter {
MemPool* m_pool = nullptr;
public:
Deleter() = default;
Deleter(MemPool* pool) : m_pool{pool} {}
void operator()(T* ptr) const { m_pool->free(ptr); }
};
using UniquePtr = std::unique_ptr<T, Deleter>;
void* alloc_raw() { return m_storage.alloc(Const<>::ELEM_SIZE); }
void free_raw(void* ptr) { m_storage.free(ptr); }
template <typename... Args>
T* alloc(Args&&... args) {
return new (alloc_raw()) T(std::forward<Args>(args)...);
}
template <typename... Args>
UniquePtr alloc_unique(Args&&... args) {
auto ptr = alloc(std::forward<Args>(args)...);
return {ptr, {this}};
}
void free(T* ptr) {
ptr->~T();
free_raw(ptr);
}
//! reorder free list for cache friendly in future alloc
void reorder_free() { m_storage.reorder_free(); }
//! clear all the storage without calling the destructors
void clear() { m_storage.clear(); }
/*!
* \brief disable free list for memory reuse
*
* This is only useful in the destructor of an enclosing object, so
* no extra memory allocation is needed to hold the released objects
*/
void disable_freelist() { m_storage.disable_freelist(); }
};
} // namespace mgb
// vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}