Skip to content

Commit bc830a3

Browse files
trivedivivektrivedivivek
committed
[ET-VK] Adding MemoryPool class to allocate chunks of memory and use it for smaller allocations.
Pull Request resolved: #7338 ghstack-source-id: 258102837 @exported-using-ghexport Differential Revision: [D67220112](https://our.internmc.facebook.com/intern/diff/D67220112/)
1 parent 1b75687 commit bc830a3

File tree

2 files changed

+219
-0
lines changed

2 files changed

+219
-0
lines changed

‎backends/vulkan/runtime/utils/MemoryPool.cpp‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
*
4+
* This source code is licensed under the MIT license found in the
5+
* LICENSE file in the root directory of this source tree.
6+
*/
7+
8+
#include <executorch/backends/vulkan/runtime/utils/MemoryPool.h>
9+
10+
namespace vkcompute::utils {
11+
12+
// constexpr const size_t kUniquePtrAllocatorInitBlockSize = 16 * 1024;
13+
// constexpr const size_t kUniquePtrAllocatorAlignment = 8;
14+
15+
// void MemoryPool::reset() noexcept(false) {
16+
// for (auto& block : blocks_) {
17+
// block.reset();
18+
// }
19+
// currentBlockIndex_ = 0;
20+
// }
21+
22+
// MemoryPool::BlockInfo::BlockInfo(BlockInfo&& other) noexcept {
23+
// *this = other;
24+
// }
25+
26+
// MemoryPool::BlockInfo& MemoryPool::BlockInfo::operator=(
27+
// MemoryPool::BlockInfo&& other) noexcept {
28+
// totalCapacity_ = other.totalCapacity_;
29+
// currentCapacity_ = other.currentCapacity_;
30+
// data_ = std::move(other.data_);
31+
// return *this;
32+
// }
33+
34+
// MemoryPool::BlockInfo& MemoryPool::BlockInfo::operator=(
35+
// MemoryPool::BlockInfo& other) {
36+
// totalCapacity_ = other.totalCapacity_;
37+
// currentCapacity_ = other.currentCapacity_;
38+
// data_ = std::move(other.data_);
39+
// return *this;
40+
// }
41+
42+
// void MemoryPool::BlockInfo::reset() noexcept(false) {
43+
// currentCapacity_ = totalCapacity_;
44+
// }
45+
46+
// void* MemoryPool::alloc(size_t size) {
47+
// size = ((size + kUniquePtrAllocatorAlignment - 1) /
48+
// kUniquePtrAllocatorAlignment) *
49+
// kUniquePtrAllocatorAlignment;
50+
// if (blocks_.empty() || blocks_[currentBlockIndex_].currentCapacity_ < size)
51+
// {
52+
// // find a block that has enough space
53+
// bool found = false;
54+
// for (uint32_t i = currentBlockIndex_ + 1; i < blocks_.size(); ++i) {
55+
// if (blocks_[i].currentCapacity_ >= size) {
56+
// currentBlockIndex_ = i;
57+
// found = true;
58+
// break;
59+
// }
60+
// }
61+
62+
// if (!found) { // no block has enough space, create a new one
63+
// auto& blockInfo = blocks_.emplace_back(BlockInfo{});
64+
// blockInfo.totalCapacity_ =
65+
// std::max(size, kUniquePtrAllocatorInitBlockSize);
66+
// blockInfo.currentCapacity_ = blockInfo.totalCapacity_;
67+
// blockInfo.data_ =
68+
// std::make_unique<uint8_t[]>(blockInfo.totalCapacity_);
69+
// currentBlockIndex_ = blocks_.size() - 1;
70+
// }
71+
// }
72+
73+
// auto& blockInfo = blocks_[currentBlockIndex_];
74+
// auto ret = blockInfo.data_.get() +
75+
// (blockInfo.totalCapacity_ - blockInfo.currentCapacity_);
76+
// // blockInfo.allocations_.push_back(reinterpret_cast<AllocatorBase*>(ret));
77+
// blockInfo.currentCapacity_ -= size;
78+
// return ret;
79+
// }
80+
81+
} // namespace vkcompute::utils

