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 #include "tensorflow/contrib/lite/arena_planner.h"
16
17 namespace tflite {
18
19 namespace {
20
21 // Memory allocation tuning
22 constexpr const int kDefaultArenaAlignment = 64;
23 constexpr const int kDefaultTensorAlignment = 4;
24
25 } // namespace
26
27 struct AllocationInfo {
28 // The node index requesting this allocation.
29 int node;
30 // The tensor index to be allocated or deallocated.
31 int tensor;
32 // Whether to allocate or deallocate
33 enum { ALLOC, DEALLOC } type;
34 };
35
ArenaPlanner(TfLiteContext * context,std::unique_ptr<GraphInfo> graph_info)36 ArenaPlanner::ArenaPlanner(TfLiteContext* context,
37 std::unique_ptr<GraphInfo> graph_info)
38 : context_(context),
39 graph_info_(std::move(graph_info)),
40 arena_(kDefaultArenaAlignment),
41 persistent_arena_(kDefaultArenaAlignment) {}
42
~ArenaPlanner()43 ArenaPlanner::~ArenaPlanner() {}
44
BasePointer(TfLiteAllocationType type)45 int64_t ArenaPlanner::BasePointer(TfLiteAllocationType type) {
46 if (type == kTfLiteArenaRwPersistent) {
47 return persistent_arena_.BasePointer();
48 }
49 if (type == kTfLiteArenaRw) {
50 return arena_.BasePointer();
51 }
52 return 0;
53 }
54
ResetAllocations()55 TfLiteStatus ArenaPlanner::ResetAllocations() {
56 TF_LITE_ENSURE_STATUS(arena_.Clear());
57 TF_LITE_ENSURE_STATUS(persistent_arena_.Clear());
58 allocs_.clear();
59 allocs_.resize(graph_info_->num_tensors());
60 return kTfLiteOk;
61 }
62
PlanAllocations()63 TfLiteStatus ArenaPlanner::PlanAllocations() {
64 // Invalidate any existing data.
65 TF_LITE_ENSURE_STATUS(ResetAllocations());
66
67 // Keeps track of references to each tensor.
68 std::vector<int> refcounts(graph_info_->num_tensors(), 0);
69
70 // There will be an entry in alloc_queue_ for the allocation of each tensor
71 // and another for their deallocation.
72 alloc_queue_.reserve(2 * graph_info_->num_tensors());
73
74 // We must make sure the output tensors are never overwritten. We do that by
75 // artificially adding one to their ref-counts so they are never selected
76 // for deallocation.
77 for (int tensor_index : graph_info_->outputs()) {
78 refcounts[tensor_index]++;
79 }
80
81 // Count references to node input tensors.
82 for (int i = 0; i < graph_info_->num_nodes(); ++i) {
83 const TfLiteNode& node = graph_info_->node(i);
84 TfLiteIntArray* node_inputs = node.inputs;
85 for (int j = 0; j < node_inputs->size; ++j) {
86 int tensor_index = node_inputs->data[j];
87 if (tensor_index != kOptionalTensor) {
88 refcounts[tensor_index]++;
89 }
90 }
91 }
92
93 // Queue all graph inputs for allocation.
94 for (int tensor_index : graph_info_->inputs()) {
95 if (tensor_index != kOptionalTensor) {
96 alloc_queue_.push_back({0, tensor_index, AllocationInfo::ALLOC});
97 }
98 }
99
100 // Go through the graph in execution order.
101 for (int i = 0; i < graph_info_->num_nodes(); ++i) {
102 const TfLiteNode& node = graph_info_->node(i);
103
104 // First queue output tensors for allocation.
105 TfLiteIntArray* node_outputs = node.outputs;
106 for (int j = 0; j < node_outputs->size; ++j) {
107 int tensor_index = node_outputs->data[j];
108 alloc_queue_.push_back({i, tensor_index, AllocationInfo::ALLOC});
109 }
110
111 // Then update the ref-counts of the node's inputs, and if necessary queue
112 // them for deallocation.
113 TfLiteIntArray* node_inputs = node.inputs;
114 for (int j = 0; j < node_inputs->size; ++j) {
115 int tensor_index = node_inputs->data[j];
116 if (tensor_index != kOptionalTensor) {
117 refcounts[tensor_index]--;
118 if (refcounts[tensor_index] == 0) {
119 alloc_queue_.push_back({i, tensor_index, AllocationInfo::DEALLOC});
120 }
121 }
122 }
123 }
124
125 // Note that graph outputs will never be scheduled for deallocation. We
126 // could do that here for completeness, but it won't have any effect.
127 return kTfLiteOk;
128 }
129
ExecuteAllocations(int first_node,int last_node)130 TfLiteStatus ArenaPlanner::ExecuteAllocations(int first_node, int last_node) {
131 TF_LITE_ENSURE_STATUS(CalculateAllocations(first_node, last_node));
132 TF_LITE_ENSURE_STATUS(Commit());
133
134 for (int i = 0; i < graph_info_->num_tensors(); ++i) {
135 // TODO(ahentz): we could do this only for the tensors that were modified
136 // in CalculateAllocations(), instead of redoing it for tensors that
137 // already had proper pointers. However we must be very careful, because
138 // SimpleMemoryArena::Commit() could move the base pointer.
139 TF_LITE_ENSURE_STATUS(ResolveTensorAllocation(i));
140 }
141
142 return kTfLiteOk;
143 }
144
Commit()145 TfLiteStatus ArenaPlanner::Commit() {
146 TF_LITE_ENSURE_STATUS(arena_.Commit(context_));
147 TF_LITE_ENSURE_STATUS(persistent_arena_.Commit(context_));
148 return kTfLiteOk;
149 }
150
CalculateAllocations(int first_node,int last_node)151 TfLiteStatus ArenaPlanner::CalculateAllocations(int first_node, int last_node) {
152 int active_node = first_node;
153 // When dynamic tensors are present this method is called multiple times.
154 // The items in the alloc_queue_ referring to nodes before first_node were
155 // processed previously and should be skipped. Entries after last_node are
156 // not yet ready to be handled.
157 for (const auto& alloc_info : alloc_queue_) {
158 if (alloc_info.node < first_node) continue;
159 if (alloc_info.node > last_node) break;
160 if (alloc_info.node == active_node) {
161 // This is the first allocation/deallocation for a given node. It is
162 // time to deallocate the previous temporaries and allocate new ones.
163 if (active_node != first_node) {
164 TF_LITE_ENSURE_STATUS(
165 CalculateDeallocationOfInternalTensors(active_node - 1));
166 }
167 TF_LITE_ENSURE_STATUS(CalculateAllocationOfInternalTensors(active_node));
168 ++active_node;
169 }
170 // Handle the current item.
171 if (alloc_info.type == AllocationInfo::ALLOC) {
172 TF_LITE_ENSURE_STATUS(CalculateTensorAllocation(alloc_info.tensor));
173 } else {
174 TF_LITE_ENSURE_STATUS(CalculateTensorDeallocation(alloc_info.tensor));
175 }
176 }
177
178 // Don't forget to deallocate temporaries of last node.
179 TF_LITE_ENSURE_STATUS(
180 CalculateDeallocationOfInternalTensors(active_node - 1));
181
182 return kTfLiteOk;
183 }
184
ResolveTensorAllocation(int tensor_index)185 TfLiteStatus ArenaPlanner::ResolveTensorAllocation(int tensor_index) {
186 TfLiteTensor& tensor = *graph_info_->tensor(tensor_index);
187 if (tensor.allocation_type == kTfLiteArenaRw) {
188 // Skip resolution if the size of the tensor is zero, leaving it as a
189 // nullptr.
190 if (allocs_[tensor_index].size != 0) {
191 TF_LITE_ENSURE_STATUS(arena_.ResolveAlloc(context_, allocs_[tensor_index],
192 &tensor.data.raw));
193 }
194 }
195 if (tensor.allocation_type == kTfLiteArenaRwPersistent) {
196 TF_LITE_ENSURE_STATUS(persistent_arena_.ResolveAlloc(
197 context_, allocs_[tensor_index], &tensor.data.raw));
198 }
199 return kTfLiteOk;
200 }
201
CalculateTensorAllocation(int tensor_index)202 TfLiteStatus ArenaPlanner::CalculateTensorAllocation(int tensor_index) {
203 TfLiteTensor& tensor = *graph_info_->tensor(tensor_index);
204 if (tensor.allocation_type == kTfLiteArenaRw) {
205 TF_LITE_ENSURE_STATUS(arena_.Allocate(context_, kDefaultTensorAlignment,
206 tensor.bytes,
207 &allocs_[tensor_index]));
208 }
209 if (tensor.allocation_type == kTfLiteArenaRwPersistent) {
210 TF_LITE_ENSURE_STATUS(
211 persistent_arena_.Allocate(context_, kDefaultTensorAlignment,
212 tensor.bytes, &allocs_[tensor_index]));
213 }
214 return kTfLiteOk;
215 }
216
CalculateTensorDeallocation(int tensor_index)217 TfLiteStatus ArenaPlanner::CalculateTensorDeallocation(int tensor_index) {
218 TfLiteTensor& tensor = *graph_info_->tensor(tensor_index);
219 if (tensor.allocation_type == kTfLiteArenaRw) {
220 TF_LITE_ENSURE_STATUS(arena_.Deallocate(context_, allocs_[tensor_index]));
221 }
222 return kTfLiteOk;
223 }
224
CalculateAllocationOfInternalTensors(int node_index)225 TfLiteStatus ArenaPlanner::CalculateAllocationOfInternalTensors(
226 int node_index) {
227 if (node_index < graph_info_->num_nodes()) {
228 const TfLiteNode& node = graph_info_->node(node_index);
229 TfLiteIntArray* node_temporaries = node.temporaries;
230 for (int i = 0; i < node_temporaries->size; ++i) {
231 int tensor_index = node_temporaries->data[i];
232 TF_LITE_ENSURE_STATUS(CalculateTensorAllocation(tensor_index));
233 }
234 }
235 return kTfLiteOk;
236 }
237
CalculateDeallocationOfInternalTensors(int node_index)238 TfLiteStatus ArenaPlanner::CalculateDeallocationOfInternalTensors(
239 int node_index) {
240 if (node_index < graph_info_->num_nodes()) {
241 const TfLiteNode& node = graph_info_->node(node_index);
242 TfLiteIntArray* node_temporaries = node.temporaries;
243 for (int i = 0; i < node_temporaries->size; ++i) {
244 int tensor_index = node_temporaries->data[i];
245 TF_LITE_ENSURE_STATUS(CalculateTensorDeallocation(tensor_index));
246 }
247 }
248 return kTfLiteOk;
249 }
250
251 } // namespace tflite
252