xref: /aosp_15_r20/external/executorch/runtime/executor/test/managed_memory_manager.h (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1 /*
2  * Copyright (c) Meta Platforms, Inc. and affiliates.
3  * All rights reserved.
4  *
5  * This source code is licensed under the BSD-style license found in the
6  * LICENSE file in the root directory of this source tree.
7  */
8 
9 #pragma once
10 
11 #include <memory>
12 #include <vector>
13 
14 #include <executorch/runtime/core/hierarchical_allocator.h>
15 #include <executorch/runtime/core/memory_allocator.h>
16 #include <executorch/runtime/executor/memory_manager.h>
17 
18 namespace executorch {
19 namespace runtime {
20 namespace testing {
21 
22 /**
23  * Creates and owns a MemoryManager and the allocators that it points to. Easier
24  * to manage than creating the allocators separately.
25  */
26 class ManagedMemoryManager {
27  public:
28   ManagedMemoryManager(
29       size_t planned_memory_bytes,
30       size_t method_allocator_bytes,
31       MemoryAllocator* temp_allocator = nullptr)
planned_memory_buffer_(new uint8_t[planned_memory_bytes])32       : planned_memory_buffer_(new uint8_t[planned_memory_bytes]),
33         planned_memory_span_(
34             planned_memory_buffer_.get(),
35             planned_memory_bytes),
36         planned_memory_({&planned_memory_span_, 1}),
37         method_allocator_pool_(new uint8_t[method_allocator_bytes]),
38         method_allocator_(method_allocator_bytes, method_allocator_pool_.get()),
39         memory_manager_(&method_allocator_, &planned_memory_, temp_allocator) {}
40 
get()41   MemoryManager& get() {
42     return memory_manager_;
43   }
44 
45  private:
46   std::unique_ptr<uint8_t[]> planned_memory_buffer_;
47   Span<uint8_t> planned_memory_span_;
48   HierarchicalAllocator planned_memory_;
49 
50   std::unique_ptr<uint8_t[]> method_allocator_pool_;
51   MemoryAllocator method_allocator_;
52 
53   MemoryManager memory_manager_;
54 };
55 
56 } // namespace testing
57 } // namespace runtime
58 } // namespace executorch
59 
60 namespace torch {
61 namespace executor {
62 namespace testing {
63 // TODO(T197294990): Remove these deprecated aliases once all users have moved
64 // to the new `::executorch` namespaces.
65 using ::executorch::runtime::testing::ManagedMemoryManager;
66 } // namespace testing
67 } // namespace executor
68 } // namespace torch
69