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 "flatbuffers/flexbuffers.h" // TF:flatbuffers
19 #include "tensorflow/lite/core/subgraph.h"
20 #include "tensorflow/lite/kernels/kernel_util.h"
21 #include "tensorflow/lite/kernels/register.h"
22 #include "tensorflow/lite/kernels/test_util.h"
23 #include "tensorflow/lite/model.h"
24
25 namespace tflite {
26
27 namespace ops {
28 namespace builtin {
29 // ADD and MUL are used to test simple branch.
30 TfLiteRegistration* Register_ADD();
31 TfLiteRegistration* Register_MUL();
32 // ADD and MUL are used to test dynamic sized subgraphs.
33 TfLiteRegistration* Register_PAD();
34 TfLiteRegistration* Register_LESS_EQUAL();
35 } // namespace builtin
36 namespace custom {
37 TfLiteRegistration* Register_IF();
38 TfLiteRegistration* Register_WHILE();
39 } // namespace custom
40 } // namespace ops
41
42 namespace subgraph_test_util {
43
44 namespace {
45
SetupTensor(Subgraph * subgraph,int tensor_index,TfLiteType type)46 void SetupTensor(Subgraph* subgraph, int tensor_index, TfLiteType type) {
47 ASSERT_EQ(subgraph->SetTensorParametersReadWrite(tensor_index, type, "", 0,
48 nullptr, {}, false),
49 kTfLiteOk);
50 }
51
52 } // namespace
53
~SubgraphBuilder()54 SubgraphBuilder::~SubgraphBuilder() {
55 for (auto buffer : buffers_) {
56 free(buffer);
57 }
58 }
59
BuildAddSubgraph(Subgraph * subgraph)60 void SubgraphBuilder::BuildAddSubgraph(Subgraph* subgraph) {
61 const int kInput1 = 0;
62 const int kInput2 = 1;
63 const int kOutput = 2;
64 const int kTensorCount = 3;
65 // kInput1(0) --> +---+
66 // |ADD| --> kOutput(2)
67 // kInput2(1) --> +---+
68
69 int first_new_tensor_index;
70 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
71 kTfLiteOk);
72 ASSERT_EQ(first_new_tensor_index, 0);
73 ASSERT_EQ(subgraph->SetInputs({kInput1, kInput2}), kTfLiteOk);
74 ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk);
75
76 SetupTensor(subgraph, kInput1, kTfLiteInt32);
77 SetupTensor(subgraph, kInput2, kTfLiteInt32);
78 SetupTensor(subgraph, kOutput, kTfLiteInt32);
79
80 TfLiteAddParams* params =
81 reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
82 params->activation = kTfLiteActNone;
83 int node_index;
84 subgraph->AddNodeWithParameters(
85 {kInput1, kInput2}, {kOutput}, nullptr, 0, params,
86 ::tflite::ops::builtin::Register_ADD(), &node_index);
87 }
88
89 // Build a subgraph with an mul op. Helper function for testing.
BuildMulSubgraph(Subgraph * subgraph)90 void SubgraphBuilder::BuildMulSubgraph(Subgraph* subgraph) {
91 const int kInput1 = 0;
92 const int kInput2 = 1;
93 const int kOutput = 2;
94 const int kTensorCount = 3;
95 // kInput1(0) --> +---+
96 // |MUL| --> kOutput(2)
97 // kInput2(1) --> +---+
98
99 int first_new_tensor_index;
100 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
101 kTfLiteOk);
102 ASSERT_EQ(first_new_tensor_index, 0);
103 ASSERT_EQ(subgraph->SetInputs({kInput1, kInput2}), kTfLiteOk);
104 ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk);
105
106 SetupTensor(subgraph, kInput1, kTfLiteInt32);
107 SetupTensor(subgraph, kInput2, kTfLiteInt32);
108 SetupTensor(subgraph, kOutput, kTfLiteInt32);
109
110 TfLiteMulParams* params =
111 reinterpret_cast<TfLiteMulParams*>(malloc(sizeof(TfLiteMulParams)));
112 params->activation = kTfLiteActNone;
113 int node_index;
114 subgraph->AddNodeWithParameters(
115 {kInput1, kInput2}, {kOutput}, nullptr, 0, params,
116 ::tflite::ops::builtin::Register_MUL(), &node_index);
117 }
118
119 // Build a subgraph with a pad op. Helper function for testing.
BuildPadSubgraph(Subgraph * subgraph)120 void SubgraphBuilder::BuildPadSubgraph(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 // |PAD| --> 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 TfLitePadParams* params =
141 reinterpret_cast<TfLitePadParams*>(malloc(sizeof(TfLitePadParams)));
142 int node_index;
143 subgraph->AddNodeWithParameters(
144 {kInput1, kInput2}, {kOutput}, nullptr, 0, params,
145 ::tflite::ops::builtin::Register_PAD(), &node_index);
146 }
147
BuildIfSubgraph(Subgraph * subgraph)148 void SubgraphBuilder::BuildIfSubgraph(Subgraph* subgraph) {
149 const int kCondInput = 0;
150 const int kInput1 = 1;
151 const int kInput2 = 2;
152 const int kOutput = 3;
153 const int kTensorCount = 4;
154
155 // kCondInput(0) --> +----+
156 // kInput1(1) ----> | IF | --> kOutput(3)
157 // kInput2(2) ----> +----+
158
159 int first_new_tensor_index;
160 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
161 kTfLiteOk);
162 ASSERT_EQ(first_new_tensor_index, 0);
163 ASSERT_EQ(subgraph->SetInputs({kCondInput, kInput1, kInput2}), kTfLiteOk);
164 ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk);
165
166 SetupTensor(subgraph, kCondInput, kTfLiteBool);
167 SetupTensor(subgraph, kInput1, kTfLiteInt32);
168 SetupTensor(subgraph, kInput2, kTfLiteInt32);
169 SetupTensor(subgraph, kOutput, kTfLiteInt32);
170
171 flexbuffers::Builder fbb;
172 fbb.Map([&]() {
173 fbb.Int("then_subgraph_index", 1);
174 fbb.Int("else_subgraph_index", 2);
175 });
176 fbb.Finish();
177 const auto& buffer = fbb.GetBuffer();
178
179 int node_index;
180 subgraph->AddNodeWithParameters(
181 {kCondInput, kInput1, kInput2}, {kOutput},
182 reinterpret_cast<const char*>(buffer.data()), buffer.size(), nullptr,
183 ::tflite::ops::custom::Register_IF(), &node_index);
184 }
185
BuildLessEqualCondSubgraph(Subgraph * subgraph,int rhs)186 void SubgraphBuilder::BuildLessEqualCondSubgraph(Subgraph* subgraph, int rhs) {
187 const int kInput1 = 0;
188 const int kInput2 = 1;
189 const int kOutput = 2;
190 const int kConstRhs = 3;
191 const int kTensorCount = 4;
192
193 // kInput1(0) ----> +------------+
194 // | LESS_EQUAL | --> kOutput(2)
195 // kConstRhs(3) --> +------------+
196 //
197 // kInput2(1) --> (unused)
198
199 int first_new_tensor_index;
200 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
201 kTfLiteOk);
202 ASSERT_EQ(first_new_tensor_index, 0);
203 ASSERT_EQ(subgraph->SetInputs({kInput1, kInput2}), kTfLiteOk);
204 ASSERT_EQ(subgraph->SetOutputs({kOutput}), kTfLiteOk);
205
206 SetupTensor(subgraph, kInput1, kTfLiteInt32);
207 SetupTensor(subgraph, kInput2, kTfLiteInt32);
208 SetupTensor(subgraph, kOutput, kTfLiteBool);
209
210 CreateConstantInt32Tensor(subgraph, kConstRhs, {1}, {rhs});
211 int node_index;
212 subgraph->AddNodeWithParameters(
213 {kInput1, kConstRhs}, {kOutput}, nullptr, 0, nullptr,
214 ::tflite::ops::builtin::Register_LESS_EQUAL(), &node_index);
215 }
216
BuildAccumulateLoopBodySubgraph(Subgraph * subgraph)217 void SubgraphBuilder::BuildAccumulateLoopBodySubgraph(Subgraph* subgraph) {
218 const int kInputCounter = 0;
219 const int kInputValue = 1;
220 const int kOutputCounter = 2;
221 const int kOutputValue = 3;
222 const int kConstStep = 4;
223 const int kTensorCount = 5;
224
225 // kInputCounter(0) --> +-----+
226 // | ADD | --> kOutputCounter(2)
227 // kConstStep(4) -----> +-----+ |
228 // |
229 // v
230 // +-----+
231 // | ADD | --> kOutputValue(3)
232 // kInputValue(1) ----------------------+-----+
233
234 int first_new_tensor_index;
235 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
236 kTfLiteOk);
237 ASSERT_EQ(first_new_tensor_index, 0);
238 ASSERT_EQ(subgraph->SetInputs({kInputCounter, kInputValue}), kTfLiteOk);
239 ASSERT_EQ(subgraph->SetOutputs({kOutputCounter, kOutputValue}), kTfLiteOk);
240
241 SetupTensor(subgraph, kInputCounter, kTfLiteInt32);
242 SetupTensor(subgraph, kInputValue, kTfLiteInt32);
243 SetupTensor(subgraph, kOutputCounter, kTfLiteInt32);
244 SetupTensor(subgraph, kOutputValue, kTfLiteInt32);
245 CreateConstantInt32Tensor(subgraph, kConstStep, {1}, {1});
246
247 int node_index;
248 TfLiteAddParams* params =
249 reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
250 params->activation = kTfLiteActNone;
251 subgraph->AddNodeWithParameters({0, 4}, {2}, nullptr, 0, params,
252 ::tflite::ops::builtin::Register_ADD(),
253 &node_index);
254 params = reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
255 params->activation = kTfLiteActNone;
256 subgraph->AddNodeWithParameters({2, 1}, {3}, nullptr, 0, params,
257 ::tflite::ops::builtin::Register_ADD(),
258 &node_index);
259 }
260
BuildPadLoopBodySubgraph(Subgraph * subgraph,const std::vector<int> padding)261 void SubgraphBuilder::BuildPadLoopBodySubgraph(Subgraph* subgraph,
262 const std::vector<int> padding) {
263 const int kInputCounter = 0;
264 const int kInputValue = 1;
265 const int kOutputCounter = 2;
266 const int kOutputValue = 3;
267 const int kConstStep = 4;
268 const int kConstPadding = 5;
269 const int kTensorCount = 6;
270
271 // kInputCounter(0) --> +-----+
272 // | ADD | --> kOutputCounter(2)
273 // kConstStep(4) -----> +-----+
274 //
275 // kInputValue(1) ----> +-----+
276 // | PAD | --> kOutputValue(3)
277 // kConstPadding(5) --> +-----+
278
279 int first_new_tensor_index;
280 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
281 kTfLiteOk);
282 ASSERT_EQ(first_new_tensor_index, 0);
283 ASSERT_EQ(subgraph->SetInputs({kInputCounter, kInputValue}), kTfLiteOk);
284 ASSERT_EQ(subgraph->SetOutputs({kOutputCounter, kOutputValue}), kTfLiteOk);
285
286 SetupTensor(subgraph, kInputCounter, kTfLiteInt32);
287 SetupTensor(subgraph, kInputValue, kTfLiteInt32);
288 SetupTensor(subgraph, kOutputCounter, kTfLiteInt32);
289 SetupTensor(subgraph, kOutputValue, kTfLiteInt32);
290
291 CreateConstantInt32Tensor(subgraph, kConstStep, {1}, {1});
292 ASSERT_EQ(padding.size() % 2, 0);
293 int padding_dims = padding.size();
294 CreateConstantInt32Tensor(subgraph, kConstPadding, {1, padding_dims},
295 padding);
296
297 int node_index;
298 TfLiteAddParams* add_params =
299 reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
300 add_params->activation = kTfLiteActNone;
301 subgraph->AddNodeWithParameters(
302 {kInputCounter, kConstStep}, {kOutputCounter}, nullptr, 0, add_params,
303 ::tflite::ops::builtin::Register_ADD(), &node_index);
304 TfLitePadParams* pad_params =
305 reinterpret_cast<TfLitePadParams*>(malloc(sizeof(TfLiteAddParams)));
306 subgraph->AddNodeWithParameters(
307 {kInputValue, kConstPadding}, {kOutputValue}, nullptr, 0, pad_params,
308 ::tflite::ops::builtin::Register_PAD(), &node_index);
309 }
310
BuildWhileSubgraph(Subgraph * subgraph)311 void SubgraphBuilder::BuildWhileSubgraph(Subgraph* subgraph) {
312 const int kInput1 = 0;
313 const int kInput2 = 1;
314 const int kOutput1 = 2;
315 const int kOutput2 = 3;
316 const int kTensorCount = 4;
317
318 // kInput1(0) --> +-------+ --> kOutput1(2)
319 // | WHILE |
320 // kInput2(1) --> +-------+ --> kOutput2(3)
321
322 int first_new_tensor_index;
323 ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
324 kTfLiteOk);
325 ASSERT_EQ(first_new_tensor_index, 0);
326 ASSERT_EQ(subgraph->SetInputs({kInput1, kInput2}), kTfLiteOk);
327 ASSERT_EQ(subgraph->SetOutputs({kOutput1, kOutput2}), kTfLiteOk);
328
329 SetupTensor(subgraph, kInput1, kTfLiteInt32);
330 SetupTensor(subgraph, kInput2, kTfLiteInt32);
331 SetupTensor(subgraph, kOutput1, kTfLiteInt32);
332 SetupTensor(subgraph, kOutput2, kTfLiteInt32);
333
334 flexbuffers::Builder fbb;
335 fbb.Map([&]() {
336 fbb.Int("cond_subgraph_index", 1);
337 fbb.Int("body_subgraph_index", 2);
338 });
339 fbb.Finish();
340 const auto& buffer = fbb.GetBuffer();
341
342 int node_index;
343 subgraph->AddNodeWithParameters(
344 {0, 1}, {2, 3}, reinterpret_cast<const char*>(buffer.data()),
345 buffer.size(), nullptr, ::tflite::ops::custom::Register_WHILE(),
346 &node_index);
347 }
348
CreateConstantInt32Tensor(Subgraph * subgraph,int tensor_index,const std::vector<int> & shape,const std::vector<int> & data)349 void SubgraphBuilder::CreateConstantInt32Tensor(Subgraph* subgraph,
350 int tensor_index,
351 const std::vector<int>& shape,
352 const std::vector<int>& data) {
353 ASSERT_GT(shape.size(), 0);
354 int num_elements = 1;
355 for (int dim : shape) {
356 num_elements *= dim;
357 }
358 ASSERT_EQ(data.size(), num_elements);
359 size_t size_in_bytes = sizeof(int32_t) * num_elements;
360 // Maybe aligned.
361 int32_t* buffer = reinterpret_cast<int32_t*>(malloc(size_in_bytes));
362 for (int i = 0; i < num_elements; ++i) {
363 buffer[i] = data[i];
364 }
365 buffers_.push_back(buffer);
366 ASSERT_EQ(subgraph->SetTensorParametersReadOnly(
367 tensor_index, kTfLiteInt32, "", shape, {},
368 reinterpret_cast<const char*>(buffer), size_in_bytes),
369 kTfLiteOk);
370 }
371
FillIntTensor(TfLiteTensor * tensor,const std::vector<int32_t> & data)372 void FillIntTensor(TfLiteTensor* tensor, const std::vector<int32_t>& data) {
373 int count = NumElements(tensor);
374 ASSERT_EQ(count, data.size());
375 for (int i = 0; i < count; ++i) {
376 tensor->data.i32[i] = data[i];
377 }
378 }
379
CheckIntTensor(const TfLiteTensor * tensor,const std::vector<int> & shape,const std::vector<int32_t> & data)380 void CheckIntTensor(const TfLiteTensor* tensor, const std::vector<int>& shape,
381 const std::vector<int32_t>& data) {
382 ASSERT_EQ(tensor->dims->size, shape.size());
383 for (int i = 0; i < tensor->dims->size; ++i) {
384 ASSERT_EQ(tensor->dims->data[i], shape[i]);
385 }
386 ASSERT_EQ(tensor->type, kTfLiteInt32);
387 int count = NumElements(tensor);
388 ASSERT_EQ(count, data.size());
389 for (int i = 0; i < count; ++i) {
390 EXPECT_EQ(tensor->data.i32[i], data[i]);
391 }
392 }
393
CheckBoolTensor(const TfLiteTensor * tensor,const std::vector<int> & shape,const std::vector<bool> & data)394 void CheckBoolTensor(const TfLiteTensor* tensor, const std::vector<int>& shape,
395 const std::vector<bool>& data) {
396 ASSERT_EQ(tensor->dims->size, shape.size());
397 for (int i = 0; i < tensor->dims->size; ++i) {
398 ASSERT_EQ(tensor->dims->data[i], shape[i]);
399 }
400 ASSERT_EQ(tensor->type, kTfLiteBool);
401 int count = NumElements(tensor);
402 ASSERT_EQ(count, data.size());
403 for (int i = 0; i < count; ++i) {
404 EXPECT_EQ(tensor->data.b[i], data[i]);
405 }
406 }
407
408 } // namespace subgraph_test_util
409 } // namespace tflite
410