1 /* Copyright 2019 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/kernels/subgraph_test_util.h"
17
18 #include <stddef.h>
19 #include <stdint.h>
20 #include <stdlib.h>
21
22 #include <random>
23 #include <vector>
24
25 #include <gtest/gtest.h>
26 #include "tensorflow/lite/builtin_ops.h"
27 #include "tensorflow/lite/c/builtin_op_data.h"
28 #include "tensorflow/lite/c/common.h"
29 #include "tensorflow/lite/core/subgraph.h"
30 #include "tensorflow/lite/kernels/builtin_op_kernels.h"
31 #include "tensorflow/lite/kernels/kernel_util.h"
32 #include "tensorflow/lite/string_util.h"
33
34 namespace tflite {
35
36 // Forward declaration for op kernels.
37 namespace ops {
38 namespace custom {
39 namespace random_int {
40
Prepare(TfLiteContext * context,TfLiteNode * node)41 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
42 TF_LITE_ENSURE_EQ(context, NumInputs(node), 0);
43 TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
44
45 TfLiteTensor* output = GetOutput(context, node, 0);
46 TfLiteIntArray* outputSize = TfLiteIntArrayCreate(1);
47 outputSize->data[0] = 1;
48 return context->ResizeTensor(context, output, outputSize);
49 }
50
Eval(TfLiteContext * context,TfLiteNode * node)51 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
52 TfLiteTensor& output = context->tensors[node->outputs->data[0]];
53
54 std::random_device rd;
55 std::uniform_int_distribution<int> dist(1, 32768);
56 output.data.i32[0] = dist(rd);
57 return kTfLiteOk;
58 }
59
60 } // namespace random_int
61
Register_RANDOM_INT()62 TfLiteRegistration* Register_RANDOM_INT() {
63 static TfLiteRegistration r = {nullptr, nullptr, random_int::Prepare,
64 random_int::Eval};
65 return &r;
66 }
67
68 } // namespace custom
69 } // namespace ops
70
71 namespace subgraph_test_util {
72
73 namespace {
74
SetupTensor(Subgraph * subgraph,int tensor_index,TfLiteType type)75 void SetupTensor(Subgraph* subgraph, int tensor_index, TfLiteType type) {
76 ASSERT_EQ(subgraph->SetTensorParametersReadWrite(tensor_index, type, "", 0,
77 nullptr, {}, false),
78 kTfLiteOk);
79 }
80
81 } // namespace
82
~SubgraphBuilder()83 SubgraphBuilder::~SubgraphBuilder() {
84 for (auto buffer : buffers_) {
85 free(buffer);
86 }
87 }
88
BuildAddSubgraph(Subgraph * subgraph)89 void SubgraphBuilder::BuildAddSubgraph(Subgraph* subgraph) {
90 const int kInput1 = 0;
91 const int kInput2 = 1;
92 const int kOutput = 2;
93 const int kTensorCount = 3;
94 // kInput1(0) --> +---+
95 // |ADD| --> kOutput(2)
96 // kInput2(1) --> +---+
97
98 int first_new_tensor_index;
99 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
100 kTfLiteOk);
101 ASSERT_EQ(first_new_tensor_index, 0);
102 ASSERT_EQ(subgraph->SetInputs({kInput1, kInput2}), kTfLiteOk);
103 ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk);
104
105 SetupTensor(subgraph, kInput1, kTfLiteInt32);
106 SetupTensor(subgraph, kInput2, kTfLiteInt32);
107 SetupTensor(subgraph, kOutput, kTfLiteInt32);
108
109 TfLiteAddParams* params =
110 reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
111 params->activation = kTfLiteActNone;
112 auto* add_reg = ops::builtin::Register_ADD();
113 add_reg->builtin_code = kTfLiteBuiltinAdd;
114 int node_index;
115 subgraph->AddNodeWithParameters({kInput1, kInput2}, {kOutput}, {}, nullptr, 0,
116 params, add_reg, &node_index);
117 }
118
119 // Build a subgraph with an mul op. Helper function for testing.
BuildMulSubgraph(Subgraph * subgraph)120 void SubgraphBuilder::BuildMulSubgraph(Subgraph* subgraph) {
121 const int kInput1 = 0;
122 const int kInput2 = 1;
123 const int kOutput = 2;
124 const int kTensorCount = 3;
125 // kInput1(0) --> +---+
126 // |MUL| --> kOutput(2)
127 // kInput2(1) --> +---+
128
129 int first_new_tensor_index;
130 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
131 kTfLiteOk);
132 ASSERT_EQ(first_new_tensor_index, 0);
133 ASSERT_EQ(subgraph->SetInputs({kInput1, kInput2}), kTfLiteOk);
134 ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk);
135
136 SetupTensor(subgraph, kInput1, kTfLiteInt32);
137 SetupTensor(subgraph, kInput2, kTfLiteInt32);
138 SetupTensor(subgraph, kOutput, kTfLiteInt32);
139
140 TfLiteMulParams* params =
141 reinterpret_cast<TfLiteMulParams*>(malloc(sizeof(TfLiteMulParams)));
142 params->activation = kTfLiteActNone;
143 auto* mul_reg = ops::builtin::Register_MUL();
144 mul_reg->builtin_code = kTfLiteBuiltinMul;
145 int node_index;
146 subgraph->AddNodeWithParameters({kInput1, kInput2}, {kOutput}, {}, nullptr, 0,
147 params, mul_reg, &node_index);
148 }
149
150 // Build a subgraph with a pad op. Helper function for testing.
BuildPadSubgraph(Subgraph * subgraph)151 void SubgraphBuilder::BuildPadSubgraph(Subgraph* subgraph) {
152 const int kInput1 = 0;
153 const int kInput2 = 1;
154 const int kOutput = 2;
155 const int kTensorCount = 3;
156 // kInput1(0) --> +---+
157 // |PAD| --> kOutput(2)
158 // kInput2(1) --> +---+
159
160 int first_new_tensor_index;
161 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
162 kTfLiteOk);
163 ASSERT_EQ(first_new_tensor_index, 0);
164 ASSERT_EQ(subgraph->SetInputs({kInput1, kInput2}), kTfLiteOk);
165 ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk);
166
167 SetupTensor(subgraph, kInput1, kTfLiteInt32);
168 SetupTensor(subgraph, kInput2, kTfLiteInt32);
169 SetupTensor(subgraph, kOutput, kTfLiteInt32);
170
171 TfLitePadParams* params =
172 reinterpret_cast<TfLitePadParams*>(malloc(sizeof(TfLitePadParams)));
173 auto* pad_reg = ops::builtin::Register_PAD();
174 pad_reg->builtin_code = kTfLiteBuiltinPad;
175 int node_index;
176 subgraph->AddNodeWithParameters({kInput1, kInput2}, {kOutput}, {}, nullptr, 0,
177 params, pad_reg, &node_index);
178 }
179
BuildIfSubgraph(Subgraph * subgraph)180 void SubgraphBuilder::BuildIfSubgraph(Subgraph* subgraph) {
181 const int kCondInput = 0;
182 const int kInput1 = 1;
183 const int kInput2 = 2;
184 const int kOutput = 3;
185 const int kTensorCount = 4;
186
187 // kCondInput(0) --> +----+
188 // kInput1(1) ----> | IF | --> kOutput(3)
189 // kInput2(2) ----> +----+
190
191 int first_new_tensor_index;
192 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
193 kTfLiteOk);
194 ASSERT_EQ(first_new_tensor_index, 0);
195 ASSERT_EQ(subgraph->SetInputs({kCondInput, kInput1, kInput2}), kTfLiteOk);
196 ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk);
197
198 SetupTensor(subgraph, kCondInput, kTfLiteBool);
199 SetupTensor(subgraph, kInput1, kTfLiteInt32);
200 SetupTensor(subgraph, kInput2, kTfLiteInt32);
201 SetupTensor(subgraph, kOutput, kTfLiteInt32);
202
203 TfLiteIfParams* params =
204 reinterpret_cast<TfLiteIfParams*>(malloc(sizeof(TfLiteIfParams)));
205 params->then_subgraph_index = 1;
206 params->else_subgraph_index = 2;
207 auto* if_reg = ops::builtin::Register_IF();
208 if_reg->builtin_code = kTfLiteBuiltinIf;
209
210 int node_index;
211 subgraph->AddNodeWithParameters({kCondInput, kInput1, kInput2}, {kOutput}, {},
212 nullptr, 0, params, if_reg, &node_index);
213 }
214
BuildLessEqualCondSubgraph(Subgraph * subgraph,int rhs)215 void SubgraphBuilder::BuildLessEqualCondSubgraph(Subgraph* subgraph, int rhs) {
216 const int kInput1 = 0;
217 const int kInput2 = 1;
218 const int kOutput = 2;
219 const int kConstRhs = 3;
220 const int kTensorCount = 4;
221
222 // kInput1(0) ----> +------------+
223 // | LESS_EQUAL | --> kOutput(2)
224 // kConstRhs(3) --> +------------+
225 //
226 // kInput2(1) --> (unused)
227
228 int first_new_tensor_index;
229 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
230 kTfLiteOk);
231 ASSERT_EQ(first_new_tensor_index, 0);
232 ASSERT_EQ(subgraph->SetInputs({kInput1, kInput2}), kTfLiteOk);
233 ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk);
234
235 SetupTensor(subgraph, kInput1, kTfLiteInt32);
236 SetupTensor(subgraph, kInput2, kTfLiteInt32);
237 SetupTensor(subgraph, kOutput, kTfLiteBool);
238
239 auto* le_reg = ops::builtin::Register_LESS_EQUAL();
240 le_reg->builtin_code = kTfLiteBuiltinLessEqual;
241
242 CreateConstantInt32Tensor(subgraph, kConstRhs, {1}, {rhs});
243 int node_index;
244 subgraph->AddNodeWithParameters({kInput1, kConstRhs}, {kOutput}, {}, nullptr,
245 0, nullptr, le_reg, &node_index);
246 }
247
BuildAccumulateLoopBodySubgraph(Subgraph * subgraph)248 void SubgraphBuilder::BuildAccumulateLoopBodySubgraph(Subgraph* subgraph) {
249 const int kInputCounter = 0;
250 const int kInputValue = 1;
251 const int kOutputCounter = 2;
252 const int kOutputValue = 3;
253 const int kConstStep = 4;
254 const int kTensorCount = 5;
255
256 // kInputCounter(0) --> +-----+
257 // | ADD | --> kOutputCounter(2)
258 // kConstStep(4) -----> +-----+ |
259 // |
260 // v
261 // +-----+
262 // | ADD | --> kOutputValue(3)
263 // kInputValue(1) ----------------------+-----+
264
265 int first_new_tensor_index;
266 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
267 kTfLiteOk);
268 ASSERT_EQ(first_new_tensor_index, 0);
269 ASSERT_EQ(subgraph->SetInputs({kInputCounter, kInputValue}), kTfLiteOk);
270 ASSERT_EQ(subgraph->SetOutputs({kOutputCounter, kOutputValue}), kTfLiteOk);
271
272 SetupTensor(subgraph, kInputCounter, kTfLiteInt32);
273 SetupTensor(subgraph, kInputValue, kTfLiteInt32);
274 SetupTensor(subgraph, kOutputCounter, kTfLiteInt32);
275 SetupTensor(subgraph, kOutputValue, kTfLiteInt32);
276 CreateConstantInt32Tensor(subgraph, kConstStep, {1}, {1});
277
278 int node_index;
279 TfLiteAddParams* params =
280 reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
281 params->activation = kTfLiteActNone;
282 params->pot_scale_int16 = false;
283 auto* add_reg = ops::builtin::Register_ADD();
284 add_reg->builtin_code = kTfLiteBuiltinAdd;
285 subgraph->AddNodeWithParameters({0, 4}, {2}, {}, nullptr, 0, params, add_reg,
286 &node_index);
287 params = reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
288 params->activation = kTfLiteActNone;
289 params->pot_scale_int16 = false;
290 subgraph->AddNodeWithParameters({2, 1}, {3}, {}, nullptr, 0, params, add_reg,
291 &node_index);
292 }
293
BuildPadLoopBodySubgraph(Subgraph * subgraph,const std::vector<int> padding)294 void SubgraphBuilder::BuildPadLoopBodySubgraph(Subgraph* subgraph,
295 const std::vector<int> padding) {
296 const int kInputCounter = 0;
297 const int kInputValue = 1;
298 const int kOutputCounter = 2;
299 const int kOutputValue = 3;
300 const int kConstStep = 4;
301 const int kConstPadding = 5;
302 const int kTensorCount = 6;
303
304 // kInputCounter(0) --> +-----+
305 // | ADD | --> kOutputCounter(2)
306 // kConstStep(4) -----> +-----+
307 //
308 // kInputValue(1) ----> +-----+
309 // | PAD | --> kOutputValue(3)
310 // kConstPadding(5) --> +-----+
311
312 int first_new_tensor_index;
313 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
314 kTfLiteOk);
315 ASSERT_EQ(first_new_tensor_index, 0);
316 ASSERT_EQ(subgraph->SetInputs({kInputCounter, kInputValue}), kTfLiteOk);
317 ASSERT_EQ(subgraph->SetOutputs({kOutputCounter, kOutputValue}), kTfLiteOk);
318
319 SetupTensor(subgraph, kInputCounter, kTfLiteInt32);
320 SetupTensor(subgraph, kInputValue, kTfLiteInt32);
321 SetupTensor(subgraph, kOutputCounter, kTfLiteInt32);
322 SetupTensor(subgraph, kOutputValue, kTfLiteInt32);
323
324 CreateConstantInt32Tensor(subgraph, kConstStep, {1}, {1});
325 ASSERT_EQ(padding.size() % 2, 0);
326 int padding_dims = padding.size();
327 CreateConstantInt32Tensor(subgraph, kConstPadding, {1, padding_dims},
328 padding);
329
330 int node_index;
331 TfLiteAddParams* add_params =
332 reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
333 add_params->activation = kTfLiteActNone;
334 auto* add_reg = ops::builtin::Register_ADD();
335 add_reg->builtin_code = kTfLiteBuiltinAdd;
336 subgraph->AddNodeWithParameters({kInputCounter, kConstStep}, {kOutputCounter},
337 {}, nullptr, 0, add_params, add_reg,
338 &node_index);
339 TfLitePadParams* pad_params =
340 reinterpret_cast<TfLitePadParams*>(malloc(sizeof(TfLiteAddParams)));
341 auto* pad_reg = ops::builtin::Register_PAD();
342 pad_reg->builtin_code = kTfLiteBuiltinPad;
343 subgraph->AddNodeWithParameters({kInputValue, kConstPadding}, {kOutputValue},
344 {}, nullptr, 0, pad_params, pad_reg,
345 &node_index);
346 }
347
BuildWhileSubgraph(Subgraph * subgraph)348 void SubgraphBuilder::BuildWhileSubgraph(Subgraph* subgraph) {
349 const int kInput1 = 0;
350 const int kInput2 = 1;
351 const int kOutput1 = 2;
352 const int kOutput2 = 3;
353 const int kTensorCount = 4;
354
355 // kInput1(0) --> +-------+ --> kOutput1(2)
356 // | WHILE |
357 // kInput2(1) --> +-------+ --> kOutput2(3)
358
359 int first_new_tensor_index;
360 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
361 kTfLiteOk);
362 ASSERT_EQ(first_new_tensor_index, 0);
363 ASSERT_EQ(subgraph->SetInputs({kInput1, kInput2}), kTfLiteOk);
364 ASSERT_EQ(subgraph->SetOutputs({kOutput1, kOutput2}), kTfLiteOk);
365
366 SetupTensor(subgraph, kInput1, kTfLiteInt32);
367 SetupTensor(subgraph, kInput2, kTfLiteInt32);
368 SetupTensor(subgraph, kOutput1, kTfLiteInt32);
369 SetupTensor(subgraph, kOutput2, kTfLiteInt32);
370
371 TfLiteWhileParams* params =
372 reinterpret_cast<TfLiteWhileParams*>(malloc(sizeof(TfLiteWhileParams)));
373 params->cond_subgraph_index = 1;
374 params->body_subgraph_index = 2;
375 auto* while_reg = ops::builtin::Register_WHILE();
376 while_reg->builtin_code = kTfLiteBuiltinWhile;
377
378 int node_index;
379 subgraph->AddNodeWithParameters({0, 1}, {2, 3}, {}, nullptr, 0, params,
380 while_reg, &node_index);
381 }
382
BuildAssignRandomValueToVariableSubgraph(Subgraph * subgraph)383 void SubgraphBuilder::BuildAssignRandomValueToVariableSubgraph(
384 Subgraph* subgraph) {
385 const int kConstResourceId = 0;
386 const int kRandomValue = 1;
387 const int kTensorCount = 3;
388
389 // Construct a graph like ths:
390 // %1 = random_int()
391 // variable_assign(%0, %1)
392
393 int first_new_tensor_index;
394 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
395 kTfLiteOk);
396 ASSERT_EQ(subgraph->SetInputs({}), kTfLiteOk);
397 ASSERT_EQ(subgraph->SetOutputs({}), kTfLiteOk);
398
399 SetupTensor(subgraph, kRandomValue, kTfLiteInt32);
400 CreateConstantInt32Tensor(subgraph, kConstResourceId, {1}, {1024});
401
402 int node_index;
403 subgraph->AddNodeWithParameters({}, {kRandomValue}, {}, nullptr, 0, nullptr,
404 ::tflite::ops::custom::Register_RANDOM_INT(),
405 &node_index);
406 subgraph->AddNodeWithParameters(
407 {kConstResourceId, kRandomValue}, {}, {}, nullptr, 0, nullptr,
408 ::tflite::ops::builtin::Register_ASSIGN_VARIABLE(), &node_index);
409 }
410
BuildCallOnceAndReadVariableSubgraph(Subgraph * subgraph)411 void SubgraphBuilder::BuildCallOnceAndReadVariableSubgraph(Subgraph* subgraph) {
412 const int kConstResourceId = 0;
413 const int kOutput = 1;
414 const int kTensorCount = 2;
415
416 // Construct a graph like ths:
417 // Output: %1
418 // %1 = read_variable(%0)
419
420 int first_new_tensor_index;
421 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
422 kTfLiteOk);
423 ASSERT_EQ(subgraph->SetInputs({}), kTfLiteOk);
424 ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk);
425
426 SetupTensor(subgraph, kOutput, kTfLiteInt32);
427 CreateConstantInt32Tensor(subgraph, kConstResourceId, {1}, {1024});
428
429 TfLiteCallOnceParams* params = reinterpret_cast<TfLiteCallOnceParams*>(
430 malloc(sizeof(TfLiteCallOnceParams)));
431 params->init_subgraph_index = 1;
432
433 int node_index;
434 subgraph->AddNodeWithParameters({}, {}, {}, nullptr, 0, params,
435 ::tflite::ops::builtin::Register_CALL_ONCE(),
436 &node_index);
437 subgraph->AddNodeWithParameters(
438 {kConstResourceId}, {kOutput}, {}, nullptr, 0, nullptr,
439 ::tflite::ops::builtin::Register_READ_VARIABLE(), &node_index);
440 }
441
BuildCallOnceAndReadVariablePlusOneSubgraph(Subgraph * subgraph)442 void SubgraphBuilder::BuildCallOnceAndReadVariablePlusOneSubgraph(
443 Subgraph* subgraph) {
444 const int kConstResourceId = 0;
445 const int kConstOne = 1;
446 const int kReadVariableResult = 2;
447 const int kOutput = 3;
448 const int kTensorCount = 4;
449
450 // Construct a graph like ths:
451 // Output: %3
452 // %2 = read_variable(%0)
453 // %3 = add(%2, %1)
454
455 int first_new_tensor_index;
456 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
457 kTfLiteOk);
458 ASSERT_EQ(subgraph->SetInputs({}), kTfLiteOk);
459 ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk);
460
461 SetupTensor(subgraph, kReadVariableResult, kTfLiteInt32);
462 SetupTensor(subgraph, kOutput, kTfLiteInt32);
463 CreateConstantInt32Tensor(subgraph, kConstResourceId, {1}, {1024});
464 CreateConstantInt32Tensor(subgraph, kConstOne, {1}, {1});
465
466 TfLiteCallOnceParams* params = reinterpret_cast<TfLiteCallOnceParams*>(
467 malloc(sizeof(TfLiteCallOnceParams)));
468 params->init_subgraph_index = 1;
469
470 int node_index;
471 subgraph->AddNodeWithParameters({}, {}, {}, nullptr, 0, params,
472 ::tflite::ops::builtin::Register_CALL_ONCE(),
473 &node_index);
474 subgraph->AddNodeWithParameters(
475 {kConstResourceId}, {kReadVariableResult}, {}, nullptr, 0, nullptr,
476 ::tflite::ops::builtin::Register_READ_VARIABLE(), &node_index);
477
478 TfLiteAddParams* add_params =
479 reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
480 add_params->activation = kTfLiteActNone;
481 subgraph->AddNodeWithParameters(
482 {kReadVariableResult, kConstOne}, {kOutput}, {}, nullptr, 0, add_params,
483 ::tflite::ops::builtin::Register_ADD(), &node_index);
484 }
485
BuildLessEqualCondSubgraphWithDynamicTensor(Subgraph * subgraph,int rhs)486 void SubgraphBuilder::BuildLessEqualCondSubgraphWithDynamicTensor(
487 Subgraph* subgraph, int rhs) {
488 const int kStringInput1 = 0;
489 const int kStringInput2 = 1;
490 const int kIntegerInput = 2;
491 const int kOutput = 3;
492 const int kConstRhs = 4;
493 const int kTensorCount = 5;
494
495 // kIntegerInput(2) --> +------------+
496 // | LESS_EQUAL | --> kOutput(3)
497 // kConstRhs(4) --> +------------+
498 //
499 // kStringInput1(0) --> (unused)
500 // kStringInput2(1) --> (unused)
501
502 int first_new_tensor_index;
503 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
504 kTfLiteOk);
505 ASSERT_EQ(first_new_tensor_index, 0);
506 ASSERT_EQ(subgraph->SetInputs({kStringInput1, kStringInput2, kIntegerInput}),
507 kTfLiteOk);
508 ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk);
509
510 SetupTensor(subgraph, kStringInput1, kTfLiteString);
511 SetupTensor(subgraph, kStringInput2, kTfLiteString);
512 SetupTensor(subgraph, kIntegerInput, kTfLiteInt32);
513 SetupTensor(subgraph, kOutput, kTfLiteBool);
514
515 auto* le_reg = ops::builtin::Register_LESS_EQUAL();
516 le_reg->builtin_code = kTfLiteBuiltinLessEqual;
517
518 CreateConstantInt32Tensor(subgraph, kConstRhs, {1}, {rhs});
519 int node_index;
520 subgraph->AddNodeWithParameters({kIntegerInput, kConstRhs}, {kOutput}, {},
521 nullptr, 0, nullptr, le_reg, &node_index);
522 }
523
BuildBodySubgraphWithDynamicTensor(Subgraph * subgraph)524 void SubgraphBuilder::BuildBodySubgraphWithDynamicTensor(Subgraph* subgraph) {
525 const int kStringInput1 = 0;
526 const int kStringInput2 = 1;
527 const int kIntegerInput = 2;
528 const int kStringOutput1 = 0; // Forwarded of the `kStringInput1` tensor.
529 const int kStringOutput2 = 4;
530 const int kIntegerOutput = 5;
531 const int kConst = 6;
532 const int kTensorCount = 7;
533
534 // Construct a graph like this:
535 // %5 = tf.Add(%2, 1)
536 // %4 = tf.Fill(%0, %5)
537 // yield(%0, %4, %5)
538
539 int first_new_tensor_index;
540 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
541 kTfLiteOk);
542 ASSERT_EQ(first_new_tensor_index, 0);
543 ASSERT_EQ(subgraph->SetInputs({kStringInput1, kStringInput2, kIntegerInput}),
544 kTfLiteOk);
545 ASSERT_EQ(
546 subgraph->SetOutputs({kStringOutput1, kStringOutput2, kIntegerOutput}),
547 kTfLiteOk);
548
549 SetupTensor(subgraph, kStringInput1, kTfLiteString);
550 SetupTensor(subgraph, kStringInput2, kTfLiteString);
551 SetupTensor(subgraph, kIntegerInput, kTfLiteInt32);
552 SetupTensor(subgraph, kStringOutput1, kTfLiteString);
553 SetupTensor(subgraph, kStringOutput2, kTfLiteString);
554 SetupTensor(subgraph, kIntegerOutput, kTfLiteInt32);
555 SetupTensor(subgraph, kConst, kTfLiteInt32);
556
557 TfLiteAddParams* add_params =
558 reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
559 add_params->activation = kTfLiteActNone;
560
561 auto* add_reg = ops::builtin::Register_ADD();
562 add_reg->builtin_code = kTfLiteBuiltinAdd;
563
564 CreateConstantInt32Tensor(subgraph, kConst, {1}, {1});
565 int node_index;
566 subgraph->AddNodeWithParameters({kIntegerInput, kConst}, {kIntegerOutput}, {},
567 nullptr, 0, add_params, add_reg, &node_index);
568
569 auto* fill_reg = ops::builtin::Register_FILL();
570 fill_reg->builtin_code = kTfLiteBuiltinFill;
571 subgraph->AddNodeWithParameters({kIntegerOutput, kStringInput1},
572 {kStringOutput2}, {}, nullptr, 0, nullptr,
573 fill_reg, &node_index);
574 }
575
BuildWhileSubgraphWithDynamicTensor(Subgraph * subgraph)576 void SubgraphBuilder::BuildWhileSubgraphWithDynamicTensor(Subgraph* subgraph) {
577 const int kStringInput1 = 0;
578 const int kStringInput2 = 1;
579 const int kIntegerInput = 2;
580 const int kStringOutput1 = 3;
581 const int kStringOutput2 = 4;
582 const int kIntegerOutput = 5;
583 const int kTensorCount = 6;
584
585 // Create a while op with 2 string tensor and 1 integer tensor.
586 int first_new_tensor_index;
587 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
588 kTfLiteOk);
589 ASSERT_EQ(first_new_tensor_index, 0);
590 ASSERT_EQ(subgraph->SetInputs({kStringInput1, kStringInput2, kIntegerInput}),
591 kTfLiteOk);
592 ASSERT_EQ(
593 subgraph->SetOutputs({kStringOutput1, kStringOutput2, kIntegerOutput}),
594 kTfLiteOk);
595
596 SetupTensor(subgraph, kStringInput1, kTfLiteString);
597 SetupTensor(subgraph, kStringInput2, kTfLiteString);
598 SetupTensor(subgraph, kIntegerInput, kTfLiteInt32);
599 SetupTensor(subgraph, kStringOutput1, kTfLiteString);
600 SetupTensor(subgraph, kStringOutput2, kTfLiteString);
601 SetupTensor(subgraph, kIntegerOutput, kTfLiteInt32);
602
603 TfLiteWhileParams* params =
604 reinterpret_cast<TfLiteWhileParams*>(malloc(sizeof(TfLiteWhileParams)));
605 params->cond_subgraph_index = 1;
606 params->body_subgraph_index = 2;
607 auto* while_reg = ops::builtin::Register_WHILE();
608 while_reg->builtin_code = kTfLiteBuiltinWhile;
609
610 int node_index;
611 subgraph->AddNodeWithParameters(
612 {kStringInput1, kStringInput2, kIntegerInput},
613 {kStringOutput1, kStringOutput2, kIntegerOutput}, {}, nullptr, 0, params,
614 while_reg, &node_index);
615 }
616
CreateConstantInt32Tensor(Subgraph * subgraph,int tensor_index,const std::vector<int> & shape,const std::vector<int> & data)617 void SubgraphBuilder::CreateConstantInt32Tensor(Subgraph* subgraph,
618 int tensor_index,
619 const std::vector<int>& shape,
620 const std::vector<int>& data) {
621 ASSERT_GT(shape.size(), 0);
622 int num_elements = 1;
623 for (int dim : shape) {
624 num_elements *= dim;
625 }
626 ASSERT_EQ(data.size(), num_elements);
627 size_t size_in_bytes = sizeof(int32_t) * num_elements;
628 // Maybe aligned.
629 int32_t* buffer = reinterpret_cast<int32_t*>(malloc(size_in_bytes));
630 for (int i = 0; i < num_elements; ++i) {
631 buffer[i] = data[i];
632 }
633 buffers_.push_back(buffer);
634 ASSERT_EQ(subgraph->SetTensorParametersReadOnly(
635 tensor_index, kTfLiteInt32, "", shape, {},
636 reinterpret_cast<const char*>(buffer), size_in_bytes),
637 kTfLiteOk);
638 }
639
FillIntTensor(TfLiteTensor * tensor,const std::vector<int32_t> & data)640 void FillIntTensor(TfLiteTensor* tensor, const std::vector<int32_t>& data) {
641 int count = NumElements(tensor);
642 ASSERT_EQ(count, data.size());
643 for (int i = 0; i < count; ++i) {
644 tensor->data.i32[i] = data[i];
645 }
646 }
647
FillScalarStringTensor(TfLiteTensor * tensor,const std::string & data)648 void FillScalarStringTensor(TfLiteTensor* tensor, const std::string& data) {
649 StringRef str_ref;
650 str_ref.str = data.c_str();
651 str_ref.len = data.size();
652 DynamicBuffer buf;
653 buf.AddString(str_ref);
654 buf.WriteToTensor(tensor, /*new_shape=*/TfLiteIntArrayCreate(0));
655 }
656
CheckScalarStringTensor(const TfLiteTensor * tensor,const std::string & data)657 void CheckScalarStringTensor(const TfLiteTensor* tensor,
658 const std::string& data) {
659 ASSERT_EQ(tensor->dims->size, 0);
660 ASSERT_EQ(tensor->type, kTfLiteString);
661 StringRef str_ref = GetString(tensor, 0);
662 EXPECT_EQ(std::string(str_ref.str, str_ref.len), data);
663 }
664
CheckStringTensor(const TfLiteTensor * tensor,const std::vector<int> & shape,const std::vector<std::string> & data)665 void CheckStringTensor(const TfLiteTensor* tensor,
666 const std::vector<int>& shape,
667 const std::vector<std::string>& data) {
668 ASSERT_EQ(tensor->dims->size, shape.size());
669 for (int i = 0; i < tensor->dims->size; ++i) {
670 ASSERT_EQ(tensor->dims->data[i], shape[i]);
671 }
672 ASSERT_EQ(tensor->type, kTfLiteString);
673 int count = GetStringCount(tensor);
674 ASSERT_EQ(count, data.size());
675 for (int i = 0; i < count; ++i) {
676 StringRef str_ref = GetString(tensor, i);
677 EXPECT_EQ(std::string(str_ref.str, str_ref.len), data[i]);
678 }
679 }
CheckIntTensor(const TfLiteTensor * tensor,const std::vector<int> & shape,const std::vector<int32_t> & data)680 void CheckIntTensor(const TfLiteTensor* tensor, const std::vector<int>& shape,
681 const std::vector<int32_t>& data) {
682 ASSERT_EQ(tensor->dims->size, shape.size());
683 for (int i = 0; i < tensor->dims->size; ++i) {
684 ASSERT_EQ(tensor->dims->data[i], shape[i]);
685 }
686 ASSERT_EQ(tensor->type, kTfLiteInt32);
687 int count = NumElements(tensor);
688 ASSERT_EQ(count, data.size());
689 for (int i = 0; i < count; ++i) {
690 EXPECT_EQ(tensor->data.i32[i], data[i]);
691 }
692 }
693
CheckBoolTensor(const TfLiteTensor * tensor,const std::vector<int> & shape,const std::vector<bool> & data)694 void CheckBoolTensor(const TfLiteTensor* tensor, const std::vector<int>& shape,
695 const std::vector<bool>& data) {
696 ASSERT_EQ(tensor->dims->size, shape.size());
697 for (int i = 0; i < tensor->dims->size; ++i) {
698 ASSERT_EQ(tensor->dims->data[i], shape[i]);
699 }
700 ASSERT_EQ(tensor->type, kTfLiteBool);
701 int count = NumElements(tensor);
702 ASSERT_EQ(count, data.size());
703 for (int i = 0; i < count; ++i) {
704 EXPECT_EQ(tensor->data.b[i], data[i]);
705 }
706 }
707
708 } // namespace subgraph_test_util
709 } // namespace tflite
710