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 <executorch/runtime/core/hierarchical_allocator.h> 12 #include <executorch/runtime/core/memory_allocator.h> 13 14 namespace executorch { 15 namespace runtime { 16 17 /** 18 * A container class for allocators used during Method load and execution. 19 * 20 * This class consolidates all dynamic memory needs for Method load and 21 * execution. This can allow for heap-based as well as heap-less execution 22 * (relevant to some embedded scenarios), and overall provides more control over 23 * memory use. 24 * 25 * This class, however, cannot ensure all allocation is accounted for since 26 * kernel and backend implementations are free to use a separate way to allocate 27 * memory (e.g., for things like scratch space). But we do suggest that backends 28 * and kernels use these provided allocators whenever possible. 29 */ 30 class MemoryManager final { 31 public: 32 /** 33 * Constructs a new MemoryManager. 34 * 35 * @param[in] method_allocator The allocator to use when loading a Method and 36 * allocating its internal structures. Must outlive the Method that uses 37 * it. 38 * @param[in] planned_memory The memory-planned buffers to use for mutable 39 * tensor data when executing a Method. Must outlive the Method that uses 40 * it. May be `nullptr` if the Method does not use any memory-planned 41 * tensor data. The sizes of the buffers in this HierarchicalAllocator 42 * must agree with the corresponding 43 * `MethodMeta::num_memory_planned_buffers()` and 44 * `MethodMeta::memory_planned_buffer_size(N)` values, which are embedded 45 * in the Program. 46 * @param[in] temp_allocator The allocator to use when allocating temporary 47 * data during kernel or delegate execution. Must outlive the Method that 48 * uses it. May be `nullptr` if the Method does not use kernels or 49 * delegates that allocate temporary data. This allocator will be reset 50 * after every kernel or delegate call during execution. 51 */ 52 explicit MemoryManager( 53 MemoryAllocator* method_allocator, 54 HierarchicalAllocator* planned_memory = nullptr, 55 MemoryAllocator* temp_allocator = nullptr) method_allocator_(method_allocator)56 : method_allocator_(method_allocator), 57 planned_memory_(planned_memory), 58 temp_allocator_(temp_allocator) { 59 ET_CHECK_MSG( 60 method_allocator != temp_allocator, 61 "method allocator cannot be the same as temp allocator"); 62 } 63 64 /** 65 * DEPRECATED: Use the constructor without `constant_allocator` instead. 66 * 67 * TODO(T162089316): Remove this once all users migrate to the new ctor. 68 */ MemoryManager(MemoryAllocator * constant_allocator,HierarchicalAllocator * non_constant_allocator,MemoryAllocator * runtime_allocator,MemoryAllocator * temporary_allocator)69 ET_DEPRECATED MemoryManager( 70 MemoryAllocator* constant_allocator, 71 HierarchicalAllocator* non_constant_allocator, 72 MemoryAllocator* runtime_allocator, 73 MemoryAllocator* temporary_allocator) 74 : MemoryManager( 75 /*method_allocator=*/runtime_allocator, 76 /*planned_memory=*/non_constant_allocator, 77 /*temp_allocator=*/temporary_allocator) { 78 (void)constant_allocator; // Suppress unused variable warning 79 } 80 81 /** 82 * Returns the allocator that the runtime will use to allocate internal 83 * structures while loading a Method. Must not be used after its associated 84 * Method has been loaded. 85 */ method_allocator()86 MemoryAllocator* method_allocator() const { 87 return method_allocator_; 88 } 89 90 /** 91 * Returns the memory-planned buffers to use for mutable tensor data. 92 */ planned_memory()93 HierarchicalAllocator* planned_memory() const { 94 return planned_memory_; 95 } 96 97 /** 98 * Returns the allocator to use for allocating temporary data during kernel or 99 * delegate execution. 100 * 101 * This allocator will be reset after every kernel or delegate call during 102 * execution. 103 */ temp_allocator()104 MemoryAllocator* temp_allocator() const { 105 return temp_allocator_; 106 } 107 108 private: 109 MemoryAllocator* method_allocator_; 110 HierarchicalAllocator* planned_memory_; 111 MemoryAllocator* temp_allocator_; 112 }; 113 114 } // namespace runtime 115 } // namespace executorch 116 117 namespace torch { 118 namespace executor { 119 // TODO(T197294990): Remove these deprecated aliases once all users have moved 120 // to the new `::executorch` namespaces. 121 using ::executorch::runtime::MemoryManager; 122 } // namespace executor 123 } // namespace torch 124