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 { |
37 | 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 | |
45 | 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) |
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 | allocs_erased_(false) {} |
73 | |
74 | // Erases all allocs which have been marked for deletion. This must be called |
75 | // after Deallocate. |
76 | void ResolveDeallocations(); |
77 | |
78 | // Schedule memory allocation for a tensor with a given size, assuming that it |
79 | // needs to be allocated before the execution of first_node, and deallocated |
80 | // after the execution of last_node. |
81 | TfLiteStatus Allocate(TfLiteContext* context, size_t alignment, size_t size, |
82 | int32_t tensor, int32_t first_node, int32_t last_node, |
83 | ArenaAllocWithUsageInterval* new_alloc); |
84 | |
85 | // Marks `alloc` for deletion by setting tensor to -1. After one or |
86 | // more calls to this function, `ResolveDeallocations` must be called to |
87 | // delete all allocs which have been marked, otherwise the arena memory cannot |
88 | // be re-used, increasing overall memory usage. |
89 | TfLiteStatus Deallocate(TfLiteContext* context, |
90 | ArenaAllocWithUsageInterval& alloc); |
91 | |
92 | // Deletes all allocs which are allocated by nodes after `node`. |
93 | // This is equivalent, but much more efficient than calling Deallocate for |
94 | // each alloc individually. |
95 | void DeallocateAfter(int32_t node); |
96 | |
97 | inline size_t RequiredBufferSize() { |
98 | // Add in a small amount of padding to reduce the chance of resize events |
99 | // for small allocations. |
100 | size_t padding = arena_alignment_; |
101 | return arena_alignment_ + high_water_mark_ + padding; |
102 | } |
103 | |
104 | TfLiteStatus Commit(TfLiteContext* context, bool* arena_reallocated); |
105 | |
106 | TfLiteStatus ResolveAlloc(TfLiteContext* context, |
107 | const ArenaAllocWithUsageInterval& alloc, |
108 | char** output_ptr); |
109 | |
110 | // This clears allocation details but does not release the underlying buffer. |
111 | // New allocations should be committed & resolved before using this arena |
112 | // again. |
113 | TfLiteStatus ClearPlan(); |
114 | |
115 | // This releases the underlying buffer but does not clear the allocation plan. |
116 | // Since all associated pointers are invalidated, the arena cannot be used |
117 | // again until Commit() is called & tensor allocations are resolved. |
118 | TfLiteStatus ReleaseBuffer(); |
119 | |
120 | size_t GetBufferSize() const { return underlying_buffer_size_; } |
121 | |
122 | std::intptr_t BasePointer() const { |
123 | return reinterpret_cast<std::intptr_t>(underlying_buffer_aligned_ptr_); |
124 | } |
125 | |
126 | // Dumps the memory allocation information of this memory arena (which could |
127 | // be differentiated from others by the `name`) against the specified op node |
128 | // execution plan (i.e. `execution_plan`) for the purpose of debugging. |
129 | // Note: in order to have minimal binary increase caused by this debug info |
130 | // dump implementation for the TfLite library, and allow users to plug-in |
131 | // their own memory planner debugger, we have utilized weak symbols to meet |
132 | // these two requirementsements. By default, there is no debugging info |
133 | // dumped. To override this, provide a strong defintion of |
134 | // tflite::DumpArenaInfo(...) whose weak defintion is in |
135 | // simple_memory_arena.cc. TfLite provides a sample one as |
136 | // "lite:simple_memory_arena_debug_dump". When this dep is added to the |
137 | // program, calling this function will output information of this memory arena |
138 | // about tenosrs and ops, such as memory arena utilization rate, live tensors |
139 | // at each op etc. |
140 | void DumpDebugInfo(const std::string& name, |
141 | const std::vector<int>& execution_plan) const; |
142 | |
143 | protected: |
144 | int subgraph_index_; |
145 | |
146 | private: |
147 | bool committed_; |
148 | size_t arena_alignment_; |
149 | size_t high_water_mark_; |
150 | std::unique_ptr<char[]> underlying_buffer_; |
151 | size_t underlying_buffer_size_; |
152 | char* underlying_buffer_aligned_ptr_; |
153 | std::vector<ArenaAllocWithUsageInterval> ordered_allocs_; |
154 | bool allocs_erased_; |
155 | }; |
156 | |
157 | } // namespace tflite |
158 | |
159 | #endif // TENSORFLOW_LITE_SIMPLE_MEMORY_ARENA_H_ |
160 | |