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/interpreter.h"
17
18 #include <stddef.h>
19 #include <stdint.h>
20 #include <stdlib.h>
21 #include <string.h>
22
23 #include <map>
24 #include <memory>
25 #include <new>
26 #include <string>
27 #include <utility>
28 #include <vector>
29
30 #include <gmock/gmock.h>
31 #include <gtest/gtest.h>
32 #include "third_party/eigen3/Eigen/Core"
33 #include "tensorflow/lite/c/builtin_op_data.h"
34 #include "tensorflow/lite/external_cpu_backend_context.h"
35 #include "tensorflow/lite/interpreter_test_util.h"
36 #include "tensorflow/lite/kernels/builtin_op_kernels.h"
37 #include "tensorflow/lite/kernels/internal/compatibility.h"
38 #include "tensorflow/lite/kernels/kernel_util.h"
39 #include "tensorflow/lite/string_type.h"
40 #include "tensorflow/lite/string_util.h"
41 #include "tensorflow/lite/testing/util.h"
42 #include "tensorflow/lite/util.h"
43
44 namespace tflite {
45
46 namespace ops {
47 namespace builtin {
48 TfLiteRegistration* Register_PADV2();
49 TfLiteRegistration* Register_NEG();
50 } // namespace builtin
51 } // namespace ops
52
53 namespace {
54
55 using ::testing::IsEmpty;
56
57 // Make an interpreter that has no tensors and no nodes
TEST(BasicInterpreter,ZeroInterpreter)58 TEST(BasicInterpreter, ZeroInterpreter) {
59 testing::internal::CaptureStderr();
60
61 Interpreter interpreter;
62
63 #ifndef NDEBUG
64 const char* kExpectedLog = "INFO: Initialized TensorFlow Lite runtime";
65 #else
66 const char* kExpectedLog = "";
67 #endif
68 EXPECT_THAT(testing::internal::GetCapturedStderr(),
69 testing::HasSubstr(kExpectedLog));
70
71 interpreter.SetInputs({});
72 interpreter.SetOutputs({});
73 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
74 ASSERT_EQ(interpreter.Invoke(), kTfLiteOk);
75
76 // Creating a new interpreter should not redundantly log runtime init.
77 testing::internal::CaptureStderr();
78 Interpreter interpreter2;
79 EXPECT_THAT(testing::internal::GetCapturedStderr(), IsEmpty());
80 }
81
82 // Test various error conditions.
TEST(BasicInterpreter,InvokeInvalidModel)83 TEST(BasicInterpreter, InvokeInvalidModel) {
84 Interpreter interpreter;
85 ASSERT_NE(interpreter.Invoke(), kTfLiteOk);
86 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
87 ASSERT_EQ(interpreter.Invoke(), kTfLiteOk);
88 }
89
TEST(BasicInterpreter,TestAllocateTensorsResetVariableTensorsFloatAndHyrbid)90 TEST(BasicInterpreter, TestAllocateTensorsResetVariableTensorsFloatAndHyrbid) {
91 Interpreter interpreter;
92 int tensor_index;
93 ASSERT_EQ(interpreter.AddTensors(1, &tensor_index), kTfLiteOk);
94 constexpr int kTensorSize = 16;
95 TfLiteQuantizationParams quant;
96 interpreter.SetTensorParametersReadWrite(tensor_index, kTfLiteFloat32, "",
97 {kTensorSize}, quant,
98 /*is_variable=*/true);
99 interpreter.SetVariables({tensor_index});
100 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
101 TfLiteTensor* tensor = interpreter.tensor(tensor_index);
102 // Ensure that variable tensors are reset to zero.
103 for (int i = 0; i < kTensorSize; ++i) {
104 ASSERT_EQ(tensor->data.f[i], 0.0f);
105 }
106 }
107
TEST(BasicInterpreter,TestAllocateTensorsResetVariableTensorsInt8)108 TEST(BasicInterpreter, TestAllocateTensorsResetVariableTensorsInt8) {
109 Interpreter interpreter;
110 int tensor_index;
111 ASSERT_EQ(interpreter.AddTensors(1, &tensor_index), kTfLiteOk);
112 constexpr int kTensorSize = 16;
113 TfLiteQuantizationParams quant;
114 quant.scale = 0.15;
115 quant.zero_point = -3;
116 interpreter.SetTensorParametersReadWrite(tensor_index, kTfLiteInt8, "",
117 {kTensorSize}, quant,
118 /*is_variable=*/true);
119 interpreter.SetVariables({tensor_index});
120 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
121 TfLiteTensor* tensor = interpreter.tensor(tensor_index);
122 // Ensure that variable tensors are reset to zero point.
123 for (int i = 0; i < kTensorSize; ++i) {
124 ASSERT_EQ(tensor->data.int8[i], -3);
125 }
126 }
127
128 // Test size accessor functions.
TEST(BasicInterpreter,TestSizeFunctions)129 TEST(BasicInterpreter, TestSizeFunctions) {
130 Interpreter interpreter;
131 int base_index;
132 ASSERT_EQ(interpreter.nodes_size(), 0);
133 ASSERT_EQ(interpreter.tensors_size(), 0);
134 ASSERT_EQ(interpreter.AddTensors(2, &base_index), kTfLiteOk);
135 ASSERT_EQ(interpreter.tensors_size(), 2);
136 ASSERT_EQ(base_index, 0);
137 ASSERT_EQ(interpreter.AddTensors(3, &base_index), kTfLiteOk);
138 ASSERT_EQ(interpreter.tensors_size(), 5);
139 ASSERT_EQ(interpreter.AddTensors(1), kTfLiteOk);
140 ASSERT_EQ(interpreter.tensors_size(), 6);
141 ASSERT_EQ(base_index, 2);
142 }
143
144 // Test if invalid indices make a model inconsistent (and conversely if
145 // valid indices keep a model consistent).
TEST(BasicInterpreter,InconsistentModel)146 TEST(BasicInterpreter, InconsistentModel) {
147 // Invalid inputs
148 {
149 Interpreter interpreter;
150 ASSERT_NE(interpreter.SetInputs({5}), kTfLiteOk);
151 ASSERT_NE(interpreter.AllocateTensors(), kTfLiteOk);
152 ASSERT_NE(interpreter.Invoke(), kTfLiteOk);
153 ASSERT_EQ(interpreter.inputs(), std::vector<int>());
154 }
155 // Invalid outputs
156 {
157 Interpreter interpreter;
158 ASSERT_NE(interpreter.SetOutputs({5}), kTfLiteOk);
159 ASSERT_NE(interpreter.AllocateTensors(), kTfLiteOk);
160 ASSERT_NE(interpreter.Invoke(), kTfLiteOk);
161 ASSERT_EQ(interpreter.outputs(), std::vector<int>());
162 }
163 // Invalid node inputs
164 {
165 Interpreter interpreter;
166 TfLiteRegistration registration = {nullptr, nullptr, nullptr, nullptr};
167 ASSERT_NE(interpreter.AddNodeWithParameters({3}, {0}, nullptr, 0, nullptr,
168 ®istration),
169 kTfLiteOk);
170 ASSERT_NE(interpreter.AllocateTensors(), kTfLiteOk);
171 ASSERT_NE(interpreter.Invoke(), kTfLiteOk);
172 }
173 // Valid inputs and outputs and a node with valid inputs and outputs
174 {
175 Interpreter interpreter;
176 ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk);
177 TfLiteRegistration registration = {nullptr, nullptr, nullptr, nullptr};
178 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
179 ASSERT_EQ(interpreter.SetOutputs({0}), kTfLiteOk);
180 ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr,
181 ®istration),
182 kTfLiteOk);
183 }
184 }
185
186 // Make an interpreter that has one tensor but no ops
TEST(BasicInterpreter,CheckAllocate)187 TEST(BasicInterpreter, CheckAllocate) {
188 struct {
189 TfLiteType type;
190 size_t size;
191 } cases[] = {
192 {kTfLiteFloat32, sizeof(float)}, {kTfLiteInt32, sizeof(int32_t)},
193 {kTfLiteUInt32, sizeof(uint32_t)}, {kTfLiteUInt8, sizeof(uint8_t)},
194 {kTfLiteInt64, sizeof(int64_t)}, {kTfLiteInt16, sizeof(int16_t)},
195 {kTfLiteFloat16, sizeof(TfLiteFloat16)},
196 };
197
198 for (auto test : cases) {
199 Interpreter interpreter;
200 ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk);
201 interpreter.SetInputs({0, 1});
202 interpreter.SetOutputs({});
203 TfLiteQuantizationParams quant;
204
205 interpreter.SetTensorParametersReadWrite(0, test.type, "", {3}, quant);
206 interpreter.SetTensorParametersReadWrite(1, test.type, "", {4}, quant);
207 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
208 ASSERT_EQ(interpreter.tensor(0)->bytes, 3 * test.size);
209 ASSERT_NE(interpreter.tensor(0)->data.raw, nullptr);
210 ASSERT_EQ(interpreter.tensor(1)->bytes, 4 * test.size);
211 ASSERT_NE(interpreter.tensor(1)->data.raw, nullptr);
212 }
213 }
214
TEST(BasicInterpreter,CheckQuantization)215 TEST(BasicInterpreter, CheckQuantization) {
216 Interpreter interpreter;
217 ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk);
218 interpreter.SetInputs({0, 1});
219 interpreter.SetOutputs({});
220 TfLiteType tensor_type = kTfLiteInt8;
221 const uint8_t int8s[] = {3, 4};
222 float scale = 0.5f;
223 int32_t zero_point = 12;
224
225 TfLiteQuantization rw_quantization;
226 rw_quantization.type = kTfLiteAffineQuantization;
227 auto* rw_affine_quantization = static_cast<TfLiteAffineQuantization*>(
228 malloc(sizeof(TfLiteAffineQuantization)));
229 rw_affine_quantization->scale = TfLiteFloatArrayCreate(1);
230 rw_affine_quantization->zero_point = TfLiteIntArrayCreate(1);
231 rw_affine_quantization->scale->data[0] = scale;
232 rw_affine_quantization->zero_point->data[0] = zero_point;
233 rw_quantization.params = rw_affine_quantization;
234
235 TfLiteQuantization ro_quantization;
236 ro_quantization.type = kTfLiteAffineQuantization;
237 auto* ro_affine_quantization = static_cast<TfLiteAffineQuantization*>(
238 malloc(sizeof(TfLiteAffineQuantization)));
239 ro_affine_quantization->scale = TfLiteFloatArrayCreate(1);
240 ro_affine_quantization->zero_point = TfLiteIntArrayCreate(1);
241 ro_affine_quantization->scale->data[0] = scale;
242 ro_affine_quantization->zero_point->data[0] = zero_point;
243 ro_quantization.params = ro_affine_quantization;
244
245 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(0, tensor_type, "", {3},
246 rw_quantization),
247 kTfLiteOk);
248 ASSERT_EQ(interpreter.SetTensorParametersReadOnly(
249 1, tensor_type, "", {2}, ro_quantization,
250 reinterpret_cast<const char*>(int8s), 2),
251 kTfLiteOk);
252 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
253 // Check that the legacy scale and zero_point are set correctly.
254 ASSERT_EQ(interpreter.tensor(0)->params.scale, scale);
255 ASSERT_EQ(interpreter.tensor(0)->params.zero_point, zero_point);
256 ASSERT_EQ(interpreter.tensor(0)->quantization.type, rw_quantization.type);
257 ASSERT_EQ(interpreter.tensor(1)->params.scale, scale);
258 ASSERT_EQ(interpreter.tensor(1)->params.zero_point, zero_point);
259 ASSERT_EQ(interpreter.tensor(1)->quantization.type, ro_quantization.type);
260 }
261
TEST(BasicInterpreter,CheckResize)262 TEST(BasicInterpreter, CheckResize) {
263 const float floats[] = {-3., -4.};
264 const int32_t int32s[] = {-3, -4};
265 const uint32_t uint32s[] = {3, 4};
266 const uint8_t uint8s[] = {3, 4};
267 const int64_t int64s[] = {6, -7};
268 const int16_t int16s[] = {8, -9};
269 const Eigen::half float16s[] = {Eigen::half(-3.f), Eigen::half(-4.f)};
270
271 struct {
272 TfLiteType type;
273 size_t size;
274 const char* array;
275 } cases[] = {
276 {kTfLiteFloat32, sizeof(float), reinterpret_cast<const char*>(floats)},
277 {kTfLiteInt32, sizeof(int32_t), reinterpret_cast<const char*>(int32s)},
278 {kTfLiteUInt32, sizeof(uint32_t), reinterpret_cast<const char*>(uint32s)},
279 {kTfLiteUInt8, sizeof(uint8_t), reinterpret_cast<const char*>(uint8s)},
280 {kTfLiteInt64, sizeof(int64_t), reinterpret_cast<const char*>(int64s)},
281 {kTfLiteInt16, sizeof(int16_t), reinterpret_cast<const char*>(int16s)},
282 {kTfLiteFloat16, sizeof(TfLiteFloat16),
283 reinterpret_cast<const char*>(float16s)},
284 };
285
286 for (auto test : cases) {
287 Interpreter interpreter;
288
289 ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk);
290 interpreter.SetInputs({0, 1});
291 interpreter.SetOutputs({});
292 TfLiteQuantizationParams quant;
293
294 ASSERT_EQ(
295 interpreter.SetTensorParametersReadWrite(0, test.type, "", {3}, quant),
296 kTfLiteOk);
297 ASSERT_EQ(interpreter.SetTensorParametersReadOnly(
298 1, test.type, "", {2}, quant, test.array, 2 * test.size),
299 kTfLiteOk);
300 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
301 ASSERT_EQ(interpreter.ResizeInputTensor(0, {1, 2}), kTfLiteOk);
302 // Resizing a mmapped tensor is not allowed and should produce error.
303 ASSERT_NE(interpreter.ResizeInputTensor(1, {3}), kTfLiteOk);
304 // Set the tensor to be mmapped but with a buffer size that is insufficient
305 // to match the dimensionality.
306 ASSERT_NE(interpreter.SetTensorParametersReadOnly(
307 1, test.type, "", {2}, quant, test.array, 1 * test.size),
308 kTfLiteOk);
309 // Allocating should work since we should have our last correct array
310 // values in place.
311 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
312 }
313 }
314
TEST(BasicInterpreter,CheckAlignment)315 TEST(BasicInterpreter, CheckAlignment) {
316 struct {
317 TfLiteType type;
318 } cases[] = {{kTfLiteFloat32}, {kTfLiteInt32}, {kTfLiteUInt32},
319 {kTfLiteUInt8}, {kTfLiteInt64}, {kTfLiteInt16},
320 {kTfLiteFloat16}};
321
322 for (auto test : cases) {
323 Interpreter interpreter;
324
325 ASSERT_EQ(interpreter.AddTensors(4), kTfLiteOk);
326
327 for (int i = 0; i < 4; i++) {
328 TfLiteQuantizationParams quant;
329 interpreter.SetTensorParametersReadWrite(i, test.type, "", {2 * i + 1},
330 quant);
331 }
332 interpreter.AllocateTensors();
333 for (int i = 0; i < 4; i++) {
334 const TfLiteTensor& tensor = *interpreter.tensor(i);
335 ASSERT_EQ(reinterpret_cast<intptr_t>(tensor.data.raw) % 4, 0);
336 }
337 }
338 }
339
TEST(BasicInterpreter,CheckArenaAllocation)340 TEST(BasicInterpreter, CheckArenaAllocation) {
341 Interpreter interpreter;
342 ASSERT_EQ(interpreter.AddTensors(10), kTfLiteOk);
343
344 TfLiteQuantizationParams quant;
345 TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
346
347 std::vector<int> sizes{2048, 4096, 1023, 2047, 1021,
348 2047, 1023, 2046, 0, 2048};
349 for (size_t i = 0; i < sizes.size(); ++i) {
350 interpreter.SetTensorParametersReadWrite(static_cast<int>(i), kTfLiteUInt8,
351 "", {sizes[i]}, quant);
352 }
353 interpreter.SetInputs({0, 1});
354 interpreter.SetOutputs({9, 4});
355 interpreter.AddNodeWithParameters({0, 1}, {2, 3}, nullptr, 0, nullptr, ®);
356 interpreter.AddNodeWithParameters({2, 1}, {4, 5}, nullptr, 0, nullptr, ®);
357 interpreter.AddNodeWithParameters({4, 3}, {6, 7}, nullptr, 0, nullptr, ®);
358 interpreter.AddNodeWithParameters({6, 5}, {8}, nullptr, 0, nullptr, ®);
359 interpreter.AddNodeWithParameters({8, 7}, {9}, nullptr, 0, nullptr, ®);
360
361 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
362
363 ASSERT_LT(interpreter.tensor(0)->data.raw, interpreter.tensor(1)->data.raw);
364 ASSERT_LT(interpreter.tensor(1)->data.raw, interpreter.tensor(3)->data.raw);
365 ASSERT_EQ(interpreter.tensor(3)->data.raw, interpreter.tensor(9)->data.raw);
366 ASSERT_LT(interpreter.tensor(3)->data.raw, interpreter.tensor(5)->data.raw);
367 ASSERT_LT(interpreter.tensor(5)->data.raw, interpreter.tensor(2)->data.raw);
368 ASSERT_EQ(interpreter.tensor(2)->data.raw, interpreter.tensor(7)->data.raw);
369 ASSERT_LT(interpreter.tensor(2)->data.raw, interpreter.tensor(4)->data.raw);
370 // #4 is the one with the largest pointer.
371 ASSERT_EQ(interpreter.tensor(8)->data.raw, nullptr);
372 }
373
TEST(BasicInterpreter,BufferAccess)374 TEST(BasicInterpreter, BufferAccess) {
375 Interpreter interpreter;
376 ASSERT_EQ(interpreter.AddTensors(1), kTfLiteOk);
377 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
378
379 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(
380 0, kTfLiteFloat32, "", {3}, TfLiteQuantizationParams()),
381 kTfLiteOk);
382 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
383 // Verify we get a valid pointer.
384 ASSERT_NE(interpreter.typed_tensor<float>(0), nullptr);
385 // Verify incorrect pointer is not returned.
386 ASSERT_EQ(interpreter.typed_tensor<int>(0), nullptr);
387 // Verify that raw c interface ptr matches safe interface.
388 ASSERT_EQ(interpreter.typed_tensor<float>(0), interpreter.tensor(0)->data.f);
389 }
390
TEST(BasicInterpreter,NoOpInterpreter)391 TEST(BasicInterpreter, NoOpInterpreter) {
392 Interpreter interpreter;
393 ASSERT_EQ(interpreter.AddTensors(1), kTfLiteOk);
394 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
395 ASSERT_EQ(interpreter.SetOutputs({0}), kTfLiteOk);
396
397 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(
398 0, kTfLiteFloat32, "", {3}, TfLiteQuantizationParams()),
399 kTfLiteOk);
400
401 ASSERT_EQ(interpreter.ResizeInputTensor(interpreter.inputs()[0], {1, 2, 3}),
402 kTfLiteOk);
403 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
404 ASSERT_EQ(interpreter.Invoke(), kTfLiteOk);
405 }
406
TEST(BasicInterpreter,RedundantAllocateTensors)407 TEST(BasicInterpreter, RedundantAllocateTensors) {
408 Interpreter interpreter;
409 ASSERT_EQ(interpreter.AddTensors(1), kTfLiteOk);
410 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
411
412 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(
413 0, kTfLiteFloat32, "", {3}, TfLiteQuantizationParams()),
414 kTfLiteOk);
415
416 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
417 const auto data_raw = interpreter.tensor(0)->data.raw;
418 ASSERT_NE(data_raw, nullptr);
419
420 // A redundant allocation request should have no impact.
421 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
422 ASSERT_EQ(interpreter.tensor(0)->data.raw, data_raw);
423 }
424
TEST(BasicInterpreter,RedundantAllocateTensorsWithDynamicInputs)425 TEST(BasicInterpreter, RedundantAllocateTensorsWithDynamicInputs) {
426 Interpreter interpreter;
427 TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
428 ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk);
429 interpreter.SetInputs({0});
430 interpreter.SetOutputs({1});
431 interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr, ®);
432
433 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(
434 0, kTfLiteFloat32, "", {3}, TfLiteQuantizationParams()),
435 kTfLiteOk);
436 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(
437 1, kTfLiteFloat32, "", {3}, TfLiteQuantizationParams()),
438 kTfLiteOk);
439
440 // Configure the input tensor as dynamic.
441 interpreter.tensor(0)->data.raw = nullptr;
442 interpreter.tensor(0)->allocation_type = kTfLiteDynamic;
443
444 ASSERT_EQ(interpreter.ResizeInputTensor(interpreter.inputs()[0], {1, 2, 3}),
445 kTfLiteOk);
446 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
447 ASSERT_NE(interpreter.tensor(1)->data.raw, nullptr);
448
449 // Reset the output tensor's buffer.
450 interpreter.tensor(1)->data.raw = nullptr;
451
452 // A redundant allocation request should be honored, as the input tensor
453 // was marked dynamic.
454 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
455 ASSERT_NE(interpreter.tensor(1)->data.raw, nullptr);
456 }
457
TEST(BasicInterpreter,ResizingTensors)458 TEST(BasicInterpreter, ResizingTensors) {
459 Interpreter interpreter;
460 ASSERT_EQ(interpreter.AddTensors(1), kTfLiteOk);
461 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
462 ASSERT_EQ(interpreter.SetOutputs({0}), kTfLiteOk);
463
464 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(
465 0, kTfLiteFloat32, "", {3}, TfLiteQuantizationParams()),
466 kTfLiteOk);
467
468 int t = interpreter.inputs()[0];
469 TfLiteTensor* tensor = interpreter.tensor(t);
470
471 ASSERT_EQ(interpreter.ResizeInputTensor(t, {1, 2, 3}), kTfLiteOk);
472 EXPECT_EQ(tensor->bytes, 6 * sizeof(float));
473 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
474
475 tensor->data.f[5] = 0.123f;
476
477 // Changing from kTfLiteArenaRw to kTfLiteDynamic is quite complicate: we need
478 // to unset data.raw, otherwise Realloc will try to free that memory.
479 tensor->data.raw = nullptr;
480 tensor->allocation_type = kTfLiteDynamic;
481
482 ASSERT_EQ(interpreter.ResizeInputTensor(t, {1, 2, 4}), kTfLiteOk);
483 EXPECT_EQ(tensor->bytes, 8 * sizeof(float));
484 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
485
486 ASSERT_EQ(interpreter.ResizeInputTensor(t, {}), kTfLiteOk);
487 EXPECT_EQ(tensor->bytes, 1 * sizeof(float));
488 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
489
490 ASSERT_EQ(interpreter.ResizeInputTensor(t, {0}), kTfLiteOk);
491 EXPECT_EQ(tensor->bytes, 0);
492 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
493
494 ASSERT_EQ(interpreter.ResizeInputTensor(t, {1, 2, 0}), kTfLiteOk);
495 EXPECT_EQ(tensor->bytes, 0);
496 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
497
498 // TODO(ahentz): We shouldn't have to force reallocation, but
499 // ResizeInputTensor doesn't realloc dynamic tensors. Also note that
500 // TfLiteTensorRealloc(tensor->bytes, tensor) is a no-op.
501 TfLiteTensorRealloc(9 * sizeof(float), tensor);
502 tensor->data.f[7] = 0.123f;
503
504 ASSERT_EQ(interpreter.ResizeInputTensor(t, {2, 2, 4}), kTfLiteOk);
505 EXPECT_EQ(tensor->bytes, 16 * sizeof(float));
506 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
507
508 // TODO(ahentz): We shouldn't have to force reallocation, but
509 // ResizeInputTensor doesn't realloc dynamic tensors. Also note that
510 // TfLiteTensorRealloc(tensor->bytes, tensor) is a no-op.
511 TfLiteTensorRealloc(17 * sizeof(float), tensor);
512 tensor->data.f[15] = 0.123f;
513 }
514
TEST(BasicInterpreter,NoopResizingTensors)515 TEST(BasicInterpreter, NoopResizingTensors) {
516 Interpreter interpreter;
517 ASSERT_EQ(interpreter.AddTensors(1), kTfLiteOk);
518 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
519 ASSERT_EQ(interpreter.SetOutputs({0}), kTfLiteOk);
520
521 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(
522 0, kTfLiteFloat32, "", {3}, TfLiteQuantizationParams()),
523 kTfLiteOk);
524
525 int t = interpreter.inputs()[0];
526 TfLiteTensor* tensor = interpreter.tensor(t);
527
528 ASSERT_EQ(interpreter.ResizeInputTensor(t, {1, 2, 3}), kTfLiteOk);
529 EXPECT_EQ(tensor->bytes, 6 * sizeof(float));
530 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
531 tensor->data.f[5] = 0.123f;
532
533 // Resizing to the same size should not trigger re-allocation.
534 ASSERT_EQ(interpreter.ResizeInputTensor(t, {1, 2, 3}), kTfLiteOk);
535 EXPECT_EQ(tensor->bytes, 6 * sizeof(float));
536 ASSERT_NE(tensor->data.raw, nullptr);
537 ASSERT_EQ(tensor->data.f[5], 0.123f);
538
539 // Explicitly allocating should be a no-op, as no resize was performed.
540 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
541 EXPECT_EQ(tensor->bytes, 6 * sizeof(float));
542 ASSERT_NE(tensor->data.raw, nullptr);
543 ASSERT_EQ(tensor->data.f[5], 0.123f);
544 }
545
TEST(BasicInterpreter,ResizingTensorsStrictInvalid)546 TEST(BasicInterpreter, ResizingTensorsStrictInvalid) {
547 // Tests ResizeInputTensorStrict where `dims_signature` is not specified.
548 Interpreter interpreter;
549 ASSERT_EQ(interpreter.AddTensors(1), kTfLiteOk);
550 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
551 ASSERT_EQ(interpreter.SetOutputs({0}), kTfLiteOk);
552
553 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(
554 0, kTfLiteFloat32, "", {1, 1, 3}, TfLiteQuantizationParams()),
555 kTfLiteOk);
556
557 int t = interpreter.inputs()[0];
558 TfLiteTensor* tensor = interpreter.tensor(t);
559
560 ASSERT_EQ(interpreter.ResizeInputTensorStrict(t, {1, 1, 3}), kTfLiteOk);
561 EXPECT_EQ(tensor->bytes, 3 * sizeof(float));
562 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
563
564 // Invalid becuase `dims_signature` is not specified.
565 ASSERT_EQ(interpreter.ResizeInputTensorStrict(t, {1, 2, 3}), kTfLiteError);
566 EXPECT_EQ(tensor->bytes, 3 * sizeof(float));
567 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
568
569 // Assert that ResizeInputTensor works for this value.
570 ASSERT_EQ(interpreter.ResizeInputTensor(t, {1, 2, 3}), kTfLiteOk);
571 EXPECT_EQ(tensor->bytes, 6 * sizeof(float));
572 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
573 }
574
TEST(BasicInterpreter,ResizingTensorsStrict)575 TEST(BasicInterpreter, ResizingTensorsStrict) {
576 // Tests ResizeInputTensorStrict where `dims_signature` is specified.
577 Interpreter interpreter;
578 ASSERT_EQ(interpreter.AddTensors(1), kTfLiteOk);
579 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
580 ASSERT_EQ(interpreter.SetOutputs({0}), kTfLiteOk);
581
582 std::vector<int> dims_signature = {-1, -1, 3};
583 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(
584 0, kTfLiteFloat32, "", {1, 1, 3}, TfLiteQuantizationParams(),
585 false, &dims_signature),
586 kTfLiteOk);
587
588 int t = interpreter.inputs()[0];
589 TfLiteTensor* tensor = interpreter.tensor(t);
590
591 ASSERT_EQ(interpreter.ResizeInputTensorStrict(t, {1, 2, 3}), kTfLiteOk);
592 EXPECT_EQ(tensor->bytes, 6 * sizeof(float));
593 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
594
595 ASSERT_EQ(interpreter.ResizeInputTensorStrict(t, {1, 2, 4}), kTfLiteError);
596 EXPECT_EQ(tensor->bytes, 6 * sizeof(float));
597 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
598
599 // Assert that ResizeInputTensor works for this value.
600 ASSERT_EQ(interpreter.ResizeInputTensor(t, {1, 2, 4}), kTfLiteOk);
601 EXPECT_EQ(tensor->bytes, 8 * sizeof(float));
602 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
603 }
604
605 // Simple op that does input = output.
GetPassthroughOpRegistration()606 TfLiteRegistration GetPassthroughOpRegistration() {
607 TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
608 reg.init = [](TfLiteContext* context, const char*, size_t) -> void* {
609 auto* first_new_tensor = new int;
610 context->AddTensors(context, 2, first_new_tensor);
611 return first_new_tensor;
612 };
613 reg.free = [](TfLiteContext* context, void* buffer) {
614 delete static_cast<int*>(buffer);
615 };
616 reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
617 auto* first_new_tensor = static_cast<int*>(node->user_data);
618
619 const TfLiteTensor* tensor0;
620 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &tensor0));
621 TfLiteTensor* tensor1;
622 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &tensor1));
623
624 TfLiteIntArray* newSize = TfLiteIntArrayCopy(tensor0->dims);
625 TF_LITE_ENSURE_STATUS(context->ResizeTensor(context, tensor1, newSize));
626
627 TfLiteIntArrayFree(node->temporaries);
628 node->temporaries = TfLiteIntArrayCreate(2);
629 for (int i = 0; i < 2; ++i) {
630 node->temporaries->data[i] = *(first_new_tensor) + i;
631 }
632
633 auto setup_temporary = [&](int id) {
634 TfLiteTensor* tmp = &context->tensors[id];
635 tmp->type = kTfLiteFloat32;
636 tmp->allocation_type = kTfLiteArenaRw;
637 return context->ResizeTensor(context, tmp,
638 TfLiteIntArrayCopy(tensor0->dims));
639 };
640 TF_LITE_ENSURE_STATUS(setup_temporary(node->temporaries->data[0]));
641 TF_LITE_ENSURE_STATUS(setup_temporary(node->temporaries->data[1]));
642
643 return kTfLiteOk;
644 };
645 reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
646 const TfLiteTensor* a0;
647 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
648
649 auto populate = [&](int id) {
650 TfLiteTensor* t = &context->tensors[id];
651 int num = a0->dims->data[0];
652 for (int i = 0; i < num; i++) {
653 t->data.f[i] = a0->data.f[i];
654 }
655 };
656
657 populate(node->outputs->data[0]);
658 populate(node->temporaries->data[0]);
659 populate(node->temporaries->data[1]);
660 return kTfLiteOk;
661 };
662
663 return reg;
664 }
665
TEST(BasicInterpreter,OneOpInterpreter)666 TEST(BasicInterpreter, OneOpInterpreter) {
667 Interpreter interpreter;
668 ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk);
669 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
670 ASSERT_EQ(interpreter.SetOutputs({1}), kTfLiteOk);
671
672 TfLiteQuantizationParams quantized;
673 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(0, kTfLiteFloat32, "in1",
674 {3}, quantized),
675 kTfLiteOk);
676 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(1, kTfLiteFloat32, "out0",
677 {3}, quantized),
678 kTfLiteOk);
679
680 ASSERT_EQ(interpreter.GetInputName(0), "in1");
681 ASSERT_EQ(interpreter.GetOutputName(0), "out0");
682
683 TfLiteRegistration reg = GetPassthroughOpRegistration();
684
685 ASSERT_EQ(
686 interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr, ®),
687 kTfLiteOk);
688 ASSERT_EQ(interpreter.ResizeInputTensor(0, {3}), kTfLiteOk);
689 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
690
691 ASSERT_EQ(interpreter.Invoke(), kTfLiteOk);
692 }
693
TEST(BasicInterpreter,ReleaseNonPersistentMemory)694 TEST(BasicInterpreter, ReleaseNonPersistentMemory) {
695 Interpreter interpreter;
696 ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk);
697 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
698 ASSERT_EQ(interpreter.SetOutputs({1}), kTfLiteOk);
699
700 TfLiteQuantizationParams quantized;
701 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(0, kTfLiteFloat32, "in1",
702 {3}, quantized),
703 kTfLiteOk);
704 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(1, kTfLiteFloat32, "out0",
705 {3}, quantized),
706 kTfLiteOk);
707
708 TfLiteRegistration reg = GetPassthroughOpRegistration();
709
710 ASSERT_EQ(
711 interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr, ®),
712 kTfLiteOk);
713 ASSERT_EQ(interpreter.ResizeInputTensor(0, {3}), kTfLiteOk);
714
715 // AllocateTensors() hasn't been called yet, so this should be a no-op.
716 ASSERT_EQ(interpreter.ReleaseNonPersistentMemory(), kTfLiteOk);
717
718 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
719 ASSERT_EQ(interpreter.Invoke(), kTfLiteOk);
720
721 ASSERT_EQ(interpreter.ReleaseNonPersistentMemory(), kTfLiteOk);
722 // Invoke() now fails because non-persistent arenas have been released.
723 ASSERT_NE(interpreter.Invoke(), kTfLiteOk);
724
725 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
726 ASSERT_EQ(interpreter.Invoke(), kTfLiteOk);
727
728 // ResizeInputTensors just after ReleaseNonPersistentMemory should also need
729 // AllocateTensors, without causing any unexpected crashes.
730 ASSERT_EQ(interpreter.ReleaseNonPersistentMemory(), kTfLiteOk);
731 ASSERT_EQ(interpreter.ResizeInputTensor(0, {4}), kTfLiteOk);
732 ASSERT_NE(interpreter.Invoke(), kTfLiteOk);
733 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
734 ASSERT_EQ(interpreter.Invoke(), kTfLiteOk);
735 }
736
737 // Forcefully divides tensor allocation in three steps: one before invocation
738 // and two more at invocation time. This happens because we use string tensors
739 // and their sizes can't be determined until invocation time.
TEST(BasicInterpreter,ThreeStepAllocate)740 TEST(BasicInterpreter, ThreeStepAllocate) {
741 Interpreter interpreter;
742 ASSERT_EQ(interpreter.AddTensors(5), kTfLiteOk);
743 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
744 ASSERT_EQ(interpreter.SetOutputs({4}), kTfLiteOk);
745
746 TfLiteQuantizationParams quantized;
747
748 // String tensor with one string of length 3
749 union {
750 char raw_bytes[15];
751 struct {
752 int32_t num_strs;
753 int32_t offsets[2];
754 char str_data[3];
755 } tensor_data;
756 } data;
757 data.tensor_data = {1, {12, 15}, {'A', 'B', 'C'}};
758
759 // Read only string tensor.
760 ASSERT_EQ(interpreter.SetTensorParametersReadOnly(0, kTfLiteString, "", {1},
761 quantized, data.raw_bytes,
762 sizeof(data.raw_bytes)),
763 kTfLiteOk);
764 // Read-write string tensor.
765 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(1, kTfLiteString, "", {1},
766 quantized),
767 kTfLiteOk);
768 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(2, kTfLiteInt32, "", {1},
769 quantized),
770 kTfLiteOk);
771 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(3, kTfLiteString, "", {1},
772 quantized),
773 kTfLiteOk);
774 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(4, kTfLiteInt32, "", {1},
775 quantized),
776 kTfLiteOk);
777
778 // String-in String-out node.
779 TfLiteRegistration reg_copy = {nullptr, nullptr, nullptr, nullptr};
780 reg_copy.invoke = [](TfLiteContext* context, TfLiteNode* node) {
781 const TfLiteTensor* input;
782 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
783 TfLiteTensor* output;
784 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
785 DynamicBuffer buf;
786 StringRef str_ref = GetString(input, 0);
787 buf.AddString(str_ref);
788 buf.WriteToTensorAsVector(output);
789 return kTfLiteOk;
790 };
791
792 // String-in Int-out node.
793 TfLiteRegistration reg_len = {nullptr, nullptr, nullptr, nullptr};
794 reg_len.prepare = [](TfLiteContext* context, TfLiteNode* node) {
795 TfLiteTensor* output;
796 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
797 TfLiteIntArray* outputSize = TfLiteIntArrayCreate(1);
798 outputSize->data[0] = 1;
799 return context->ResizeTensor(context, output, outputSize);
800 };
801 reg_len.invoke = [](TfLiteContext* context, TfLiteNode* node) {
802 const TfLiteTensor* a0;
803 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
804 TfLiteTensor* a1;
805 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &a1));
806 a1->data.i32[0] = a0->bytes;
807 return kTfLiteOk;
808 };
809
810 ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr,
811 ®_copy),
812 kTfLiteOk);
813 ASSERT_EQ(interpreter.AddNodeWithParameters({1}, {2}, nullptr, 0, nullptr,
814 ®_len),
815 kTfLiteOk);
816 ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {3}, nullptr, 0, nullptr,
817 ®_copy),
818 kTfLiteOk);
819 ASSERT_EQ(interpreter.AddNodeWithParameters({3}, {4}, nullptr, 0, nullptr,
820 ®_len),
821 kTfLiteOk);
822
823 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
824 ASSERT_EQ(interpreter.Invoke(), kTfLiteOk);
825
826 ASSERT_EQ(interpreter.tensor(0)->bytes, 15);
827 ASSERT_NE(interpreter.tensor(0)->data.raw, nullptr);
828 ASSERT_EQ(interpreter.tensor(1)->bytes, 15);
829 ASSERT_NE(interpreter.tensor(1)->data.raw, nullptr);
830 ASSERT_EQ(interpreter.tensor(3)->bytes, 15);
831 ASSERT_NE(interpreter.tensor(4)->data.raw, nullptr);
832 ASSERT_EQ(interpreter.tensor(2)->bytes, 4);
833 ASSERT_EQ(interpreter.tensor(2)->data.i32[0], 15);
834 ASSERT_EQ(interpreter.tensor(4)->bytes, 4);
835 ASSERT_EQ(interpreter.tensor(4)->data.i32[0], 15);
836 }
837
TEST(BasicInterpreter,AllocateTwice)838 TEST(BasicInterpreter, AllocateTwice) {
839 Interpreter interpreter;
840 ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk);
841 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
842 ASSERT_EQ(interpreter.SetOutputs({1}), kTfLiteOk);
843
844 TfLiteQuantizationParams quantized;
845 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(0, kTfLiteFloat32, "", {3},
846 quantized),
847 kTfLiteOk);
848 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(1, kTfLiteFloat32, "", {3},
849 quantized),
850 kTfLiteOk);
851
852 TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
853 reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
854 const TfLiteTensor* tensor0;
855 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &tensor0));
856 TfLiteTensor* tensor1;
857 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &tensor1));
858 TfLiteIntArray* newSize = TfLiteIntArrayCopy(tensor0->dims);
859 return context->ResizeTensor(context, tensor1, newSize);
860 };
861 reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
862 const TfLiteTensor* a0;
863 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
864 TfLiteTensor* a1;
865 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &a1));
866 int num = a0->dims->data[0];
867 for (int i = 0; i < num; i++) {
868 a1->data.f[i] = a0->data.f[i];
869 }
870 return kTfLiteOk;
871 };
872 ASSERT_EQ(
873 interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr, ®),
874 kTfLiteOk);
875 ASSERT_EQ(interpreter.ResizeInputTensor(0, {3}), kTfLiteOk);
876 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
877 ASSERT_EQ(interpreter.Invoke(), kTfLiteOk);
878 char* old_tensor0_ptr = interpreter.tensor(0)->data.raw;
879 char* old_tensor1_ptr = interpreter.tensor(1)->data.raw;
880
881 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
882 ASSERT_EQ(interpreter.Invoke(), kTfLiteOk);
883 ASSERT_EQ(old_tensor0_ptr, interpreter.tensor(0)->data.raw);
884 ASSERT_EQ(old_tensor1_ptr, interpreter.tensor(1)->data.raw);
885 }
886
TEST(BasicInterpreter,TestNullErrorReporter)887 TEST(BasicInterpreter, TestNullErrorReporter) {
888 TestErrorReporter reporter;
889 Interpreter interpreter;
890 }
891
TEST(BasicInterpreter,TestCustomErrorReporter)892 TEST(BasicInterpreter, TestCustomErrorReporter) {
893 TestErrorReporter reporter;
894 Interpreter interpreter(&reporter);
895 ASSERT_NE(interpreter.Invoke(), kTfLiteOk);
896 ASSERT_EQ(reporter.error_messages(),
897 "Invoke called on model that is not ready.");
898 ASSERT_EQ(reporter.num_calls(), 1);
899 }
900
TEST(BasicInterpreter,TestOverflow)901 TEST(BasicInterpreter, TestOverflow) {
902 TestErrorReporter reporter;
903 Interpreter interpreter(&reporter);
904 TfLiteQuantizationParams quantized;
905
906 ASSERT_EQ(interpreter.AddTensors(1), kTfLiteOk);
907 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
908 ASSERT_EQ(interpreter.SetOutputs({0}), kTfLiteOk);
909 // Overflow testing is pointer word size dependent.
910 if (sizeof(size_t) == 8) {
911 // #bits for bytecount = 30 + 30 + 2 = 62 < 64
912 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(
913 0, kTfLiteFloat32, "in1", {1 << 30, 1 << 30}, quantized),
914 kTfLiteOk);
915 // #bits for element count = 30 + 30 + 2 = 62 < 64 (no overflow)
916 // #bits for byte count = 30 + 30 + 2 + 2 = 64 == 64 (overflow)
917 ASSERT_NE(
918 interpreter.SetTensorParametersReadWrite(
919 0, kTfLiteFloat32, "in1", {1 << 30, 1 << 30, 1 << 2}, quantized),
920 kTfLiteOk);
921 EXPECT_THAT(
922 reporter.error_messages(),
923 testing::EndsWith("BytesRequired number of bytes overflowed.\n"));
924 // #bits for element count = 30 + 30 + 2 + 4 = 66 > 64 (overflow).
925 // #bits for byte count = 30 + 30 + 2 + 4 + 2 = 68 > 64 (overflow).
926 reporter.Reset();
927 ASSERT_NE(interpreter.SetTensorParametersReadWrite(
928 0, kTfLiteFloat32, "in1", {1 << 30, 1 << 30, 1 << 2, 1 << 4},
929 quantized),
930 kTfLiteOk);
931 EXPECT_THAT(
932 reporter.error_messages(),
933 testing::EndsWith("BytesRequired number of elements overflowed.\n"));
934
935 } else if (sizeof(size_t) == 4) {
936 // #bits for bytecount = 14 + 14 + 2 = 30 < 32
937 ASSERT_EQ(interpreter.SetTensorParametersReadWrite(
938 0, kTfLiteFloat32, "in1", {1 << 14, 1 << 14}, quantized),
939 kTfLiteOk);
940 // #bits for element count = 14 + 14 + 3 = 31 < 32 (no overflow).
941 // #bits for byte count = 14 + 14 + 3 + 2 = 33 > 32 (overflow).
942 ASSERT_NE(
943 interpreter.SetTensorParametersReadWrite(
944 0, kTfLiteFloat32, "in1", {1 << 14, 1 << 14, 1 << 3}, quantized),
945 kTfLiteOk);
946 EXPECT_THAT(
947 reporter.error_messages(),
948 testing::EndsWith("BytesRequired number of bytes overflowed.\n"));
949 // #bits for element count = 14 + 14 + 4 = 32 == 32 (overflow).
950 // byte count also overflows, but we don't get to that check.
951 reporter.Reset();
952 ASSERT_NE(
953 interpreter.SetTensorParametersReadWrite(
954 0, kTfLiteFloat32, "in1", {1 << 14, 1 << 14, 1 << 4}, quantized),
955 kTfLiteOk);
956 EXPECT_THAT(
957 reporter.error_messages(),
958 testing::EndsWith("BytesRequired number of elements overflowed.\n"));
959 } else {
960 // This test failing means that we are using a non 32/64 bit architecture.
961 ASSERT_TRUE(false);
962 }
963 }
964
TEST(BasicInterpreter,TestUnsupportedDelegateFunctions)965 TEST(BasicInterpreter, TestUnsupportedDelegateFunctions) {
966 Interpreter interpreter;
967 ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk);
968 TfLiteRegistration registration = {nullptr, nullptr, nullptr, nullptr};
969 // These functions are only supported inside Delegate's Prepare function.
970 // The test verifies that these functions returns `kTfLiteError`, but not
971 // `kTfLiteOk` or just crashes.
972 registration.prepare = [](TfLiteContext* context, TfLiteNode* node) {
973 {
974 TfLiteIntArray* execution_plan;
975 EXPECT_EQ(context->GetExecutionPlan(context, &execution_plan),
976 kTfLiteError);
977 }
978 {
979 TfLiteNode* node;
980 TfLiteRegistration* registration;
981 EXPECT_EQ(
982 context->GetNodeAndRegistration(context, 0, &node, ®istration),
983 kTfLiteError);
984 }
985 {
986 TfLiteRegistration delegate_registration = {nullptr, nullptr, nullptr,
987 nullptr};
988 TfLiteIntArray nodes_to_replace;
989 nodes_to_replace.size = 0;
990 EXPECT_EQ(context->ReplaceNodeSubsetsWithDelegateKernels(
991 context, delegate_registration, &nodes_to_replace, nullptr),
992 kTfLiteError);
993 }
994 return kTfLiteError;
995 };
996 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
997 ASSERT_EQ(interpreter.SetOutputs({0}), kTfLiteOk);
998 ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr,
999 ®istration),
1000 kTfLiteOk);
1001 EXPECT_EQ(interpreter.AllocateTensors(), kTfLiteError);
1002 }
1003
TEST(BasicInterpreter,DynamicTensorsResizeDescendants)1004 TEST(BasicInterpreter, DynamicTensorsResizeDescendants) {
1005 // Assemble a graph with a node that has dynamically sized output (via the
1006 // pad op), followed by a node with a standard element-wise op (negate).
1007 Interpreter interpreter;
1008 interpreter.AddTensors(4);
1009 interpreter.SetInputs({0, 1});
1010 interpreter.SetOutputs({3});
1011 TfLiteQuantizationParams quant;
1012 interpreter.SetTensorParametersReadWrite(0, kTfLiteFloat32, "", {2, 2, 1, 1},
1013 quant);
1014 interpreter.SetTensorParametersReadWrite(1, kTfLiteInt32, "", {4, 2}, quant);
1015 interpreter.SetTensorParametersReadWrite(2, kTfLiteFloat32, "", {}, quant);
1016 interpreter.SetTensorParametersReadWrite(3, kTfLiteFloat32, "", {}, quant);
1017
1018 TfLiteRegistration* pad_op = tflite::ops::builtin::Register_PADV2();
1019 TfLiteRegistration* neg_op = tflite::ops::builtin::Register_NEG();
1020 interpreter.AddNodeWithParameters({0, 1}, {2}, nullptr, 0, nullptr, pad_op);
1021 interpreter.AddNodeWithParameters({2}, {3}, nullptr, 0, nullptr, neg_op);
1022 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
1023
1024 // Configure [[2,2],[4,4]] padding and execute the graph.
1025 interpreter.typed_tensor<int>(1)[0] = 2;
1026 interpreter.typed_tensor<int>(1)[1] = 2;
1027 interpreter.typed_tensor<int>(1)[2] = 2;
1028 interpreter.typed_tensor<int>(1)[3] = 2;
1029 interpreter.typed_tensor<int>(1)[4] = 0;
1030 interpreter.typed_tensor<int>(1)[5] = 0;
1031 interpreter.typed_tensor<int>(1)[6] = 0;
1032 interpreter.typed_tensor<int>(1)[7] = 0;
1033 ASSERT_EQ(interpreter.Invoke(), kTfLiteOk);
1034
1035 // Both the output and intermediate tensor sizes should reflect the output
1036 // from the dynamic pad operation.
1037 ASSERT_EQ(interpreter.tensor(2)->bytes, sizeof(float) * 6 * 6);
1038 ASSERT_EQ(interpreter.tensor(3)->bytes, sizeof(float) * 6 * 6);
1039
1040 // Now configure [[4,4],[6,6]] padding and execute the graph.
1041 interpreter.typed_tensor<int>(1)[0] = 4;
1042 interpreter.typed_tensor<int>(1)[1] = 4;
1043 interpreter.typed_tensor<int>(1)[2] = 6;
1044 interpreter.typed_tensor<int>(1)[3] = 6;
1045 interpreter.typed_tensor<int>(1)[4] = 0;
1046 interpreter.typed_tensor<int>(1)[5] = 0;
1047 interpreter.typed_tensor<int>(1)[6] = 0;
1048 interpreter.typed_tensor<int>(1)[7] = 0;
1049 ASSERT_EQ(interpreter.Invoke(), kTfLiteOk);
1050
1051 // Again, the output and intermediate tensor sizes should reflect the *new*
1052 // resize from the latest pad operation.
1053 ASSERT_EQ(interpreter.tensor(2)->bytes, sizeof(float) * 10 * 14);
1054 ASSERT_EQ(interpreter.tensor(3)->bytes, sizeof(float) * 10 * 14);
1055 }
1056
TEST(InterpreterTensorsCapacityTest,TestWithinHeadroom)1057 TEST(InterpreterTensorsCapacityTest, TestWithinHeadroom) {
1058 Interpreter interpreter;
1059 ASSERT_EQ(interpreter.AddTensors(Interpreter::kTensorsReservedCapacity),
1060 kTfLiteOk);
1061 TfLiteRegistration registration = {nullptr, nullptr, nullptr, nullptr};
1062 registration.prepare = [](TfLiteContext* context, TfLiteNode* node) {
1063 TfLiteTensor* first_tensor = context->tensors;
1064
1065 int new_tensor_index;
1066 context->AddTensors(context, Interpreter::kTensorsCapacityHeadroom,
1067 &new_tensor_index);
1068 EXPECT_EQ(first_tensor, context->tensors);
1069 return kTfLiteOk;
1070 };
1071 ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr,
1072 ®istration),
1073 kTfLiteOk);
1074 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
1075 }
1076
TEST(InterpreterTensorsCapacityTest,TestExceedHeadroom)1077 TEST(InterpreterTensorsCapacityTest, TestExceedHeadroom) {
1078 Interpreter interpreter;
1079 ASSERT_EQ(interpreter.AddTensors(Interpreter::kTensorsReservedCapacity),
1080 kTfLiteOk);
1081 TfLiteRegistration registration = {nullptr, nullptr, nullptr, nullptr};
1082 registration.prepare = [](TfLiteContext* context, TfLiteNode* node) {
1083 TfLiteTensor* first_tensor = context->tensors;
1084
1085 int new_tensor_index;
1086 // Add enough tensors to trigger buffer re-allocation.
1087 context->AddTensors(
1088 context,
1089 (context->tensors_size + Interpreter::kTensorsCapacityHeadroom + 1) * 2,
1090 &new_tensor_index);
1091 EXPECT_NE(first_tensor, context->tensors);
1092 return kTfLiteOk;
1093 };
1094 ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr,
1095 ®istration),
1096 kTfLiteOk);
1097 ASSERT_EQ(interpreter.AllocateTensors(), kTfLiteOk);
1098 }
1099
1100 struct TestExternalContext : public TfLiteExternalContext {
1101 static constexpr TfLiteExternalContextType kType = kTfLiteGemmLowpContext;
1102
Gettflite::__anona9bbd6c50111::TestExternalContext1103 static TestExternalContext* Get(TfLiteContext* context) {
1104 return reinterpret_cast<TestExternalContext*>(
1105 context->GetExternalContext(context, kType));
1106 }
1107
Settflite::__anona9bbd6c50111::TestExternalContext1108 static void Set(TfLiteContext* context, TestExternalContext* value) {
1109 context->SetExternalContext(context, kType, value);
1110 }
1111
1112 int num_refreshes = 0;
1113 };
1114
TEST_F(InterpreterTest,GetSetResetExternalContexts)1115 TEST_F(InterpreterTest, GetSetResetExternalContexts) {
1116 auto* context = GetInterpreterContext();
1117
1118 TestExternalContext external_context;
1119 external_context.Refresh = [](TfLiteContext* context) {
1120 auto* ptr = TestExternalContext::Get(context);
1121 if (ptr != nullptr) {
1122 ++ptr->num_refreshes;
1123 }
1124 return kTfLiteOk;
1125 };
1126
1127 EXPECT_EQ(TestExternalContext::Get(context), nullptr);
1128 ASSERT_EQ(interpreter_.SetNumThreads(4), kTfLiteOk);
1129
1130 TestExternalContext::Set(context, &external_context);
1131 EXPECT_EQ(TestExternalContext::Get(context), &external_context);
1132 ASSERT_EQ(interpreter_.SetNumThreads(4), kTfLiteOk);
1133 ASSERT_EQ(interpreter_.SetNumThreads(5), kTfLiteOk);
1134 EXPECT_EQ(external_context.num_refreshes, 2);
1135
1136 // Reset refresh count to 0
1137 external_context.num_refreshes = 0;
1138 // Below should not call external context refresh
1139 ASSERT_EQ(interpreter_.SetNumThreads(-2), kTfLiteError);
1140 EXPECT_EQ(external_context.num_refreshes, 0);
1141
1142 ASSERT_EQ(interpreter_.SetNumThreads(-1), kTfLiteOk);
1143 EXPECT_EQ(external_context.num_refreshes, 1);
1144
1145 TestExternalContext::Set(context, nullptr);
1146 EXPECT_EQ(TestExternalContext::Get(context), nullptr);
1147 ASSERT_EQ(interpreter_.SetNumThreads(4), kTfLiteOk);
1148 }
1149
TEST_F(InterpreterTest,SetNumThreadsSucceedsWithZero)1150 TEST_F(InterpreterTest, SetNumThreadsSucceedsWithZero) {
1151 ASSERT_EQ(interpreter_.SetNumThreads(0), kTfLiteOk);
1152 // num_threads == 0 has the same effect as num_threads == 1.
1153 EXPECT_EQ(interpreter_.subgraph(0)->context()->recommended_num_threads, 1);
1154 }
1155
1156 struct TestCpuBackendContext : public TfLiteInternalBackendContext {
1157 // Count the number of calls to ClearCaches for the backend context.
ClearCachestflite::__anona9bbd6c50111::TestCpuBackendContext1158 void ClearCaches() override { ++num_calls; }
SetMaxNumThreadstflite::__anona9bbd6c50111::TestCpuBackendContext1159 void SetMaxNumThreads(int num_threads) override {}
1160 int num_calls = 0;
1161 };
1162
TEST_F(InterpreterTest,ExternalBackendContextClearsCachesOnDelete)1163 TEST_F(InterpreterTest, ExternalBackendContextClearsCachesOnDelete) {
1164 ExternalCpuBackendContext external_cpu_context;
1165 TestCpuBackendContext* cpu_backend_context = new TestCpuBackendContext();
1166 external_cpu_context.set_internal_backend_context(
1167 std::unique_ptr<TfLiteInternalBackendContext>(cpu_backend_context));
1168
1169 {
1170 // Create an interpreter with an external Cpu backend context and ensure
1171 // it goes out of scope.
1172 Interpreter interpreter;
1173 interpreter.SetExternalContext(kTfLiteCpuBackendContext,
1174 &external_cpu_context);
1175 EXPECT_EQ(cpu_backend_context->num_calls, 0);
1176 }
1177 EXPECT_EQ(cpu_backend_context->num_calls, 1);
1178 }
1179
1180 // Test fixture that allows playing with execution plans. It creates a two
1181 // node graph that can be executed in either [0,1] order or [1,0] order.
1182 // The CopyOp records when it is invoked in the class member run_order_
1183 // so we can test whether the execution plan was honored.
1184 class TestExecutionPlan : public InterpreterTest {
1185 // Encapsulates the node ids and provides them to a C primitive data type
1186 // Allocatable with placement new, but never destructed, so make sure this
1187 // doesn't own any heap allocated data. This is then is used as op local
1188 // data to allow access to the test fixture data.
1189 class CallReporting {
1190 public:
CallReporting(int node_id,std::vector<int> * run_order)1191 CallReporting(int node_id, std::vector<int>* run_order)
1192 : node_id_(node_id), run_order_(run_order) {}
1193
Record()1194 void Record() { run_order_->push_back(node_id_); }
1195
1196 private:
1197 // The node id for this particular node
1198 int node_id_;
1199 // A pointer to the global run-order
1200 std::vector<int>* run_order_;
1201 };
1202
1203 // Build a kernel registration for an op that copies its one input
1204 // to an output
CopyOpRegistration()1205 TfLiteRegistration CopyOpRegistration() {
1206 TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
1207
1208 reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
1209 // Set output size to input size
1210 const TfLiteTensor* tensor0;
1211 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &tensor0));
1212 TfLiteTensor* tensor1;
1213 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &tensor1));
1214 TfLiteIntArray* newSize = TfLiteIntArrayCopy(tensor0->dims);
1215 return context->ResizeTensor(context, tensor1, newSize);
1216 };
1217
1218 reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
1219 CallReporting* call_reporting =
1220 static_cast<CallReporting*>(node->builtin_data);
1221 // Copy input data to output data.
1222 const TfLiteTensor* a0;
1223 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &a0));
1224 TfLiteTensor* a1;
1225 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &a1));
1226 int num = a0->dims->data[0];
1227 for (int i = 0; i < num; i++) {
1228 a1->data.f[i] = a0->data.f[i];
1229 }
1230 call_reporting->Record();
1231 return kTfLiteOk;
1232 };
1233 return reg;
1234 }
1235
1236 // Adds a copy node going from tensor `input` to output tensor `output`.
1237 // Note, input is used as the node_id. Inject run_order as op accessible
1238 // data. Note: this is a little strange of a way to do this, but it is
1239 // using op functionality to avoid static global variables.
MakeCopyNode(int input,int output)1240 void MakeCopyNode(int input, int output) {
1241 // Ownership of call_reporting is taken by interpreter (malloc is used due
1242 // to nodes being a C99 interface so free() is used).
1243 TfLiteRegistration copy_op = CopyOpRegistration();
1244 CallReporting* call_reporting_1 =
1245 static_cast<CallReporting*>(malloc(sizeof(CallReporting)));
1246 new (call_reporting_1) CallReporting(input, &run_order_);
1247 ASSERT_EQ(interpreter_.AddNodeWithParameters(
1248 {0}, {2}, nullptr, 0, static_cast<void*>(call_reporting_1),
1249 ©_op),
1250 kTfLiteOk);
1251 ASSERT_EQ(interpreter_.ResizeInputTensor(input, {3}), kTfLiteOk);
1252 }
1253
SetUp()1254 void SetUp() final {
1255 // Add two inputs and two outputs that don't depend on each other
1256 ASSERT_EQ(interpreter_.AddTensors(4), kTfLiteOk);
1257 interpreter_.SetInputs({0, 1});
1258 interpreter_.SetOutputs({2, 3});
1259 TfLiteQuantizationParams quantized;
1260 for (int tensor_index = 0; tensor_index < 4; tensor_index++) {
1261 ASSERT_EQ(interpreter_.SetTensorParametersReadWrite(
1262 tensor_index, kTfLiteFloat32, "", {3}, quantized),
1263 kTfLiteOk);
1264 }
1265
1266 // Define two copy functions that also use the user_data to report that
1267 // they were called.
1268 // i.e. tensor[2] = copy(tensor[0]); tensor[3] = copy(tensor[1]);
1269 // thus we can reorder the two nodes arbitrary and still satisfy dependency
1270 // order.
1271 MakeCopyNode(0, 2);
1272 MakeCopyNode(1, 3);
1273
1274 ASSERT_EQ(interpreter_.AllocateTensors(), kTfLiteOk);
1275 }
1276
1277 protected:
1278 // list of node_ids that were run
1279 std::vector<int> run_order_;
1280 };
1281
TEST_F(TestExecutionPlan,DefaultExecutionPlan)1282 TEST_F(TestExecutionPlan, DefaultExecutionPlan) {
1283 // Check default order
1284 ASSERT_EQ(interpreter_.Invoke(), kTfLiteOk);
1285 ASSERT_EQ(run_order_, std::vector<int>({0, 1}));
1286 }
1287
TEST_F(TestExecutionPlan,ReversedExecutionPlan)1288 TEST_F(TestExecutionPlan, ReversedExecutionPlan) {
1289 // Check reversed order
1290 SetExecutionPlan({1, 0});
1291 ASSERT_EQ(interpreter_.Invoke(), kTfLiteOk);
1292 ASSERT_EQ(run_order_, std::vector<int>({1, 0}));
1293 }
1294
TEST_F(TestExecutionPlan,SubsetExecutionPlan)1295 TEST_F(TestExecutionPlan, SubsetExecutionPlan) {
1296 // Check running only node index 1
1297 SetExecutionPlan({1});
1298 ASSERT_EQ(interpreter_.Invoke(), kTfLiteOk);
1299 ASSERT_EQ(run_order_, std::vector<int>({1}));
1300 }
1301
TEST_F(TestExecutionPlan,NullExecutionPlan)1302 TEST_F(TestExecutionPlan, NullExecutionPlan) {
1303 // Check nothing executed.
1304 SetExecutionPlan({});
1305 ASSERT_EQ(interpreter_.Invoke(), kTfLiteOk);
1306 ASSERT_EQ(run_order_, std::vector<int>());
1307 }
1308
TEST(TestDelegateOwnership,ProperlyDisposed)1309 TEST(TestDelegateOwnership, ProperlyDisposed) {
1310 struct TfLiteInterpreterOwnedDelegate : public TfLiteDelegate {
1311 TfLiteInterpreterOwnedDelegate(bool* destroyed, bool* prepared)
1312 : destroyed(destroyed), prepared(prepared) {
1313 flags = kTfLiteDelegateFlagsNone;
1314 Prepare = [](TfLiteContext*, TfLiteDelegate* delegate) -> TfLiteStatus {
1315 *static_cast<TfLiteInterpreterOwnedDelegate*>(delegate)->prepared =
1316 true;
1317 return kTfLiteOk;
1318 };
1319 }
1320 ~TfLiteInterpreterOwnedDelegate() { *destroyed = true; }
1321
1322 bool* destroyed;
1323 bool* prepared;
1324 };
1325
1326 // Construct a delegate with flags for indicating preparation/destruction.
1327 bool destroyed = false;
1328 bool prepared = false;
1329 std::unique_ptr<TfLiteInterpreterOwnedDelegate> delegate(
1330 new TfLiteInterpreterOwnedDelegate(&destroyed, &prepared));
1331 {
1332 // Create an interpreter and assemble a simple graph.
1333 Interpreter interpreter;
1334 TfLiteRegistration registration = {nullptr, nullptr, nullptr, nullptr};
1335 ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk);
1336 ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk);
1337 ASSERT_EQ(interpreter.SetOutputs({1}), kTfLiteOk);
1338 ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr,
1339 ®istration),
1340 kTfLiteOk);
1341
1342 // Pass delegate ownership to that interpreter.
1343 ASSERT_EQ(InterpreterTest::ModifyGraphWithDelegate(&interpreter,
1344 std::move(delegate)),
1345 kTfLiteOk);
1346
1347 // The delegate should be prepared as normal, and should be preserved.
1348 EXPECT_TRUE(prepared);
1349 EXPECT_FALSE(destroyed);
1350
1351 // Interpreter interaction should not impact the delegate's validity.
1352 interpreter.AllocateTensors();
1353 interpreter.Invoke();
1354 EXPECT_FALSE(destroyed);
1355 }
1356
1357 // Only after the interpreter is destroyed should the delegate be destroyed.
1358 EXPECT_TRUE(destroyed);
1359 }
1360
1361 // CancellationData contains the data required to cancel a call to Invoke().
1362 struct CancellationData {
1363 bool is_cancelled = false;
1364 };
1365
1366 // Indicates whether Invoke() has been cancelled based on the value of the
1367 // CancellationData object passed in.
CheckCancellation(void * data)1368 bool CheckCancellation(void* data) {
1369 CancellationData* cancellation_data =
1370 static_cast<struct CancellationData*>(data);
1371 return cancellation_data->is_cancelled;
1372 }
1373
1374 static struct CancellationData cancellation_data_;
1375
1376 // Test fixture to test cancellation within the Interpreter.
1377 class CancellationTest : public ::testing::Test {
1378 public:
Invoke()1379 TfLiteStatus Invoke() { return interpreter_.Invoke(); }
Cancel()1380 void Cancel() { cancellation_data_.is_cancelled = true; }
1381
1382 // Adds an CancelOp with input tensor `input` and output tensor `output`.
MakeCancelNode(int input,int output)1383 void MakeCancelNode(int input, int output) {
1384 TfLiteRegistration op = CancelOpRegistration();
1385 ASSERT_EQ(interpreter_.AddNodeWithParameters({input}, {output}, nullptr, 0,
1386 nullptr, &op),
1387 kTfLiteOk);
1388 ASSERT_EQ(interpreter_.ResizeInputTensor(input, {3}), kTfLiteOk);
1389 }
1390
1391 // Adds an OkOp with input tensor `input` and output tensor `output`.
MakeOkNode(int input,int output)1392 void MakeOkNode(int input, int output) {
1393 TfLiteRegistration op = OkOpRegistration();
1394 ASSERT_EQ(interpreter_.AddNodeWithParameters({input}, {output}, nullptr, 0,
1395 nullptr, &op),
1396 kTfLiteOk);
1397 ASSERT_EQ(interpreter_.ResizeInputTensor(input, {3}), kTfLiteOk);
1398 }
1399
1400 Interpreter interpreter_;
1401
1402 private:
1403 // Build the kernel registration for an op that cancels the operation.
CancelOpRegistration()1404 TfLiteRegistration CancelOpRegistration() {
1405 TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
1406
1407 // Set output size to the input size in CancelOp::Prepare(). Code exists to
1408 // have a framework in Prepare. The input and output tensors are not used.
1409 reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
1410 const TfLiteTensor* in_tensor;
1411 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &in_tensor));
1412 TfLiteTensor* out_tensor;
1413 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &out_tensor));
1414 TfLiteIntArray* new_size = TfLiteIntArrayCopy(in_tensor->dims);
1415 return context->ResizeTensor(context, out_tensor, new_size);
1416 };
1417
1418 reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
1419 cancellation_data_.is_cancelled = true;
1420 return kTfLiteOk;
1421 };
1422 return reg;
1423 }
1424
1425 // Build the kernel registration for an op that returns kTfLiteOk.
OkOpRegistration()1426 TfLiteRegistration OkOpRegistration() {
1427 TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr};
1428
1429 // Set output size to the input size in OkOp::Prepare(). Code exists to have
1430 // a framework in Prepare. The input and output tensors are not used.
1431 reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
1432 const TfLiteTensor* in_tensor;
1433 TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &in_tensor));
1434 TfLiteTensor* out_tensor;
1435 TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &out_tensor));
1436 TfLiteIntArray* new_size = TfLiteIntArrayCopy(in_tensor->dims);
1437 return context->ResizeTensor(context, out_tensor, new_size);
1438 };
1439
1440 reg.invoke = [](TfLiteContext* context, TfLiteNode* node) {
1441 return kTfLiteOk;
1442 };
1443 return reg;
1444 }
1445
SetUp()1446 void SetUp() final {
1447 cancellation_data_.is_cancelled = false;
1448
1449 // Set up the interpreter. Create the input and output tensors.
1450 int num_tensors = 3;
1451 ASSERT_EQ(interpreter_.AddTensors(num_tensors), kTfLiteOk);
1452 interpreter_.SetInputs({0});
1453 interpreter_.SetOutputs({2});
1454 TfLiteQuantizationParams quantized;
1455 for (int tensor_index = 0; tensor_index < num_tensors; tensor_index++) {
1456 ASSERT_EQ(interpreter_.SetTensorParametersReadWrite(
1457 tensor_index, kTfLiteFloat32, "", {3}, quantized),
1458 kTfLiteOk);
1459 }
1460 interpreter_.SetCancellationFunction(&cancellation_data_,
1461 &CheckCancellation);
1462 }
1463 };
1464
TEST_F(CancellationTest,CancelBeforeInvoke)1465 TEST_F(CancellationTest, CancelBeforeInvoke) {
1466 // Cancel prior to calling Invoke.
1467 CancellationTest::MakeOkNode(1, 2);
1468 ASSERT_EQ(interpreter_.AllocateTensors(), kTfLiteOk);
1469
1470 CancellationTest::Cancel();
1471 TfLiteStatus invoke_error_code = CancellationTest::Invoke();
1472 ASSERT_EQ(invoke_error_code, kTfLiteError);
1473 }
1474
TEST_F(CancellationTest,CancelDuringInvoke)1475 TEST_F(CancellationTest, CancelDuringInvoke) {
1476 // Tests a model which sets the cancel in order to test cancellation works
1477 // between ops.
1478 //
1479 // The first op will set the cancellation bit to true. The second op returns
1480 // `kTfLiteOk` if executed.
1481 CancellationTest::MakeCancelNode(0, 1);
1482 CancellationTest::MakeOkNode(1, 2);
1483 ASSERT_EQ(interpreter_.AllocateTensors(), kTfLiteOk);
1484
1485 TfLiteStatus invoke_error_code = CancellationTest::Invoke();
1486 ASSERT_EQ(invoke_error_code, kTfLiteError);
1487 }
1488
1489 // Tests functionality related to custom memory allocations in TFLite.
1490 class TestCustomAllocation : public ::testing::Test {
1491 protected:
SetUp()1492 void SetUp() override {
1493 // Simple model with two custom ops that add 2 float tensors each.
1494 interpreter_.reset(new Interpreter);
1495 interpreter_->AddTensors(7);
1496 interpreter_->SetInputs({0, 1});
1497 interpreter_->SetOutputs({3, 4, 6});
1498 TfLiteQuantizationParams quant;
1499 interpreter_->SetTensorParametersReadWrite(0, kTfLiteFloat32, "", {3},
1500 quant);
1501 interpreter_->SetTensorParametersReadWrite(1, kTfLiteFloat32, "", {3},
1502 quant);
1503 interpreter_->SetTensorParametersReadWrite(2, kTfLiteFloat32, "", {3},
1504 quant);
1505 interpreter_->SetTensorParametersReadWrite(3, kTfLiteFloat32, "", {3},
1506 quant);
1507 interpreter_->SetTensorParametersReadWrite(4, kTfLiteFloat32, "", {3},
1508 quant);
1509 interpreter_->SetTensorParametersReadWrite(5, kTfLiteFloat32, "", {3},
1510 quant, /*is_variable=*/true);
1511 interpreter_->SetTensorParametersReadWrite(6, kTfLiteFloat32, "", {3},
1512 quant);
1513 auto* add_reg = ops::builtin::Register_ADD();
1514 TfLiteAddParams* builtin_data0 =
1515 reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
1516 TfLiteAddParams* builtin_data1 =
1517 reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
1518 TfLiteAddParams* builtin_data2 =
1519 reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
1520 TfLiteAddParams* builtin_data3 =
1521 reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
1522 builtin_data0->activation = kTfLiteActNone;
1523 builtin_data1->activation = kTfLiteActNone;
1524 builtin_data2->activation = kTfLiteActNone;
1525 builtin_data3->activation = kTfLiteActNone;
1526 interpreter_->AddNodeWithParameters({0, 0}, {2}, nullptr, 0, builtin_data0,
1527 add_reg);
1528 interpreter_->AddNodeWithParameters({1, 1}, {3}, nullptr, 0, builtin_data1,
1529 add_reg);
1530 interpreter_->AddNodeWithParameters({2, 1}, {4}, nullptr, 0, builtin_data2,
1531 add_reg);
1532 interpreter_->AddNodeWithParameters({0, 5}, {6}, nullptr, 0, builtin_data3,
1533 add_reg);
1534 interpreter_->SetVariables({5});
1535 }
1536
AssignCustomAllocForTensor(int tensor_idx,int required_alignment)1537 void AssignCustomAllocForTensor(int tensor_idx, int required_alignment) {
1538 const TfLiteTensor* tensor = interpreter_->tensor(tensor_idx);
1539 auto tensor_alloc = NewCustomAlloc(tensor->bytes, required_alignment);
1540 ASSERT_EQ(
1541 interpreter_->SetCustomAllocationForTensor(tensor_idx, tensor_alloc),
1542 kTfLiteOk);
1543 }
1544
VerifyInvoke()1545 void VerifyInvoke() {
1546 std::vector<float> input = {1.0f, 2.0f, 3.0f};
1547 std::vector<float> variable = {0.0f, 1.0f, 2.0f};
1548 std::vector<float> expected_output = {2.0f, 4.0f, 6.0f};
1549
1550 // typed_tensor<...> should work irrespective of custom alloc, since it
1551 // accesses output_tensor.data.
1552 memcpy(interpreter_->typed_tensor<float>(interpreter_->variables()[0]),
1553 variable.data(), 3 * sizeof(float));
1554 memcpy(interpreter_->typed_tensor<float>(0), input.data(),
1555 3 * sizeof(float));
1556 memcpy(interpreter_->typed_tensor<float>(1), input.data(),
1557 3 * sizeof(float));
1558 ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
1559 TfLiteTensor* output_tensor =
1560 interpreter_->tensor(interpreter_->outputs()[0]);
1561 for (int i = 0; i < 3; ++i) {
1562 EXPECT_EQ(output_tensor->data.f[i], expected_output[i]) << i;
1563 }
1564 }
1565
1566 // Actual initialized allocation is more than num_bytes, to account for
1567 // required_allocation.
NewCustomAlloc(size_t num_bytes,int required_alignment)1568 TfLiteCustomAllocation NewCustomAlloc(size_t num_bytes,
1569 int required_alignment) {
1570 // Extra memory to ensure alignment.
1571 char* new_alloc = new char[num_bytes + required_alignment];
1572 char* new_underlying_buffer_aligned_ptr = reinterpret_cast<char*>(
1573 AlignTo(required_alignment, reinterpret_cast<intptr_t>(new_alloc)));
1574 custom_alloc_buffers_.emplace_back(new_alloc);
1575
1576 return TfLiteCustomAllocation(
1577 {new_underlying_buffer_aligned_ptr, num_bytes});
1578 }
1579
AlignTo(size_t alignment,intptr_t offset)1580 intptr_t AlignTo(size_t alignment, intptr_t offset) {
1581 return offset % alignment == 0 ? offset
1582 : offset + (alignment - offset % alignment);
1583 }
1584
TearDown()1585 void TearDown() override {
1586 interpreter_.reset();
1587 custom_alloc_buffers_.clear();
1588 }
1589
1590 protected:
1591 TfLiteAddParams add_params_;
1592 std::unique_ptr<Interpreter> interpreter_;
1593 std::vector<std::unique_ptr<char[]>> custom_alloc_buffers_;
1594 };
1595
TEST_F(TestCustomAllocation,InvalidAlignment)1596 TEST_F(TestCustomAllocation, InvalidAlignment) {
1597 const TfLiteTensor* input_tensor =
1598 interpreter_->tensor(interpreter_->inputs()[0]);
1599 intptr_t dummy_ptr = kDefaultTensorAlignment - 1;
1600 TfLiteCustomAllocation input_alloc{reinterpret_cast<void*>(dummy_ptr),
1601 input_tensor->bytes};
1602 ASSERT_EQ(interpreter_->SetCustomAllocationForTensor(
1603 interpreter_->inputs()[0], input_alloc),
1604 kTfLiteError);
1605
1606 // Allocate tensors & Invoke should still work.
1607 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1608 VerifyInvoke();
1609 }
1610
TEST_F(TestCustomAllocation,InvalidAlignment_SkipCheck)1611 TEST_F(TestCustomAllocation, InvalidAlignment_SkipCheck) {
1612 const TfLiteTensor* input_tensor =
1613 interpreter_->tensor(interpreter_->inputs()[0]);
1614 const int required_alignment = kDefaultTensorAlignment - 1;
1615 auto tensor_alloc = NewCustomAlloc(input_tensor->bytes, required_alignment);
1616 ASSERT_EQ(interpreter_->SetCustomAllocationForTensor(
1617 interpreter_->inputs()[0], tensor_alloc,
1618 /**flags**/ kTfLiteCustomAllocationFlagsSkipAlignCheck),
1619 kTfLiteOk);
1620
1621 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1622 }
1623
TEST_F(TestCustomAllocation,InsufficientBytes)1624 TEST_F(TestCustomAllocation, InsufficientBytes) {
1625 auto input_alloc = NewCustomAlloc(4, kDefaultTensorAlignment);
1626
1627 // Setting the custom alloc works, but AllocateTensors doesn't.
1628 ASSERT_EQ(interpreter_->SetCustomAllocationForTensor(
1629 interpreter_->inputs()[0], input_alloc),
1630 kTfLiteOk);
1631 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteError);
1632 ASSERT_EQ(interpreter_->Invoke(), kTfLiteError);
1633 }
1634
TEST_F(TestCustomAllocation,CustomInputAlloc)1635 TEST_F(TestCustomAllocation, CustomInputAlloc) {
1636 // Set custom allocation for one input tensor.
1637 AssignCustomAllocForTensor(interpreter_->inputs()[0],
1638 /*required_alignment=*/kDefaultTensorAlignment);
1639
1640 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1641 VerifyInvoke();
1642 }
1643
TEST_F(TestCustomAllocation,CustomInputAlloc_MultipleAssigns)1644 TEST_F(TestCustomAllocation, CustomInputAlloc_MultipleAssigns) {
1645 // Set custom allocation for one input tensor.
1646 AssignCustomAllocForTensor(interpreter_->inputs()[0],
1647 /*required_alignment=*/kDefaultTensorAlignment);
1648
1649 AssignCustomAllocForTensor(interpreter_->inputs()[0],
1650 /*required_alignment=*/kDefaultTensorAlignment);
1651 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1652 VerifyInvoke();
1653
1654 AssignCustomAllocForTensor(interpreter_->inputs()[0],
1655 /*required_alignment=*/kDefaultTensorAlignment);
1656 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1657 VerifyInvoke();
1658 }
1659
TEST_F(TestCustomAllocation,CustomInputAlloc_AllocateTensorsBefore)1660 TEST_F(TestCustomAllocation, CustomInputAlloc_AllocateTensorsBefore) {
1661 // Allocate tensors.
1662 // Allocating now will cause TFLite to reserve some extra memory, but nothing
1663 // should break.
1664 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1665
1666 AssignCustomAllocForTensor(interpreter_->inputs()[0],
1667 /*required_alignment=*/kDefaultTensorAlignment);
1668
1669 VerifyInvoke();
1670 }
1671
TEST_F(TestCustomAllocation,CustomInputAndOutputAllocs)1672 TEST_F(TestCustomAllocation, CustomInputAndOutputAllocs) {
1673 // Set custom allocations for all IO tensors.
1674 AssignCustomAllocForTensor(interpreter_->inputs()[0],
1675 /*required_alignment=*/kDefaultTensorAlignment);
1676 AssignCustomAllocForTensor(interpreter_->inputs()[1],
1677 /*required_alignment=*/kDefaultTensorAlignment);
1678 AssignCustomAllocForTensor(interpreter_->outputs()[0],
1679 /*required_alignment=*/kDefaultTensorAlignment);
1680 AssignCustomAllocForTensor(interpreter_->outputs()[1],
1681 /*required_alignment=*/kDefaultTensorAlignment);
1682
1683 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1684 VerifyInvoke();
1685 }
1686
1687 // Ensure that custom allocs work for tensors on persistent arena as well.
TEST_F(TestCustomAllocation,CustomAlloc_VariableTensor)1688 TEST_F(TestCustomAllocation, CustomAlloc_VariableTensor) {
1689 // Set custom allocation for one input tensor.
1690 AssignCustomAllocForTensor(interpreter_->variables()[0],
1691 /*required_alignment=*/kDefaultTensorAlignment);
1692
1693 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1694 VerifyInvoke();
1695
1696 AssignCustomAllocForTensor(interpreter_->variables()[0],
1697 /*required_alignment=*/kDefaultTensorAlignment);
1698 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1699
1700 std::vector<float> input = {2.0f, 3.0f, 4.0f};
1701 std::vector<float> variable = {1.0f, 2.0f, 3.0f};
1702 std::vector<float> expected_output = {3.0f, 5.0f, 7.0f};
1703 memcpy(interpreter_->typed_tensor<float>(interpreter_->variables()[0]),
1704 variable.data(), 3 * sizeof(float));
1705 memcpy(interpreter_->typed_tensor<float>(0), input.data(), 3 * sizeof(float));
1706 memcpy(interpreter_->typed_tensor<float>(1), input.data(), 3 * sizeof(float));
1707 ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
1708
1709 // expected_output = input + variable
1710 TfLiteTensor* output_tensor =
1711 interpreter_->tensor(interpreter_->outputs()[2]);
1712 for (int i = 0; i < 3; ++i) {
1713 EXPECT_EQ(output_tensor->data.f[i], expected_output[i]) << i;
1714 }
1715 }
1716
TEST_F(TestCustomAllocation,ResizeInputsWithoutEnoughMemory)1717 TEST_F(TestCustomAllocation, ResizeInputsWithoutEnoughMemory) {
1718 // Set custom allocations for all input tensors.
1719 AssignCustomAllocForTensor(interpreter_->inputs()[0],
1720 /*required_alignment=*/kDefaultTensorAlignment);
1721 AssignCustomAllocForTensor(interpreter_->inputs()[1],
1722 /*required_alignment=*/kDefaultTensorAlignment);
1723
1724 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1725
1726 // Now resize tensors to double the size.
1727 ASSERT_EQ(interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {2, 3}),
1728 kTfLiteOk);
1729 ASSERT_EQ(interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {2, 3}),
1730 kTfLiteOk);
1731
1732 // Since the custom memory previously allocated isn't enough,
1733 // AllocateTensors() will fail.
1734 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteError);
1735 // Interpreter should no longer be in invokable state, so expect failure.
1736 ASSERT_EQ(interpreter_->Invoke(), kTfLiteError);
1737 }
1738
TEST_F(TestCustomAllocation,ResizeInputsWithEnoughMemory)1739 TEST_F(TestCustomAllocation, ResizeInputsWithEnoughMemory) {
1740 // Set custom allocations for all input tensors, with double the required
1741 // memory.
1742 const TfLiteTensor* input0_tensor =
1743 interpreter_->tensor(interpreter_->inputs()[0]);
1744 auto input0_alloc =
1745 NewCustomAlloc(2 * input0_tensor->bytes, kDefaultTensorAlignment);
1746 ASSERT_EQ(interpreter_->SetCustomAllocationForTensor(
1747 interpreter_->inputs()[0], input0_alloc),
1748 kTfLiteOk);
1749 const TfLiteTensor* input1_tensor =
1750 interpreter_->tensor(interpreter_->inputs()[1]);
1751 auto input1_alloc =
1752 NewCustomAlloc(2 * input1_tensor->bytes, kDefaultTensorAlignment);
1753 ASSERT_EQ(interpreter_->SetCustomAllocationForTensor(
1754 interpreter_->inputs()[1], input1_alloc),
1755 kTfLiteOk);
1756
1757 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1758
1759 // Now resize tensors to double the size.
1760 ASSERT_EQ(interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {6, 1}),
1761 kTfLiteOk);
1762 ASSERT_EQ(interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {6, 1}),
1763 kTfLiteOk);
1764
1765 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1766
1767 std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
1768 std::vector<float> expected_output = {2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f};
1769 TfLiteTensor* tensor = interpreter_->tensor(interpreter_->outputs()[0]);
1770 memcpy(interpreter_->typed_tensor<float>(0), input.data(), 6 * sizeof(float));
1771 memcpy(interpreter_->typed_tensor<float>(1), input.data(), 6 * sizeof(float));
1772 ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
1773 for (int i = 0; i < 6; ++i) {
1774 EXPECT_EQ(tensor->data.f[i], expected_output[i]) << i;
1775 }
1776
1777 ASSERT_EQ(interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {3, 1}),
1778 kTfLiteOk);
1779 ASSERT_EQ(interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {3, 1}),
1780 kTfLiteOk);
1781
1782 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1783 VerifyInvoke();
1784 }
1785
1786 // Verify typical use-cases where tensors are resized & custom allocs need to be
1787 // set for every Invoke().
TEST_F(TestCustomAllocation,ResizeAndAllocateForEveryInvoke)1788 TEST_F(TestCustomAllocation, ResizeAndAllocateForEveryInvoke) {
1789 // First assign exactly sized allocs for all IO tensors.
1790 AssignCustomAllocForTensor(interpreter_->inputs()[0],
1791 /*required_alignment=*/kDefaultTensorAlignment);
1792 AssignCustomAllocForTensor(interpreter_->inputs()[1],
1793 /*required_alignment=*/kDefaultTensorAlignment);
1794 AssignCustomAllocForTensor(interpreter_->outputs()[0],
1795 /*required_alignment=*/kDefaultTensorAlignment);
1796 AssignCustomAllocForTensor(interpreter_->outputs()[1],
1797 /*required_alignment=*/kDefaultTensorAlignment);
1798 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1799
1800 // Now resize inputs to a smaller: {3, 1} to {1, 1}.
1801 // Total alloc sized required now: 1 float == 4 bytes.
1802 ASSERT_EQ(interpreter_->ResizeInputTensor(interpreter_->inputs()[0], {1, 1}),
1803 kTfLiteOk);
1804 ASSERT_EQ(interpreter_->ResizeInputTensor(interpreter_->inputs()[1], {1, 1}),
1805 kTfLiteOk);
1806 // Assign allocs for all I/O tensors.
1807 // Even though the smaller output tensor sizes have not been propagated yet,
1808 // custom allocation works because verification of allocs happens after
1809 // preparing all ops & tensors.
1810 auto input0_alloc =
1811 NewCustomAlloc(/**num_bytes=**/ 4, kDefaultTensorAlignment);
1812 ASSERT_EQ(interpreter_->SetCustomAllocationForTensor(
1813 interpreter_->inputs()[0], input0_alloc),
1814 kTfLiteOk);
1815 auto input1_alloc =
1816 NewCustomAlloc(/**num_bytes=**/ 4, kDefaultTensorAlignment);
1817 ASSERT_EQ(interpreter_->SetCustomAllocationForTensor(
1818 interpreter_->inputs()[1], input1_alloc),
1819 kTfLiteOk);
1820 auto output0_alloc =
1821 NewCustomAlloc(/**num_bytes=**/ 4, kDefaultTensorAlignment);
1822 ASSERT_EQ(interpreter_->SetCustomAllocationForTensor(
1823 interpreter_->outputs()[0], output0_alloc),
1824 kTfLiteOk);
1825 auto output1_alloc =
1826 NewCustomAlloc(/**num_bytes=**/ 4, kDefaultTensorAlignment);
1827 ASSERT_EQ(interpreter_->SetCustomAllocationForTensor(
1828 interpreter_->outputs()[1], output1_alloc),
1829 kTfLiteOk);
1830 // AllocateTensors works.
1831 ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
1832
1833 std::vector<float> input = {2.0f};
1834 std::vector<float> expected_output = {4.0f};
1835 TfLiteTensor* tensor = interpreter_->tensor(interpreter_->outputs()[0]);
1836 memcpy(interpreter_->typed_tensor<float>(0), input.data(), sizeof(float));
1837 memcpy(interpreter_->typed_tensor<float>(1), input.data(), sizeof(float));
1838 ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
1839 EXPECT_EQ(tensor->data.f[0], expected_output[0]);
1840 }
1841
1842 // Tests related to lazy delegate providers that are primarily used for applying
1843 // TfLite delegates by default.
1844 class TestLazyDelegateProvider : public InterpreterTest {
1845 protected:
1846 struct DummyLazyDelegateProvider : public TfLiteDelegate {
DummyLazyDelegateProvidertflite::__anona9bbd6c50111::TestLazyDelegateProvider::DummyLazyDelegateProvider1847 explicit DummyLazyDelegateProvider(int64_t support_flags) {
1848 data_ = static_cast<void*>(this);
1849 flags = support_flags;
1850 Prepare = [](TfLiteContext*, TfLiteDelegate* delegate) -> TfLiteStatus {
1851 return kTfLiteOk;
1852 };
1853 }
1854 };
1855
InitWithLazyDelegate(int64_t delegate_flags,bool create_dyanmic_tensor=false,bool return_error=false)1856 void InitWithLazyDelegate(int64_t delegate_flags,
1857 bool create_dyanmic_tensor = false,
1858 bool return_error = false) {
1859 TfLiteRegistration reg = {nullptr};
1860 if (return_error) {
1861 reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
1862 return kTfLiteError;
1863 };
1864 }
1865 ASSERT_EQ(interpreter_.AddTensors(2), kTfLiteOk);
1866 interpreter_.SetInputs({0});
1867 interpreter_.SetOutputs({1});
1868 interpreter_.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr, ®);
1869
1870 Interpreter::TfLiteDelegatePtr delegate(
1871 new DummyLazyDelegateProvider(delegate_flags),
1872 [](TfLiteDelegate* delegate) {
1873 auto* dummy =
1874 static_cast<DummyLazyDelegateProvider*>(delegate->data_);
1875 delete dummy;
1876 });
1877 mutable_lazy_delegate_providers()->push_back(std::move(delegate));
1878
1879 if (create_dyanmic_tensor) {
1880 // Mark the output as dynamic tensor.
1881 interpreter_.tensor(1)->data.raw = nullptr;
1882 interpreter_.tensor(1)->allocation_type = kTfLiteDynamic;
1883 }
1884 }
1885 };
1886
TEST_F(TestLazyDelegateProvider,ApplicationSuccess)1887 TEST_F(TestLazyDelegateProvider, ApplicationSuccess) {
1888 InitWithLazyDelegate(kTfLiteDelegateFlagsNone);
1889 EXPECT_EQ(kTfLiteOk, interpreter_.AllocateTensors());
1890 // We clear Interpreter::lazy_delegate_providers_ after they are tried out.
1891 EXPECT_TRUE(mutable_lazy_delegate_providers()->empty());
1892 EXPECT_TRUE(HasDelegates());
1893 }
1894
TEST_F(TestLazyDelegateProvider,ApplicationFailure)1895 TEST_F(TestLazyDelegateProvider, ApplicationFailure) {
1896 InitWithLazyDelegate(kTfLiteDelegateFlagsNone,
1897 false /* create_dyanmic_tensor */,
1898 true /* return_error */);
1899 EXPECT_EQ(kTfLiteError, interpreter_.AllocateTensors());
1900 // We clear Interpreter::lazy_delegate_providers_ after they are tried out.
1901 EXPECT_TRUE(mutable_lazy_delegate_providers()->empty());
1902 EXPECT_FALSE(HasDelegates());
1903 }
1904
TEST_F(TestLazyDelegateProvider,ApplicationSkipped)1905 TEST_F(TestLazyDelegateProvider, ApplicationSkipped) {
1906 InitWithLazyDelegate(kTfLiteDelegateFlagsNone,
1907 true /* create_dyanmic_tensor */);
1908 EXPECT_EQ(kTfLiteOk, interpreter_.AllocateTensors());
1909 EXPECT_TRUE(mutable_lazy_delegate_providers()->empty());
1910 // As the delegate doesn't allow dynamic tensor, the delegate won't be applied
1911 // and the interpreter doesn't have any delegate applied.
1912 EXPECT_FALSE(HasDelegates());
1913 }
1914
TEST_F(InterpreterTest,SingleSignature_get_signatures)1915 TEST_F(InterpreterTest, SingleSignature_get_signatures) {
1916 const char kSignatureKey[] = "test_method";
1917 BuildSignature(kSignatureKey, {{"Input1", 0}, {"Input2", 1}},
1918 {{"Output1", 5}});
1919 auto results = interpreter_.signature_keys();
1920 ASSERT_EQ(1, results.size());
1921 EXPECT_EQ(kSignatureKey, *results[0]);
1922 }
1923
TEST_F(InterpreterTest,SingleSignature_get_inputs)1924 TEST_F(InterpreterTest, SingleSignature_get_inputs) {
1925 const char kSignatureKey[] = "test_method";
1926 const std::map<std::string, uint32_t> inputs = {{"Input1", 0}, {"Input2", 1}};
1927 const std::map<std::string, uint32_t> outputs = {{"Output1", 5}};
1928 BuildSignature(kSignatureKey, inputs, outputs);
1929 EXPECT_THAT(interpreter_.signature_inputs(kSignatureKey),
1930 testing::Eq(inputs));
1931 EXPECT_THAT(interpreter_.signature_outputs(kSignatureKey),
1932 testing::Eq(outputs));
1933 }
1934
TEST_F(InterpreterTest,SingleSignature_validate_get_tensor)1935 TEST_F(InterpreterTest, SingleSignature_validate_get_tensor) {
1936 const char kSignatureKey[] = "test_method";
1937 const std::map<std::string, uint32_t> inputs = {{"Input1", 0}, {"Input2", 1}};
1938 const std::map<std::string, uint32_t> outputs = {{"Output1", 5}};
1939
1940 BuildSignature(kSignatureKey, inputs, outputs);
1941 ASSERT_EQ(interpreter_.AddTensors(6), kTfLiteOk);
1942 ASSERT_EQ(interpreter_.SetInputs({0, 1}), kTfLiteOk);
1943 ASSERT_EQ(interpreter_.SetOutputs({5}), kTfLiteOk);
1944 ASSERT_EQ(interpreter_.SetTensorParametersReadWrite(
1945 0, kTfLiteFloat32, "", {3}, TfLiteQuantizationParams()),
1946 kTfLiteOk);
1947 ASSERT_EQ(interpreter_.SetTensorParametersReadWrite(
1948 1, kTfLiteFloat32, "", {3}, TfLiteQuantizationParams()),
1949 kTfLiteOk);
1950 ASSERT_EQ(interpreter_.ResizeInputTensor(interpreter_.inputs()[0], {1, 2, 3}),
1951 kTfLiteOk);
1952 ASSERT_EQ(interpreter_.ResizeInputTensor(interpreter_.inputs()[1], {1, 2, 3}),
1953 kTfLiteOk);
1954 ASSERT_EQ(interpreter_.AllocateTensors(), kTfLiteOk);
1955
1956 EXPECT_TRUE(interpreter_.input_tensor_by_signature("Input1", kSignatureKey) !=
1957 nullptr);
1958 EXPECT_TRUE(interpreter_.input_tensor_by_signature("Input2", kSignatureKey) !=
1959 nullptr);
1960 EXPECT_TRUE(interpreter_.output_tensor_by_signature(
1961 "Output1", kSignatureKey) != nullptr);
1962
1963 // Invalid tensor
1964 EXPECT_EQ(interpreter_.input_tensor_by_signature("Input3", kSignatureKey),
1965 nullptr);
1966 EXPECT_EQ(interpreter_.output_tensor_by_signature("Input3", kSignatureKey),
1967 nullptr);
1968 // Invalid method
1969 EXPECT_EQ(interpreter_.input_tensor_by_signature("Input1", "InvalidMethod"),
1970 nullptr);
1971 EXPECT_EQ(interpreter_.output_tensor_by_signature("Output1", "InvalidMethod"),
1972 nullptr);
1973 }
1974
1975 } // namespace
1976 } // namespace tflite
1977