1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations 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
27namespace 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.
36struct 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.
63class 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