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 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 54 SubgraphBuilder::~SubgraphBuilder() { 55 for (auto buffer : buffers_) { 56 free(buffer); 57 } 58 } 59 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. 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. 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 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 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 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 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 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 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 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 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 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
