• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
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                                                 &registration),
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                                                 &registration),
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, &reg);
356   interpreter.AddNodeWithParameters({2, 1}, {4, 5}, nullptr, 0, nullptr, &reg);
357   interpreter.AddNodeWithParameters({4, 3}, {6, 7}, nullptr, 0, nullptr, &reg);
358   interpreter.AddNodeWithParameters({6, 5}, {8}, nullptr, 0, nullptr, &reg);
359   interpreter.AddNodeWithParameters({8, 7}, {9}, nullptr, 0, nullptr, &reg);
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, &reg);
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, &reg),
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, &reg),
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                                               &reg_copy),
812             kTfLiteOk);
813   ASSERT_EQ(interpreter.AddNodeWithParameters({1}, {2}, nullptr, 0, nullptr,
814                                               &reg_len),
815             kTfLiteOk);
816   ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {3}, nullptr, 0, nullptr,
817                                               &reg_copy),
818             kTfLiteOk);
819   ASSERT_EQ(interpreter.AddNodeWithParameters({3}, {4}, nullptr, 0, nullptr,
820                                               &reg_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, &reg),
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, &registration),
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                                               &registration),
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                                               &registration),
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                                               &registration),
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                   &copy_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                                                 &registration),
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, &reg);
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