1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #ifndef TENSORFLOW_LITE_SIMPLE_MEMORY_ARENA_H_ 16 #define TENSORFLOW_LITE_SIMPLE_MEMORY_ARENA_H_ 17 18 #include <stddef.h> 19 20 #include <cstdint> 21 #include <memory> 22 #include <string> 23 #include <vector> 24 25 #include "tensorflow/lite/c/common.h" 26 27 namespace tflite { 28 29 // This little structure holds the offset and the size for a dynamic memory 30 // allocation in the memory arena as well as first_node and last_node that use 31 // corresponding tensor. It means that continuous part of memory with this size 32 // needs to be allocated before execution of operation in the first node and can 33 // be deallocated after execution of the operation in the last_node. When the 34 // arena is committed and the underlying buffer is set, the alloc can be 35 // resolved into an actual memory pointer. 36 struct ArenaAllocWithUsageInterval { ArenaAllocWithUsageIntervalArenaAllocWithUsageInterval37 ArenaAllocWithUsageInterval() { reset(); } 38 39 size_t offset; 40 size_t size; 41 int32_t tensor; 42 int32_t first_node; 43 int32_t last_node; 44 resetArenaAllocWithUsageInterval45 inline void reset() { 46 offset = 0; 47 size = 0; 48 tensor = -1; 49 first_node = -1; 50 last_node = -1; 51 } 52 53 inline bool operator<(const ArenaAllocWithUsageInterval& other) const { 54 return offset < other.offset; 55 } 56 }; 57 58 // This small class is responsible for allocating, deallocating and reusing 59 // dynamic memory from a common underlying buffer. The arena can be used in 60 // scenarios when the pattern of memory allocations and deallocations is 61 // repetitive, e.g. running NN inference in multiple iterations. Note that 62 // zero-sized allocations are explicitly allowed, and will resolve to null. 63 class SimpleMemoryArena { 64 public: 65 explicit SimpleMemoryArena(size_t arena_alignment, int subgraph_index = 0) subgraph_index_(subgraph_index)66 : subgraph_index_(subgraph_index), 67 committed_(false), 68 arena_alignment_(arena_alignment), 69 high_water_mark_(0), 70 underlying_buffer_size_(0), 71 ordered_allocs_() {} 72 73 // Schedule memory allocation for a tensor with a given size, assuming that it 74 // needs to be allocated before the execution of first_node, and deallocated 75 // after the execution of last_node. 76 TfLiteStatus Allocate(TfLiteContext* context, size_t alignment, size_t size, 77 int32_t tensor, int32_t first_node, int32_t last_node, 78 ArenaAllocWithUsageInterval* new_alloc); 79 80 TfLiteStatus Deallocate(TfLiteContext* context, 81 const ArenaAllocWithUsageInterval& alloc); 82 RequiredBufferSize()83 inline size_t RequiredBufferSize() { 84 // Add in a small amount of padding to reduce the chance of resize events 85 // for small allocations. 86 size_t padding = arena_alignment_; 87 return arena_alignment_ + high_water_mark_ + padding; 88 } 89 90 TfLiteStatus Commit(TfLiteContext* context); 91 92 TfLiteStatus ResolveAlloc(TfLiteContext* context, 93 const ArenaAllocWithUsageInterval& alloc, 94 char** output_ptr); 95 96 // This clears allocation details but does not release the underlying buffer. 97 // New allocations should be committed & resolved before using this arena 98 // again. 99 TfLiteStatus ClearPlan(); 100 101 // This releases the underlying buffer but does not clear the allocation plan. 102 // Since all associated pointers are invalidated, the arena cannot be used 103 // again until Commit() is called & tensor allocations are resolved. 104 TfLiteStatus ReleaseBuffer(); 105 GetBufferSize()106 size_t GetBufferSize() const { return underlying_buffer_size_; } 107 BasePointer()108 std::intptr_t BasePointer() const { 109 return reinterpret_cast<std::intptr_t>(underlying_buffer_aligned_ptr_); 110 } 111 112 // Dumps the memory allocation information of this memory arena (which could 113 // be differentiated from others by the `name`) against the specified op node 114 // execution plan (i.e. `execution_plan`) for the purpose of debugging. 115 // Note: in order to have minimal binary increase caused by this debug info 116 // dump implementation for the TfLite library, and allow users to plug-in 117 // their own memory planner debugger, we have utilized weak symbols to meet 118 // these two requirementsements. By default, there is no debugging info 119 // dumped. To override this, provide a strong defintion of 120 // tflite::DumpArenaInfo(...) whose weak defintion is in 121 // simple_memory_arena.cc. TfLite provides a sample one as 122 // "lite:simple_memory_arena_debug_dump". When this dep is added to the 123 // program, calling this function will output information of this memory arena 124 // about tenosrs and ops, such as memory arena utilization rate, live tensors 125 // at each op etc. 126 void DumpDebugInfo(const std::string& name, 127 const std::vector<int>& execution_plan) const; 128 129 protected: 130 int subgraph_index_; 131 132 private: 133 bool committed_; 134 size_t arena_alignment_; 135 size_t high_water_mark_; 136 std::unique_ptr<char[]> underlying_buffer_; 137 size_t underlying_buffer_size_; 138 char* underlying_buffer_aligned_ptr_; 139 std::vector<ArenaAllocWithUsageInterval> ordered_allocs_; 140 }; 141 142 } // namespace tflite 143 144 #endif // TENSORFLOW_LITE_SIMPLE_MEMORY_ARENA_H_ 145