31
31
#include < executorch/util/bundled_program_verification.h>
32
32
#include < executorch/util/util.h>
33
33
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
36
35
static constexpr size_t kBundledAllocatorPoolSize = 16 * 1024U ;
37
36
static uint8_t bundled_allocator_pool[kBundledAllocatorPoolSize ];
38
37
@@ -138,70 +137,52 @@ int main(int argc, char** argv) {
138
137
// do it dynamically.
139
138
//
140
139
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
143
142
// loading a method of the program, which will return an error if there was
144
143
// not enough memory.
145
144
//
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
147
146
// 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" );
156
155
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
159
158
// memory-planning pasees.
160
159
//
161
160
// Each buffer typically corresponds to a different hardware memory bank. Most
162
161
// mobile environments will only have a single buffer. Some embedded
163
162
// environments may have more than one for, e.g., slow/large DRAM and
164
163
// 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 .
170
169
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});
175
174
}
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 ()});
192
177
193
178
// Assemble all of the allocators into the MemoryManager that the Executor
194
179
// 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);
200
181
201
182
//
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.
205
186
//
206
187
207
188
Result<Method> method = program->load_method (method_name, &memory_manager);
@@ -214,6 +195,8 @@ int main(int argc, char** argv) {
214
195
215
196
// Prepare the inputs.
216
197
// Use ones-initialized inputs or bundled inputs.
198
+ MemoryAllocator bundled_input_allocator{
199
+ MemoryAllocator (kBundledAllocatorPoolSize , bundled_allocator_pool)};
217
200
exec_aten::ArrayRef<void *> inputs;
218
201
if (FLAGS_bundled_program) {
219
202
// Use the inputs embedded in the bundled program.
0 commit comments