• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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