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
16 #include "tensorflow/lite/simple_memory_arena.h"
17
18 #include <stddef.h>
19 #include <stdint.h>
20
21 #include <algorithm>
22 #include <cstring>
23 #include <iterator>
24 #include <limits>
25 #include <memory>
26 #include <string>
27 #include <vector>
28
29 #include "tensorflow/lite/c/common.h"
30 #include "tensorflow/lite/core/macros.h"
31 #ifdef TF_LITE_TENSORFLOW_PROFILER
32 #include "tensorflow/lite/tensorflow_profiler_logger.h"
33 #endif // TF_LITE_TENSORFLOW_PROFILER
34
35 namespace {
36
37 template <typename T>
AlignTo(size_t alignment,T offset)38 T AlignTo(size_t alignment, T offset) {
39 return offset % alignment == 0 ? offset
40 : offset + (alignment - offset % alignment);
41 }
42
43 } // namespace
44
45 namespace tflite {
46
Allocate(TfLiteContext * context,size_t alignment,size_t size,int32_t tensor,int32_t first_node,int32_t last_node,ArenaAllocWithUsageInterval * new_alloc)47 TfLiteStatus SimpleMemoryArena::Allocate(
48 TfLiteContext* context, size_t alignment, size_t size, int32_t tensor,
49 int32_t first_node, int32_t last_node,
50 ArenaAllocWithUsageInterval* new_alloc) {
51 TF_LITE_ENSURE(context, alignment <= arena_alignment_);
52 new_alloc->tensor = tensor;
53 new_alloc->first_node = first_node;
54 new_alloc->last_node = last_node;
55 new_alloc->size = size;
56 if (size == 0) {
57 new_alloc->offset = 0;
58 return kTfLiteOk;
59 }
60
61 // If we don't find a better gap just allocate at the end of the buffer.
62 const size_t kOffsetNotAssigned = std::numeric_limits<size_t>::max();
63 size_t best_offset = kOffsetNotAssigned;
64 size_t best_offset_fit = kOffsetNotAssigned;
65
66 // Go through the sorted allocs and look at the gaps between them.
67 size_t current_offset = 0;
68 for (const auto& alloc : ordered_allocs_) {
69 if (alloc.last_node < first_node || alloc.first_node > last_node) {
70 // Usage interval of alloc doesn't intersect with current tensor's usage
71 // interval, so we skip it.
72 continue;
73 }
74 size_t aligned_current_offset = AlignTo(alignment, current_offset);
75 // If we found a gap larger than required size, and smaller than previous
76 // best fit, take it.
77 if (aligned_current_offset + size <= alloc.offset &&
78 alloc.offset - aligned_current_offset < best_offset_fit) {
79 best_offset = aligned_current_offset;
80 best_offset_fit = alloc.offset - current_offset;
81 }
82 current_offset = std::max(current_offset, alloc.offset + alloc.size);
83 }
84 if (best_offset == kOffsetNotAssigned) {
85 best_offset = AlignTo(alignment, current_offset);
86 }
87
88 // Update the required buffer size.
89 high_water_mark_ = std::max(high_water_mark_, best_offset + size);
90 new_alloc->offset = best_offset;
91
92 auto insertion_it = std::upper_bound(ordered_allocs_.begin(),
93 ordered_allocs_.end(), *new_alloc);
94 ordered_allocs_.insert(insertion_it, *new_alloc);
95 return kTfLiteOk;
96 }
97
Deallocate(TfLiteContext * context,const ArenaAllocWithUsageInterval & alloc)98 TfLiteStatus SimpleMemoryArena::Deallocate(
99 TfLiteContext* context, const ArenaAllocWithUsageInterval& alloc) {
100 if (alloc.size == 0) {
101 return kTfLiteOk;
102 }
103
104 int erased_allocs_count = 0;
105 auto it = ordered_allocs_.begin();
106 while (it != ordered_allocs_.end()) {
107 if (it->tensor == alloc.tensor) {
108 erased_allocs_count++;
109 it = ordered_allocs_.erase(it);
110 } else {
111 ++it;
112 }
113 }
114 TF_LITE_ENSURE(context, erased_allocs_count <= 1);
115 return kTfLiteOk;
116 }
117
Commit(TfLiteContext * context)118 TfLiteStatus SimpleMemoryArena::Commit(TfLiteContext* context) {
119 size_t required_size = RequiredBufferSize();
120 if (required_size > underlying_buffer_size_) {
121 #ifdef TF_LITE_TENSORFLOW_PROFILER
122 OnTfLiteArenaAlloc(subgraph_index_, reinterpret_cast<std::uintptr_t>(this),
123 required_size);
124 #endif
125 char* new_alloc = new char[required_size];
126 char* new_underlying_buffer_aligned_ptr = reinterpret_cast<char*>(
127 AlignTo(arena_alignment_, reinterpret_cast<intptr_t>(new_alloc)));
128
129 // If the arena had been previously allocated, copy over the old memory.
130 // Since Alloc pointers are offset based, they will remain valid in the new
131 // memory block.
132 if (high_water_mark_ > 0 && underlying_buffer_size_ > 0) {
133 size_t copy_amount = std::min(
134 underlying_buffer_.get() + underlying_buffer_size_ -
135 underlying_buffer_aligned_ptr_,
136 new_alloc + required_size - new_underlying_buffer_aligned_ptr);
137 memcpy(new_underlying_buffer_aligned_ptr, underlying_buffer_aligned_ptr_,
138 copy_amount);
139 }
140
141 #ifdef TF_LITE_TENSORFLOW_PROFILER
142 if (underlying_buffer_size_ > 0) {
143 OnTfLiteArenaDealloc(subgraph_index_,
144 reinterpret_cast<std::uintptr_t>(this),
145 underlying_buffer_size_);
146 }
147 #endif
148 underlying_buffer_.reset(new_alloc);
149 underlying_buffer_size_ = required_size;
150 underlying_buffer_aligned_ptr_ = new_underlying_buffer_aligned_ptr;
151 }
152 committed_ = true;
153 return underlying_buffer_ != nullptr ? kTfLiteOk : kTfLiteError;
154 }
155
ResolveAlloc(TfLiteContext * context,const ArenaAllocWithUsageInterval & alloc,char ** output_ptr)156 TfLiteStatus SimpleMemoryArena::ResolveAlloc(
157 TfLiteContext* context, const ArenaAllocWithUsageInterval& alloc,
158 char** output_ptr) {
159 TF_LITE_ENSURE(context, committed_);
160 TF_LITE_ENSURE(context, output_ptr != nullptr);
161 TF_LITE_ENSURE(context,
162 underlying_buffer_size_ >= (alloc.offset + alloc.size));
163 if (alloc.size == 0) {
164 *output_ptr = nullptr;
165 } else {
166 *output_ptr = underlying_buffer_aligned_ptr_ + alloc.offset;
167 }
168 return kTfLiteOk;
169 }
170
ClearPlan()171 TfLiteStatus SimpleMemoryArena::ClearPlan() {
172 committed_ = false;
173 high_water_mark_ = 0;
174 ordered_allocs_.clear();
175 return kTfLiteOk;
176 }
177
ReleaseBuffer()178 TfLiteStatus SimpleMemoryArena::ReleaseBuffer() {
179 committed_ = false;
180 #ifdef TF_LITE_TENSORFLOW_PROFILER
181 OnTfLiteArenaDealloc(subgraph_index_, reinterpret_cast<std::uintptr_t>(this),
182 underlying_buffer_size_);
183 #endif
184 underlying_buffer_size_ = 0;
185 underlying_buffer_aligned_ptr_ = nullptr;
186 underlying_buffer_.reset();
187 return kTfLiteOk;
188 }
189
190 // Using weak symbols to create a pluggable debugging module.
DumpArenaInfo(const std::string & name,const std::vector<int> & execution_plan,size_t arena_size,const std::vector<ArenaAllocWithUsageInterval> & allocs)191 TFLITE_ATTRIBUTE_WEAK void DumpArenaInfo(
192 const std::string& name, const std::vector<int>& execution_plan,
193 size_t arena_size, const std::vector<ArenaAllocWithUsageInterval>& allocs) {
194 }
195
DumpDebugInfo(const std::string & name,const std::vector<int> & execution_plan) const196 void SimpleMemoryArena::DumpDebugInfo(
197 const std::string& name, const std::vector<int>& execution_plan) const {
198 tflite::DumpArenaInfo(name, execution_plan, underlying_buffer_size_,
199 ordered_allocs_);
200 }
201
202 } // namespace tflite
203