‎backends/vulkan/runtime/utils/MemoryPool.h‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
*
4+
* This source code is licensed under the MIT license found in the
5+
* LICENSE file in the root directory of this source tree.
6+
*/
7+
8+
#pragma once
9+
10+
#include <memory>
11+
#include <vector>
12+
13+
namespace vkcompute::utils {
14+
15+
/**
16+
* @class This class is designed to improving memory allocation for unique
17+
*pointer allocations for a class or derivatives of a class. Allocation is done
18+
*in blocks and every unique pointer allocation takes chunk from the block.
19+
*
20+
* This class is designed to be a fast alternative to calling new and delete for
21+
*every unique pointer. Its speed is mainly due to the fact that it allocates
22+
*memory in successive location in a block and does not try to reuse indivudual
23+
*memory location after its deallocated.
24+
**/
25+
template <class BaseType, const uint32_t kInitBlockSize = 4 * 1024, const uint32_t kAllocAlignment = 16>
26+
class MemoryPool final {
27+
public:
28+
MemoryPool() = default;
29+
virtual ~MemoryPool() = default;
30+
31+
struct Deleter final // deleter
32+
{
33+
inline void operator()(BaseType* obj) const {
34+
if (obj) {
35+
obj->~BaseType();
36+
}
37+
}
38+
};
39+
40+
struct alignas(kAllocAlignment) AlignedStruct final {
41+
uint8_t data[kAllocAlignment];
42+
};
43+
44+
/**
45+
* @brief Allocates a new object of Type type by allocating memory from an
46+
*allocation block and returns a unique pointer to it.
47+
* @param ref - reference object to construct object from.
48+
**/
49+
template <class Type>
50+
std::unique_ptr<BaseType, MemoryPool::Deleter> make_unique(Type&& ref) {
51+
static_assert(
52+
kAllocAlignment >= alignof(Type),
53+
"Memory pool can only safely allocate object of alignment kAllocAlignment or less");
54+
auto ptr = std::unique_ptr<BaseType, MemoryPool::Deleter>(
55+
new (alloc(sizeof(Type))) Type(std::move(ref)), MemoryPool::Deleter());
56+
return ptr;
57+
}
58+
59+
/**
60+
* @brief Reset all block allocation info, so it can be reused.
61+
* Call this function only after all the unique pointers allcated by this
62+
* block have been deallocated. Else it may result in memory leaks.
63+
**/
64+
void reset() noexcept(false) {
65+
for (auto& block : blocks_) {
66+
block.reset();
67+
}
68+
currentBlockIndex_ = 0;
69+
}
70+
71+
private:
72+
struct BlockInfo {
73+
uint32_t totalCapacity_ = 0;
74+
uint32_t currentCapacity_ = 0;
75+
std::unique_ptr<AlignedStruct[]> data_{};
76+
77+
BlockInfo() = default;
78+
~BlockInfo() = default;
79+
80+
BlockInfo(BlockInfo&& other) noexcept {
81+
*this = other;
82+
}
83+
84+
BlockInfo& operator=(BlockInfo&& other) noexcept {
85+
totalCapacity_ = other.totalCapacity_;
86+
currentCapacity_ = other.currentCapacity_;
87+
data_ = std::move(other.data_);
88+
return *this;
89+
}
90+
91+
BlockInfo& operator=(BlockInfo& other) {
92+
totalCapacity_ = other.totalCapacity_;
93+
currentCapacity_ = other.currentCapacity_;
94+
data_ = std::move(other.data_);
95+
return *this;
96+
}
97+
98+
void reset() noexcept(false) {
99+
currentCapacity_ = totalCapacity_;
100+
}
101+
};
102+
103+
std::vector<BlockInfo> blocks_{};
104+
uint32_t currentBlockIndex_ = 0;
105+
106+
void* alloc(size_t size) {
107+
size =
108+
((size + kAllocAlignment - 1) / kAllocAlignment) *
109+
kAllocAlignment;
110+
if (blocks_.empty() ||
111+
blocks_[currentBlockIndex_].currentCapacity_ < size) {
112+
// find a block that has enough space
113+
bool found = false;
114+
for (uint32_t i = currentBlockIndex_ + 1; i < blocks_.size(); ++i) {
115+
if (blocks_[i].currentCapacity_ >= size) {
116+
currentBlockIndex_ = i;
117+
found = true;
118+
break;
119+
}
120+
}
121+
122+
if (!found) { // no block has enough space, create a new one
123+
auto& blockInfo = blocks_.emplace_back(BlockInfo{});
124+
blockInfo.totalCapacity_ = std::max(size, kInitBlockSize);
125+
blockInfo.currentCapacity_ = blockInfo.totalCapacity_;
126+
blockInfo.data_ = std::make_unique<AlignedStruct[]>((blockInfo.totalCapacity_ + sizeof(AlignedStruct) - 1) / sizeof(AlignedStruct));
127+
currentBlockIndex_ = blocks_.size() - 1;
128+
}
129+
}
130+
131+
auto& blockInfo = blocks_[currentBlockIndex_];
132+
auto ret = blockInfo.data_.get()->data + (blockInfo.totalCapacity_ - blockInfo.currentCapacity_);
133+
blockInfo.currentCapacity_ -= size;
134+
return ret;
135+
}
136+
};
137+
138+
} // namespace vkcompute::utils

0 commit comments

Comments
 (0)