Skip to content

Commit fedc04c

Browse files
dbortfacebook-github-bot
authored andcommitted
Migrate runner-like targets to use the new MemoryManager API (#403)
Summary: Pull Request resolved: #403 Use the new APIs added in D49389211. Note that this simplifies the user code, and the names should be more clear. ghstack-source-id: 201207624 exported-using-ghexport Reviewed By: JacobSzwejbka Differential Revision: D49389914 fbshipit-source-id: cafb19d86b6bad3d96e46a12cfb5ffcc056115df
1 parent 6944c45 commit fedc04c

File tree

6 files changed

+141
-222
lines changed

6 files changed

+141
-222
lines changed

examples/bundled_executor_runner/bundled_executor_runner.cpp

Lines changed: 31 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,7 @@
3131
#include <executorch/util/bundled_program_verification.h>
3232
#include <executorch/util/util.h>
3333

34-
static constexpr size_t kRuntimeMemorySize = 4 * 1024U * 1024U; // 4 MB
35-
static uint8_t runtime_pool[kRuntimeMemorySize];
34+
static uint8_t method_allocator_pool[4 * 1024U * 1024U]; // 4MB
3635
static constexpr size_t kBundledAllocatorPoolSize = 16 * 1024U;
3736
static uint8_t bundled_allocator_pool[kBundledAllocatorPoolSize];
3837

@@ -138,70 +137,52 @@ int main(int argc, char** argv) {
138137
// do it dynamically.
139138
//
140139

141-
// The runtime allocator is used to allocate all dynamic C++ metadata/objects
142-
// used to represent the loaded program. This allocator is only used during
140+
// The method allocator is used to allocate all dynamic C++ metadata/objects
141+
// used to represent the loaded method. This allocator is only used during
143142
// loading a method of the program, which will return an error if there was
144143
// not enough memory.
145144
//
146-
// The amount of memory required depends on the loaded program and the runtime
145+
// The amount of memory required depends on the loaded method and the runtime
147146
// code itself. The amount of memory here is usually determined by running the
148-
// program and seeing how much memory is actually used, though it's possible
149-
// to subclass MemoryAllocator so that it calls malloc() under the hood.
150-
151-
// In this example we using statically allocated gloabl runtime_pool of
152-
// size kRuntimeMemorySize
153-
MemoryAllocator runtime_allocator{
154-
MemoryAllocator(kRuntimeMemorySize, runtime_pool)};
155-
runtime_allocator.enable_profiling("runtime allocator");
147+
// method and seeing how much memory is actually used, though it's possible to
148+
// subclass MemoryAllocator so that it calls malloc() under the hood (see
149+
// MallocMemoryAllocator).
150+
//
151+
// In this example we use a statically allocated memory pool.
152+
MemoryAllocator method_allocator{
153+
MemoryAllocator(sizeof(method_allocator_pool), method_allocator_pool)};
154+
method_allocator.enable_profiling("method allocator");
156155

157-
// The non-const buffers will back the mutable tensors used by the method. The
158-
// sizes of these buffers were determined ahead of time during the
156+
// The memory-planned buffers will back the mutable tensors used by the
157+
// method. The sizes of these buffers were determined ahead of time during the
159158
// memory-planning pasees.
160159
//
161160
// Each buffer typically corresponds to a different hardware memory bank. Most
162161
// mobile environments will only have a single buffer. Some embedded
163162
// environments may have more than one for, e.g., slow/large DRAM and
164163
// fast/small SRAM, or for memory associated with particular cores.
165-
std::vector<std::unique_ptr<uint8_t[]>> non_const_buffers;
166-
std::vector<Span<uint8_t>> non_const_spans;
167-
size_t num_non_const_buffers = method_meta->num_non_const_buffers();
168-
for (size_t id = 0; id < num_non_const_buffers; ++id) {
169-
// .get() will always succeed because id < num_non_const_buffers.
164+
std::vector<std::unique_ptr<uint8_t[]>> planned_buffers; // Owns the memory
165+
std::vector<Span<uint8_t>> planned_spans; // Passed to the allocator
166+
size_t num_memory_planned_buffers = method_meta->num_memory_planned_buffers();
167+
for (size_t id = 0; id < num_memory_planned_buffers; ++id) {
168+
// .get() will always succeed because id < num_memory_planned_buffers.
170169
size_t buffer_size =
171-
static_cast<size_t>(method_meta->non_const_buffer_size(id).get());
172-
ET_LOG(Info, "Setting up non-const buffer %zu, size %zu.", id, buffer_size);
173-
non_const_buffers.push_back(std::make_unique<uint8_t[]>(buffer_size));
174-
non_const_spans.push_back({non_const_buffers.back().get(), buffer_size});
170+
static_cast<size_t>(method_meta->memory_planned_buffer_size(id).get());
171+
ET_LOG(Info, "Setting up planned buffer %zu, size %zu.", id, buffer_size);
172+
planned_buffers.push_back(std::make_unique<uint8_t[]>(buffer_size));
173+
planned_spans.push_back({planned_buffers.back().get(), buffer_size});
175174
}
176-
HierarchicalAllocator non_const_allocator(
177-
{non_const_spans.data(), non_const_spans.size()});
178-
179-
// Allocator for bundled input.
180-
MemoryAllocator bundled_input_allocator{
181-
MemoryAllocator(kBundledAllocatorPoolSize, bundled_allocator_pool)};
182-
183-
// The constant allocator is not currently used. Please initialize with a
184-
// zero-sized allocator.
185-
MemoryAllocator const_allocator{MemoryAllocator(0, nullptr)};
186-
const_allocator.enable_profiling("const allocator");
187-
188-
// The kernel temporary allocator is not currently used. Please initialize
189-
// with a zero-sized allocator.
190-
MemoryAllocator temp_allocator{MemoryAllocator(0, nullptr)};
191-
temp_allocator.enable_profiling("temp allocator");
175+
HierarchicalAllocator planned_memory(
176+
{planned_spans.data(), planned_spans.size()});
192177

193178
// Assemble all of the allocators into the MemoryManager that the Executor
194179
// will use.
195-
MemoryManager memory_manager(
196-
&const_allocator,
197-
&non_const_allocator,
198-
&runtime_allocator,
199-
&temp_allocator);
180+
MemoryManager memory_manager(&method_allocator, &planned_memory);
200181

201182
//
202-
// Load method from the program, using the provided
203-
// allocators. Running the method can mutate allocated non_const buffers,
204-
// so should only be used by a single thread at at time, but it can be reused.
183+
// Load the method from the program, using the provided allocators. Running
184+
// the method can mutate the memory-planned buffers, so the method should only
185+
// be used by a single thread at at time, but it can be reused.
205186
//
206187

207188
Result<Method> method = program->load_method(method_name, &memory_manager);
@@ -214,6 +195,8 @@ int main(int argc, char** argv) {
214195

215196
// Prepare the inputs.
216197
// Use ones-initialized inputs or bundled inputs.
198+
MemoryAllocator bundled_input_allocator{
199+
MemoryAllocator(kBundledAllocatorPoolSize, bundled_allocator_pool)};
217200
exec_aten::ArrayRef<void*> inputs;
218201
if (FLAGS_bundled_program) {
219202
// Use the inputs embedded in the bundled program.

examples/executor_runner/executor_runner.cpp

Lines changed: 29 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,7 @@
2929
#include <executorch/runtime/platform/runtime.h>
3030
#include <executorch/util/util.h>
3131

32-
static constexpr size_t kRuntimeMemorySize = 4 * 1024U * 1024U; // 4 MB
33-
static uint8_t runtime_pool[kRuntimeMemorySize];
32+
static uint8_t method_allocator_pool[4 * 1024U * 1024U]; // 4 MB
3433

3534
DEFINE_string(
3635
model_path,
@@ -98,66 +97,52 @@ int main(int argc, char** argv) {
9897
// do it dynamically.
9998
//
10099

101-
// The runtime allocator is used to allocate all dynamic C++ metadata/objects
102-
// used to represent the loaded program. This allocator is only used during
100+
// The method allocator is used to allocate all dynamic C++ metadata/objects
101+
// used to represent the loaded method. This allocator is only used during
103102
// loading a method of the program, which will return an error if there was
104103
// not enough memory.
105104
//
106-
// The amount of memory required depends on the loaded program and the runtime
105+
// The amount of memory required depends on the loaded method and the runtime
107106
// code itself. The amount of memory here is usually determined by running the
108-
// program and seeing how much memory is actually used, though it's possible
109-
// to subclass MemoryAllocator so that it calls malloc() under the hood.
110-
111-
// In this example we using statically allocated gloabl runtime_pool of
112-
// size kRuntimeMemorySize
113-
MemoryAllocator runtime_allocator{
114-
MemoryAllocator(kRuntimeMemorySize, runtime_pool)};
115-
runtime_allocator.enable_profiling("runtime allocator");
107+
// method and seeing how much memory is actually used, though it's possible to
108+
// subclass MemoryAllocator so that it calls malloc() under the hood (see
109+
// MallocMemoryAllocator).
110+
//
111+
// In this example we use a statically allocated memory pool.
112+
MemoryAllocator method_allocator{
113+
MemoryAllocator(sizeof(method_allocator_pool), method_allocator_pool)};
114+
method_allocator.enable_profiling("method allocator");
116115

117-
// The non-const buffers will back the mutable tensors used by the method. The
118-
// sizes of these buffers were determined ahead of time during the
116+
// The memory-planned buffers will back the mutable tensors used by the
117+
// method. The sizes of these buffers were determined ahead of time during the
119118
// memory-planning pasees.
120119
//
121120
// Each buffer typically corresponds to a different hardware memory bank. Most
122121
// mobile environments will only have a single buffer. Some embedded
123122
// environments may have more than one for, e.g., slow/large DRAM and
124123
// fast/small SRAM, or for memory associated with particular cores.
125-
std::vector<std::unique_ptr<uint8_t[]>> non_const_buffers;
126-
std::vector<Span<uint8_t>> non_const_spans;
127-
size_t num_non_const_buffers = method_meta->num_non_const_buffers();
128-
for (size_t id = 0; id < num_non_const_buffers; ++id) {
129-
// .get() will always succeed because id < num_non_const_buffers.
124+
std::vector<std::unique_ptr<uint8_t[]>> planned_buffers; // Owns the memory
125+
std::vector<Span<uint8_t>> planned_spans; // Passed to the allocator
126+
size_t num_memory_planned_buffers = method_meta->num_memory_planned_buffers();
127+
for (size_t id = 0; id < num_memory_planned_buffers; ++id) {
128+
// .get() will always succeed because id < num_memory_planned_buffers.
130129
size_t buffer_size =
131-
static_cast<size_t>(method_meta->non_const_buffer_size(id).get());
132-
ET_LOG(Info, "Setting up non-const buffer %zu, size %zu.", id, buffer_size);
133-
non_const_buffers.push_back(std::make_unique<uint8_t[]>(buffer_size));
134-
non_const_spans.push_back({non_const_buffers.back().get(), buffer_size});
130+
static_cast<size_t>(method_meta->memory_planned_buffer_size(id).get());
131+
ET_LOG(Info, "Setting up planned buffer %zu, size %zu.", id, buffer_size);
132+
planned_buffers.push_back(std::make_unique<uint8_t[]>(buffer_size));
133+
planned_spans.push_back({planned_buffers.back().get(), buffer_size});
135134
}
136-
HierarchicalAllocator non_const_allocator(
137-
{non_const_spans.data(), non_const_spans.size()});
138-
139-
// The constant allocator is not currently used. Please initialize with a
140-
// zero-sized allocator.
141-
MemoryAllocator const_allocator{MemoryAllocator(0, nullptr)};
142-
const_allocator.enable_profiling("const allocator");
143-
144-
// The kernel temporary allocator is not currently used. Please initialize
145-
// with a zero-sized allocator.
146-
MemoryAllocator temp_allocator{MemoryAllocator(0, nullptr)};
147-
temp_allocator.enable_profiling("temp allocator");
135+
HierarchicalAllocator planned_memory(
136+
{planned_spans.data(), planned_spans.size()});
148137

149138
// Assemble all of the allocators into the MemoryManager that the Executor
150139
// will use.
151-
MemoryManager memory_manager(
152-
&const_allocator,
153-
&non_const_allocator,
154-
&runtime_allocator,
155-
&temp_allocator);
140+
MemoryManager memory_manager(&method_allocator, &planned_memory);
156141

157142
//
158-
// Load method from the program, using the provided
159-
// allocators. Running the method can mutate allocated non_const buffers,
160-
// so should only be used by a single thread at at time, but it can be reused.
143+
// Load the method from the program, using the provided allocators. Running
144+
// the method can mutate the memory-planned buffers, so the method should only
145+
// be used by a single thread at at time, but it can be reused.
161146
//
162147

163148
Result<Method> method = program->load_method(method_name, &memory_manager);

exir/backend/test/demos/rpc/ExecutorBackend.cpp

Lines changed: 14 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -73,50 +73,39 @@ class ExecutorBackend final : public PyTorchBackendInterface {
7373
}
7474

7575
// Building all different allocators for the client executor
76-
auto client_const_allocator = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
77-
runtime_allocator, MemoryAllocator);
78-
new (client_const_allocator) MemoryAllocator(0, nullptr);
79-
80-
auto num_non_const_buffers = method_meta->num_non_const_buffers();
76+
auto num_memory_planned_buffers = method_meta->num_memory_planned_buffers();
8177

82-
Span<uint8_t>* non_const_buffers = ET_ALLOCATE_LIST_OR_RETURN_ERROR(
83-
runtime_allocator, Span<uint8_t>, num_non_const_buffers);
78+
Span<uint8_t>* memory_planned_buffers = ET_ALLOCATE_LIST_OR_RETURN_ERROR(
79+
runtime_allocator, Span<uint8_t>, num_memory_planned_buffers);
8480

85-
for (size_t id = 0; id < num_non_const_buffers; ++id) {
86-
size_t buffer_size =
87-
static_cast<size_t>(method_meta->non_const_buffer_size(id).get());
81+
for (size_t id = 0; id < num_memory_planned_buffers; ++id) {
82+
size_t buffer_size = static_cast<size_t>(
83+
method_meta->memory_planned_buffer_size(id).get());
8884
uint8_t* buffer_i = ET_ALLOCATE_LIST_OR_RETURN_ERROR(
8985
runtime_allocator, uint8_t, buffer_size);
90-
non_const_buffers[id] = {buffer_i, buffer_size};
86+
memory_planned_buffers[id] = {buffer_i, buffer_size};
9187
}
9288

93-
auto client_non_const_allocator = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
89+
auto client_planned_memory = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
9490
runtime_allocator, HierarchicalAllocator);
95-
new (client_non_const_allocator)
96-
HierarchicalAllocator({non_const_buffers, num_non_const_buffers});
91+
new (client_planned_memory) HierarchicalAllocator(
92+
{memory_planned_buffers, num_memory_planned_buffers});
9793

9894
// Allocate some memory from runtime allocator for the client executor, in
9995
// real case, like if it's an executor in dsp, it should allocate memory
10096
// dedicated to this specific hardware
101-
auto client_runtime_allocator = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
97+
auto client_method_allocator = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
10298
runtime_allocator, MemoryAllocator);
10399
const size_t kClientRuntimeMemorySize = 4 * 1024U;
104100
auto runtime_pool = ET_ALLOCATE_OR_RETURN_ERROR(
105101
runtime_allocator, kClientRuntimeMemorySize);
106-
new (client_runtime_allocator) MemoryAllocator(
102+
new (client_method_allocator) MemoryAllocator(
107103
kClientRuntimeMemorySize, static_cast<uint8_t*>(runtime_pool));
108104

109-
auto client_temp_allocator = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
110-
runtime_allocator, MemoryAllocator);
111-
new (client_temp_allocator) MemoryAllocator(0, nullptr);
112-
113105
auto client_memory_manager =
114106
ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(runtime_allocator, MemoryManager);
115-
new (client_memory_manager) MemoryManager(
116-
client_const_allocator,
117-
client_non_const_allocator,
118-
client_runtime_allocator,
119-
client_temp_allocator);
107+
new (client_memory_manager)
108+
MemoryManager(client_method_allocator, client_planned_memory);
120109

121110
// Construct the client Method
122111
Result<Method> method_res =

0 commit comments

Comments
 (0)