1 /* Copyright 2020 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/lite/micro/micro_interpreter.h"
16
17 #include <cstdarg>
18 #include <cstddef>
19 #include <cstdint>
20
21 #include "flatbuffers/flatbuffers.h" // from @flatbuffers
22 #include "tensorflow/lite/c/common.h"
23 #include "tensorflow/lite/core/api/error_reporter.h"
24 #include "tensorflow/lite/core/api/tensor_utils.h"
25 #include "tensorflow/lite/micro/memory_helpers.h"
26 #include "tensorflow/lite/micro/micro_allocator.h"
27 #include "tensorflow/lite/micro/micro_error_reporter.h"
28 #include "tensorflow/lite/micro/micro_op_resolver.h"
29 #include "tensorflow/lite/micro/micro_profiler.h"
30 #include "tensorflow/lite/schema/schema_generated.h"
31
32 namespace tflite {
33 namespace {
34
35 #ifndef TF_LITE_STRIP_ERROR_STRINGS
OpNameFromRegistration(const TfLiteRegistration * registration)36 const char* OpNameFromRegistration(const TfLiteRegistration* registration) {
37 if (registration->builtin_code == BuiltinOperator_CUSTOM) {
38 return registration->custom_name;
39 } else {
40 return EnumNameBuiltinOperator(BuiltinOperator(registration->builtin_code));
41 }
42 }
43 #endif // !defined(TF_LITE_STRIP_ERROR_STRINGS)
44
45 } // namespace
46
47 namespace internal {
48
ContextHelper(ErrorReporter * error_reporter,MicroAllocator * allocator,const Model * model)49 ContextHelper::ContextHelper(ErrorReporter* error_reporter,
50 MicroAllocator* allocator, const Model* model)
51 : allocator_(allocator), error_reporter_(error_reporter), model_(model) {}
52
AllocatePersistentBuffer(TfLiteContext * ctx,size_t bytes)53 void* ContextHelper::AllocatePersistentBuffer(TfLiteContext* ctx,
54 size_t bytes) {
55 return reinterpret_cast<ContextHelper*>(ctx->impl_)
56 ->allocator_->AllocatePersistentBuffer(bytes);
57 }
58
RequestScratchBufferInArena(TfLiteContext * ctx,size_t bytes,int * buffer_idx)59 TfLiteStatus ContextHelper::RequestScratchBufferInArena(TfLiteContext* ctx,
60 size_t bytes,
61 int* buffer_idx) {
62 ContextHelper* helper = reinterpret_cast<ContextHelper*>(ctx->impl_);
63 return helper->allocator_->RequestScratchBufferInArena(bytes, buffer_idx);
64 }
65
GetScratchBuffer(TfLiteContext * ctx,int buffer_idx)66 void* ContextHelper::GetScratchBuffer(TfLiteContext* ctx, int buffer_idx) {
67 ContextHelper* helper = reinterpret_cast<ContextHelper*>(ctx->impl_);
68 ScratchBufferHandle* handle = helper->scratch_buffer_handles_ + buffer_idx;
69 return handle->data;
70 }
71
ReportOpError(struct TfLiteContext * context,const char * format,...)72 void ContextHelper::ReportOpError(struct TfLiteContext* context,
73 const char* format, ...) {
74 #ifndef TF_LITE_STRIP_ERROR_STRINGS
75 ContextHelper* helper = static_cast<ContextHelper*>(context->impl_);
76 va_list args;
77 va_start(args, format);
78 TF_LITE_REPORT_ERROR(helper->error_reporter_, format, args);
79 va_end(args);
80 #endif
81 }
82
GetTensor(const struct TfLiteContext * context,int tensor_idx)83 TfLiteTensor* ContextHelper::GetTensor(const struct TfLiteContext* context,
84 int tensor_idx) {
85 ContextHelper* helper = static_cast<ContextHelper*>(context->impl_);
86 return helper->allocator_->AllocateTempTfLiteTensor(
87 helper->model_, helper->eval_tensors_, tensor_idx);
88 }
89
GetEvalTensor(const struct TfLiteContext * context,int tensor_idx)90 TfLiteEvalTensor* ContextHelper::GetEvalTensor(
91 const struct TfLiteContext* context, int tensor_idx) {
92 ContextHelper* helper = reinterpret_cast<ContextHelper*>(context->impl_);
93 return &helper->eval_tensors_[tensor_idx];
94 }
95
SetTfLiteEvalTensors(TfLiteEvalTensor * eval_tensors)96 void ContextHelper::SetTfLiteEvalTensors(TfLiteEvalTensor* eval_tensors) {
97 eval_tensors_ = eval_tensors;
98 }
99
SetScratchBufferHandles(ScratchBufferHandle * scratch_buffer_handles)100 void ContextHelper::SetScratchBufferHandles(
101 ScratchBufferHandle* scratch_buffer_handles) {
102 scratch_buffer_handles_ = scratch_buffer_handles;
103 }
104
105 } // namespace internal
106
MicroInterpreter(const Model * model,const MicroOpResolver & op_resolver,uint8_t * tensor_arena,size_t tensor_arena_size,ErrorReporter * error_reporter,MicroProfiler * profiler)107 MicroInterpreter::MicroInterpreter(const Model* model,
108 const MicroOpResolver& op_resolver,
109 uint8_t* tensor_arena,
110 size_t tensor_arena_size,
111 ErrorReporter* error_reporter,
112 MicroProfiler* profiler)
113 : model_(model),
114 op_resolver_(op_resolver),
115 error_reporter_(error_reporter),
116 allocator_(*MicroAllocator::Create(tensor_arena, tensor_arena_size,
117 error_reporter)),
118 tensors_allocated_(false),
119 initialization_status_(kTfLiteError),
120 eval_tensors_(nullptr),
121 context_helper_(error_reporter_, &allocator_, model),
122 input_tensors_(nullptr),
123 output_tensors_(nullptr) {
124 Init(profiler);
125 }
126
MicroInterpreter(const Model * model,const MicroOpResolver & op_resolver,MicroAllocator * allocator,ErrorReporter * error_reporter,MicroProfiler * profiler)127 MicroInterpreter::MicroInterpreter(const Model* model,
128 const MicroOpResolver& op_resolver,
129 MicroAllocator* allocator,
130 ErrorReporter* error_reporter,
131 MicroProfiler* profiler)
132 : model_(model),
133 op_resolver_(op_resolver),
134 error_reporter_(error_reporter),
135 allocator_(*allocator),
136 tensors_allocated_(false),
137 initialization_status_(kTfLiteError),
138 eval_tensors_(nullptr),
139 context_helper_(error_reporter_, &allocator_, model),
140 input_tensors_(nullptr),
141 output_tensors_(nullptr) {
142 Init(profiler);
143 }
144
~MicroInterpreter()145 MicroInterpreter::~MicroInterpreter() {
146 if (node_and_registrations_ != nullptr) {
147 for (size_t i = 0; i < subgraph_->operators()->size(); ++i) {
148 TfLiteNode* node = &(node_and_registrations_[i].node);
149 const TfLiteRegistration* registration =
150 node_and_registrations_[i].registration;
151 // registration is allocated outside the interpreter, so double check to
152 // make sure it's not nullptr;
153 if (registration != nullptr && registration->free != nullptr) {
154 registration->free(&context_, node->user_data);
155 }
156 }
157 }
158 }
159
Init(MicroProfiler * profiler)160 void MicroInterpreter::Init(MicroProfiler* profiler) {
161 const flatbuffers::Vector<flatbuffers::Offset<SubGraph>>* subgraphs =
162 model_->subgraphs();
163 if (subgraphs->size() != 1) {
164 TF_LITE_REPORT_ERROR(error_reporter_,
165 "Only 1 subgraph is currently supported.\n");
166 initialization_status_ = kTfLiteError;
167 return;
168 }
169 subgraph_ = (*subgraphs)[0];
170
171 context_.impl_ = static_cast<void*>(&context_helper_);
172 context_.ReportError = context_helper_.ReportOpError;
173 context_.GetTensor = context_helper_.GetTensor;
174 context_.GetEvalTensor = context_helper_.GetEvalTensor;
175 context_.recommended_num_threads = 1;
176 context_.profiler = profiler;
177
178 initialization_status_ = kTfLiteOk;
179 }
180
AllocateTensors()181 TfLiteStatus MicroInterpreter::AllocateTensors() {
182 if (allocator_.StartModelAllocation(model_, op_resolver_,
183 &node_and_registrations_,
184 &eval_tensors_) != kTfLiteOk) {
185 TF_LITE_REPORT_ERROR(error_reporter_,
186 "Failed starting model allocation.\n");
187 initialization_status_ = kTfLiteError;
188 return kTfLiteError;
189 }
190
191 // Update the pointer now that TfLiteEvalTensor allocation has completed on
192 // the context helper.
193 // TODO(b/16157777): This call would not be needed if ContextHelper rolled
194 // into the interpreter.
195 context_helper_.SetTfLiteEvalTensors(eval_tensors_);
196 context_.tensors_size = subgraph_->tensors()->size();
197
198 // Only allow AllocatePersistentBuffer in Init stage.
199 context_.AllocatePersistentBuffer = context_helper_.AllocatePersistentBuffer;
200 context_.RequestScratchBufferInArena = nullptr;
201 context_.GetScratchBuffer = nullptr;
202
203 for (size_t i = 0; i < subgraph_->operators()->size(); ++i) {
204 auto* node = &(node_and_registrations_[i].node);
205 auto* registration = node_and_registrations_[i].registration;
206 size_t init_data_size;
207 const char* init_data;
208 if (registration->builtin_code == BuiltinOperator_CUSTOM) {
209 init_data = reinterpret_cast<const char*>(node->custom_initial_data);
210 init_data_size = node->custom_initial_data_size;
211 } else {
212 init_data = reinterpret_cast<const char*>(node->builtin_data);
213 init_data_size = 0;
214 }
215 if (registration->init) {
216 node->user_data =
217 registration->init(&context_, init_data, init_data_size);
218 }
219 }
220
221 // Both AllocatePersistentBuffer and RequestScratchBufferInArena is
222 // available in Prepare stage.
223 context_.RequestScratchBufferInArena =
224 context_helper_.RequestScratchBufferInArena;
225 for (size_t i = 0; i < subgraph_->operators()->size(); ++i) {
226 auto* node = &(node_and_registrations_[i].node);
227 auto* registration = node_and_registrations_[i].registration;
228 if (registration->prepare) {
229 TfLiteStatus prepare_status = registration->prepare(&context_, node);
230 if (prepare_status != kTfLiteOk) {
231 TF_LITE_REPORT_ERROR(
232 error_reporter_,
233 "Node %s (number %df) failed to prepare with status %d",
234 OpNameFromRegistration(registration), i, prepare_status);
235 return kTfLiteError;
236 }
237 }
238 allocator_.FinishPrepareNodeAllocations(/*node_id=*/i);
239 }
240
241 // Prepare is done, we're ready for Invoke. Memory allocation is no longer
242 // allowed. Kernels can only fetch scratch buffers via GetScratchBuffer.
243 context_.AllocatePersistentBuffer = nullptr;
244 context_.RequestScratchBufferInArena = nullptr;
245 context_.GetScratchBuffer = context_helper_.GetScratchBuffer;
246
247 TF_LITE_ENSURE_OK(&context_,
248 allocator_.FinishModelAllocation(model_, eval_tensors_,
249 &scratch_buffer_handles_));
250 // TODO(b/16157777): Remove this when ContextHelper is rolled into this class.
251 context_helper_.SetScratchBufferHandles(scratch_buffer_handles_);
252
253 // TODO(b/162311891): Drop these allocations when the interpreter supports
254 // handling buffers from TfLiteEvalTensor.
255 input_tensors_ =
256 reinterpret_cast<TfLiteTensor**>(allocator_.AllocatePersistentBuffer(
257 sizeof(TfLiteTensor*) * inputs_size()));
258 if (input_tensors_ == nullptr) {
259 TF_LITE_REPORT_ERROR(
260 error_reporter_,
261 "Failed to allocate memory for context->input_tensors_, "
262 "%d bytes required",
263 sizeof(TfLiteTensor*) * inputs_size());
264 return kTfLiteError;
265 }
266
267 for (size_t i = 0; i < inputs_size(); ++i) {
268 input_tensors_[i] = allocator_.AllocatePersistentTfLiteTensor(
269 model_, eval_tensors_, inputs().Get(i));
270 if (input_tensors_[i] == nullptr) {
271 TF_LITE_REPORT_ERROR(error_reporter_,
272 "Failed to initialize input tensor %d", i);
273 return kTfLiteError;
274 }
275 }
276
277 // TODO(b/162311891): Drop these allocations when the interpreter supports
278 // handling buffers from TfLiteEvalTensor.
279 output_tensors_ =
280 reinterpret_cast<TfLiteTensor**>(allocator_.AllocatePersistentBuffer(
281 sizeof(TfLiteTensor*) * outputs_size()));
282 if (output_tensors_ == nullptr) {
283 TF_LITE_REPORT_ERROR(
284 error_reporter_,
285 "Failed to allocate memory for context->output_tensors_, "
286 "%d bytes required",
287 sizeof(TfLiteTensor*) * outputs_size());
288 return kTfLiteError;
289 }
290
291 for (size_t i = 0; i < outputs_size(); ++i) {
292 output_tensors_[i] = allocator_.AllocatePersistentTfLiteTensor(
293 model_, eval_tensors_, outputs().Get(i));
294 if (output_tensors_[i] == nullptr) {
295 TF_LITE_REPORT_ERROR(error_reporter_,
296 "Failed to initialize output tensor %d", i);
297 return kTfLiteError;
298 }
299 }
300
301 TF_LITE_ENSURE_STATUS(ResetVariableTensors());
302
303 tensors_allocated_ = true;
304 return kTfLiteOk;
305 }
306
Invoke()307 TfLiteStatus MicroInterpreter::Invoke() {
308 if (initialization_status_ != kTfLiteOk) {
309 TF_LITE_REPORT_ERROR(error_reporter_,
310 "Invoke() called after initialization failed\n");
311 return kTfLiteError;
312 }
313
314 // Ensure tensors are allocated before the interpreter is invoked to avoid
315 // difficult to debug segfaults.
316 if (!tensors_allocated_) {
317 TF_LITE_ENSURE_OK(&context_, AllocateTensors());
318 }
319
320 for (size_t i = 0; i < subgraph_->operators()->size(); ++i) {
321 auto* node = &(node_and_registrations_[i].node);
322 auto* registration = node_and_registrations_[i].registration;
323
324 // This ifdef is needed (even though ScopedMicroProfiler itself is a no-op with
325 // -DTF_LITE_STRIP_ERROR_STRINGS) because the function OpNameFromRegistration is
326 // only defined for builds with the error strings.
327 #if !defined(TF_LITE_STRIP_ERROR_STRINGS)
328 ScopedMicroProfiler scoped_profiler(
329 OpNameFromRegistration(registration),
330 reinterpret_cast<MicroProfiler*>(context_.profiler));
331 #endif
332
333 TFLITE_DCHECK(registration->invoke);
334 TfLiteStatus invoke_status = registration->invoke(&context_, node);
335
336 // All TfLiteTensor structs used in the kernel are allocated from temp
337 // memory in the allocator. This creates a chain of allocations in the
338 // temp section. The call below resets the chain of allocations to
339 // prepare for the next call.
340 allocator_.ResetTempAllocations();
341
342 if (invoke_status == kTfLiteError) {
343 TF_LITE_REPORT_ERROR(
344 error_reporter_,
345 "Node %s (number %d) failed to invoke with status %d",
346 OpNameFromRegistration(registration), i, invoke_status);
347 return kTfLiteError;
348 } else if (invoke_status != kTfLiteOk) {
349 return invoke_status;
350 }
351 }
352
353 return kTfLiteOk;
354 }
355
input(size_t index)356 TfLiteTensor* MicroInterpreter::input(size_t index) {
357 const size_t length = inputs_size();
358 if (index >= length) {
359 TF_LITE_REPORT_ERROR(error_reporter_,
360 "Input index %d out of range (length is %d)", index,
361 length);
362 return nullptr;
363 }
364 return input_tensors_[index];
365 }
366
output(size_t index)367 TfLiteTensor* MicroInterpreter::output(size_t index) {
368 const size_t length = outputs_size();
369 if (index >= length) {
370 TF_LITE_REPORT_ERROR(error_reporter_,
371 "Output index %d out of range (length is %d)", index,
372 length);
373 return nullptr;
374 }
375 return output_tensors_[index];
376 }
377
tensor(size_t index)378 TfLiteTensor* MicroInterpreter::tensor(size_t index) {
379 const size_t length = tensors_size();
380 if (index >= length) {
381 TF_LITE_REPORT_ERROR(error_reporter_,
382 "Tensor index %d out of range (length is %d)", index,
383 length);
384 return nullptr;
385 }
386 return allocator_.AllocatePersistentTfLiteTensor(model_, eval_tensors_,
387 index);
388 }
389
ResetVariableTensors()390 TfLiteStatus MicroInterpreter::ResetVariableTensors() {
391 for (size_t i = 0; i < subgraph_->tensors()->size(); ++i) {
392 auto* tensor = subgraph_->tensors()->Get(i);
393 if (tensor->is_variable()) {
394 size_t buffer_size;
395 TF_LITE_ENSURE_STATUS(
396 TfLiteEvalTensorByteLength(&eval_tensors_[i], &buffer_size));
397
398 int value = 0;
399 if (tensor->type() == tflite::TensorType_INT8) {
400 value = tensor->quantization()->zero_point()->Get(0);
401 }
402 memset(eval_tensors_[i].data.raw, value, buffer_size);
403 }
404 }
405
406 return kTfLiteOk;
407 }
408
409 } // namespace tflite
410