compiler/tflchef/core/src/ModelChef.cpp

   1 /*
   2  * Copyright (c) 2018 Samsung Electronics Co., Ltd. All Rights Reserved
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *    http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include "tflchef/ModelChef.h"
  18 #include <souschef/RangedArguments.h>
  19 #include <souschef/Registry.h>
  20
  21 #include "Convert.h"
  22
  23 #include <souschef/DataChefs.h>
  24
  25 #include "OpChef.h"
  26 #include "OpChefs.h"
  27
  28 #include <souschef/Dataset.h>
  29 #include <souschef/Dims.h>
  30
  31 #include "Log.h"
  32
  33 #include <iterator>
  34 #include <map>
  35 #include <string>
  36 #include <vector>
  37
  38 #include <cassert>
  39 #include <fstream>
  40 #include <iostream>
  41 #include <numeric>
  42 #include <sstream>
  43 #include <stdexcept>
  44
  45 using namespace souschef;
  46
  47 namespace
  48 {
  49
  50 class GeneratedModelImpl final : public tflchef::GeneratedModel::Impl
  51 {
  52 public:
  53   GeneratedModelImpl(std::unique_ptr<flatbuffers::FlatBufferBuilder> &&builder)
  54     : _builder{std::move(builder)}
  55   {
  56     // DO NOTHING
  57   }
  58
  59 public:
  60   const char *base(void) const override
  61   {
  62     // Return the base address of generated flatbuffer model
  63     return reinterpret_cast<const char *>(_builder->GetBufferPointer());
  64   }
  65
  66 public:
  67   size_t size(void) const override
  68   {
  69     // Return the size of generated flatbuffer model
  70     return _builder->GetSize();
  71   }
  72
  73 private:
  74   std::unique_ptr<flatbuffers::FlatBufferBuilder> _builder;
  75 };
  76
  77 } // namespace
  78
  79 namespace
  80 {
  81
  82 struct DataChefRegistry final : public Registry<DataChefFactory>
  83 {
  84 };
  85
  86 DataChefRegistry &data_chef_registry(const tflchef::TensorType &type)
  87 {
  88   static DataChefRegistry s32;
  89   static DataChefRegistry s64;
  90   static DataChefRegistry fp32;
  91   static DataChefRegistry u8;
  92   static DataChefRegistry string;
  93   static DataChefRegistry boolean;
  94   static DataChefRegistry s16;
  95   static DataChefRegistry fp16;
  96
  97   switch (type)
  98   {
  99     case tflchef::INT32:
 100       return s32;
 101     case tflchef::INT64:
 102       return s64;
 103     case tflchef::FLOAT32:
 104       return fp32;
 105     case tflchef::FLOAT16:
 106       return fp16;
 107     case tflchef::UINT8:
 108       return u8;
 109     case tflchef::STRING:
 110       return string;
 111     case tflchef::BOOL:
 112       return boolean;
 113     case tflchef::INT16:
 114       return s16;
 115     default:
 116       break;
 117   }
 118
 119   throw std::runtime_error{"Unknown tensor type"};
 120 }
 121
 122 struct OpChefRegistry final : public Registry<OpChefFactory>
 123 {
 124 };
 125
 126 OpChefRegistry &op_chef_registry(void)
 127 {
 128   static OpChefRegistry registry;
 129   return registry;
 130 }
 131
 132 /// @brief This will prepare a map of unique builtin codes in the model recipe
 133 std::map<tflite::BuiltinOperator, int32_t>
 134 gather_builtincode_map(const ::tflchef::ModelRecipe &model_recipe)
 135 {
 136   // Key and value of the map are BuiltinOperator and operator version
 137   std::map<tflite::BuiltinOperator, int32_t> builtin_map;
 138
 139   for (const auto &operation : model_recipe.operation())
 140   {
 141     auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
 142     if (op_chef->code() == tflite::BuiltinOperator_CUSTOM)
 143       continue;
 144
 145     // Various operation version is unified as the highest version among them
 146     if (builtin_map.find(op_chef->code()) == builtin_map.end() ||
 147         builtin_map[op_chef->code()] < operation.version())
 148       builtin_map[op_chef->code()] = operation.version();
 149   }
 150
 151   // Add ops used in Graphs(subgraphs)
 152   for (int g = 0; g < model_recipe.graph_size(); ++g)
 153   {
 154     const auto &graph = model_recipe.graph(g);
 155     for (const auto &operation : graph.operation())
 156     {
 157       auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
 158       if (op_chef->code() == tflite::BuiltinOperator_CUSTOM)
 159         continue;
 160
 161       // Various operation version is unified as the highest version among them
 162       if (builtin_map.find(op_chef->code()) == builtin_map.end() ||
 163           builtin_map[op_chef->code()] < operation.version())
 164         builtin_map[op_chef->code()] = operation.version();
 165     }
 166   }
 167
 168   return builtin_map;
 169 }
 170
 171 /// @brief This will prepare a set of unique custom codes in the mode recipe
 172 std::set<std::string> gather_customcode_set(const ::tflchef::ModelRecipe &model_recipe)
 173 {
 174   std::set<std::string> customcode_set;
 175   for (const auto &operation : model_recipe.operation())
 176   {
 177     auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
 178     if (op_chef->code() == tflite::BuiltinOperator_CUSTOM)
 179       customcode_set.insert(operation.type());
 180   }
 181
 182   // Add ops used in Graphs(subgraphs)
 183   for (int g = 0; g < model_recipe.graph_size(); ++g)
 184   {
 185     const auto &graph = model_recipe.graph(g);
 186     for (const auto &operation : graph.operation())
 187     {
 188       auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
 189       if (op_chef->code() == tflite::BuiltinOperator_CUSTOM)
 190         customcode_set.insert(operation.type());
 191     }
 192   }
 193
 194   return customcode_set;
 195 }
 196
 197 } // namespace
 198
 199 namespace
 200 {
 201
 202 struct CookParams
 203 {
 204   std::vector<flatbuffers::Offset<::tflite::Buffer>> &buffer_vec;
 205   std::vector<flatbuffers::Offset<::tflite::OperatorCode>> &code_vec;
 206   std::vector<flatbuffers::Offset<::tflite::SubGraph>> &subgraph_vec;
 207   std::unique_ptr<flatbuffers::FlatBufferBuilder> &flatbuffer_builder;
 208   std::map<tflite::BuiltinOperator, int32_t> &builtin_code_map;
 209   std::vector<std::string> &custom_code_vec;
 210   std::string noname;
 211 };
 212
 213 std::vector<flatbuffers::Offset<tflite::DimensionMetadata>>
 214 make_dim_metadata_vec(flatbuffers::FlatBufferBuilder *flatbuffer_builder, int32_t dims_count,
 215                       const std::vector<int> &traversal_order_vec,
 216                       const std::vector<sparsity::TfLiteDimensionType> &format_vec,
 217                       const std::vector<std::vector<int32_t>> &dim_metadata_src)
 218 {
 219   // Build sparsity parameter.
 220   std::vector<flatbuffers::Offset<tflite::DimensionMetadata>> dim_metadata_vec(dims_count);
 221   for (int32_t i = 0; i < dims_count; i++)
 222   {
 223     const int32_t metadata_idx = 2 * i;
 224     if (format_vec[traversal_order_vec[i]] == sparsity::kTfLiteDimSparseCSR)
 225     {
 226       auto array_segments =
 227         tflite::CreateInt32Vector(*flatbuffer_builder,
 228                                   flatbuffer_builder->CreateVector(dim_metadata_src[metadata_idx]))
 229           .Union();
 230       auto array_indices =
 231         tflite::CreateInt32Vector(
 232           *flatbuffer_builder, flatbuffer_builder->CreateVector(dim_metadata_src[metadata_idx + 1]))
 233           .Union();
 234       dim_metadata_vec[i] =
 235         tflite::CreateDimensionMetadata(*flatbuffer_builder, tflite::DimensionType_SPARSE_CSR, 0,
 236                                         tflite::SparseIndexVector_Int32Vector, array_segments,
 237                                         tflite::SparseIndexVector_Int32Vector, array_indices);
 238     }
 239     else
 240     {
 241       dim_metadata_vec[i] = tflite::CreateDimensionMetadata(
 242         *flatbuffer_builder, tflite::DimensionType_DENSE, dim_metadata_src[metadata_idx][0]);
 243     }
 244   }
 245   return dim_metadata_vec;
 246 }
 247
 248 template <typename T> std::map<std::string, int32_t> cook_graph(const T &graph, CookParams &cp)
 249 {
 250   LOGGER(l);
 251
 252   std::vector<flatbuffers::Offset<::tflite::Buffer>> &buffer_vec = cp.buffer_vec;
 253   std::vector<flatbuffers::Offset<::tflite::OperatorCode>> &code_vec = cp.code_vec;
 254   std::vector<flatbuffers::Offset<::tflite::SubGraph>> &subgraph_vec = cp.subgraph_vec;
 255   std::unique_ptr<flatbuffers::FlatBufferBuilder> &flatbuffer_builder = cp.flatbuffer_builder;
 256   std::map<tflite::BuiltinOperator, int32_t> &builtin_code_map = cp.builtin_code_map;
 257   std::vector<std::string> &custom_code_vec = cp.custom_code_vec;
 258
 259   // Operand-related
 260   std::vector<flatbuffers::Offset<::tflite::Tensor>> tensor_vec;
 261
 262   // Operation-related
 263   std::vector<flatbuffers::Offset<::tflite::Operator>> operator_vec;
 264
 265   // default name for graph
 266   std::string graph_name = cp.noname;
 267   if (graph.has_name())
 268     graph_name = graph.name();
 269
 270   // Tensor Name -> Tensor ID mapping (per Graph)
 271   std::map<std::string, int32_t> symbol_table;
 272
 273   auto lookup = [&symbol_table, &graph_name](const std::string &name) {
 274     if (symbol_table.find(name) != symbol_table.end())
 275       return symbol_table.at(name);
 276     else if (name == "")
 277       return -1; // -1 in TFLite means that optional input tensor is empty.
 278     else
 279     {
 280       std::string msg = "tflchef : input not found in " + graph_name + " graph";
 281       throw std::runtime_error(msg.c_str());
 282     }
 283   };
 284
 285   int32_t buffer_start = buffer_vec.size();
 286   int32_t buffer_index = 0;
 287
 288   // Create buffer(s) 1~n(I) for input(s)
 289   const auto size_input = graph.input_size();
 290   for (int ci = 0; ci < size_input; ++ci)
 291   {
 292     tflite::BufferBuilder buffer_builder{*flatbuffer_builder};
 293     buffer_vec.emplace_back(buffer_builder.Finish());
 294   }
 295   // Create buffer(s) n(I)+1~n(I)+n(O) for output(s)
 296   const auto size_output = graph.output_size();
 297   for (int co = 0; co < size_output; ++co)
 298   {
 299     tflite::BufferBuilder buffer_builder{*flatbuffer_builder};
 300     buffer_vec.emplace_back(buffer_builder.Finish());
 301   }
 302
 303   auto input_names = as_dataset(graph.input()).vectorize();
 304   auto output_names = as_dataset(graph.output()).vectorize();
 305
 306   for (const auto &operand : graph.operand())
 307   {
 308     assert(operand.has_name());
 309
 310     assert(operand.has_type());
 311
 312     flatbuffers::Offset<tflite::SparsityParameters> sparsity_index;
 313
 314     flatbuffers::Offset<flatbuffers::Vector<int32_t>> shape;
 315     std::vector<int32_t> dims;
 316     if (operand.has_shape())
 317     {
 318       dims = as_dims(operand.shape());
 319       shape = flatbuffer_builder->CreateVector(dims);
 320     }
 321
 322     auto name = flatbuffer_builder->CreateString(operand.name());
 323
 324     buffer_index = 0;
 325
 326     // Create Buffer if filler is specified
 327     if (operand.has_filler())
 328     {
 329       const auto &filler = operand.filler();
 330
 331       assert(filler.has_tag());
 332
 333       auto args = ranged_arguments(filler.arg().begin(), filler.arg().end());
 334       auto chef = data_chef_registry(operand.type()).lookup(filler.tag()).create(args);
 335
 336       assert(chef != nullptr);
 337
 338       // Create Data
 339       int32_t count = (element_count(dims) > 0) ? element_count(dims) : filler.arg_size();
 340       auto data_vec = chef->generate(count);
 341
 342       if (operand.has_make_sparse() && operand.make_sparse())
 343       {
 344         assert(not operand.has_sparsity());
 345         assert(operand.has_shape());
 346
 347         const int32_t dims_count = dims.size();
 348         std::vector<int> traversal_order_vec;
 349         std::vector<sparsity::TfLiteDimensionType> format_vec;
 350         for (int32_t o = 0; o < dims_count; ++o)
 351           traversal_order_vec.push_back(o);
 352         for (int32_t o = 0; o < dims_count - 1; ++o)
 353           format_vec.push_back(sparsity::kTfLiteDimDense);
 354         format_vec.push_back(sparsity::kTfLiteDimSparseCSR);
 355
 356         if (operand.type() == tflchef::FLOAT32)
 357         {
 358           ::sparsity::FormatConverter<float> converter(dims, traversal_order_vec, format_vec);
 359           converter.DenseToSparse(reinterpret_cast<const float *>(data_vec.data()));
 360           const auto &sparse_data = converter.GetData();
 361
 362           std::vector<uint8_t> sparse_uint8;
 363           for (int c = 0; c < sparse_data.size(); ++c)
 364           {
 365             const float value = sparse_data.at(c);
 366             const uint8_t *arr = reinterpret_cast<const uint8_t *>(&value);
 367             for (uint32_t b = 0; b < sizeof(float); ++b)
 368             {
 369               sparse_uint8.emplace_back(arr[b]);
 370             }
 371           }
 372           auto data = flatbuffer_builder->CreateVector(sparse_uint8);
 373
 374           // Create Buffer
 375           tflite::BufferBuilder buffer_builder{*flatbuffer_builder};
 376           buffer_builder.add_data(data);
 377           auto buffer = buffer_builder.Finish();
 378
 379           // Update Buffer Index & Vector
 380           buffer_index = buffer_vec.size();
 381           buffer_vec.emplace_back(buffer);
 382
 383           // save SparsityParameters
 384           auto traversal_order = flatbuffer_builder->CreateVector(traversal_order_vec);
 385
 386           // Create block map
 387           std::vector<int> block_map_vec{};
 388           auto block_map = flatbuffer_builder->CreateVector(block_map_vec);
 389
 390           // Create dimension metadata
 391           const auto &dim_metadata_src = converter.GetDimMetadata();
 392           auto dim_metadata_vec =
 393             make_dim_metadata_vec(flatbuffer_builder.get(), dims_count, traversal_order_vec,
 394                                   format_vec, dim_metadata_src);
 395           auto dim_metadata = flatbuffer_builder->CreateVector(dim_metadata_vec);
 396           sparsity_index = tflite::CreateSparsityParameters(*flatbuffer_builder, traversal_order,
 397                                                             block_map, dim_metadata);
 398         }
 399         else if (operand.type() == tflchef::FLOAT16)
 400         {
 401           ::sparsity::FormatConverter<uint16_t> converter(dims, traversal_order_vec, format_vec);
 402           converter.DenseToSparse(reinterpret_cast<const uint16_t *>(data_vec.data()));
 403           const auto &sparse_data = converter.GetData();
 404
 405           std::vector<uint8_t> sparse_uint8;
 406           for (int c = 0; c < sparse_data.size(); ++c)
 407           {
 408             const uint16_t value = sparse_data.at(c);
 409             const uint8_t *arr = reinterpret_cast<const uint8_t *>(&value);
 410             for (uint32_t b = 0; b < sizeof(uint16_t); ++b)
 411             {
 412               sparse_uint8.emplace_back(arr[b]);
 413             }
 414           }
 415           auto data = flatbuffer_builder->CreateVector(sparse_uint8);
 416
 417           // Create Buffer
 418           tflite::BufferBuilder buffer_builder{*flatbuffer_builder};
 419           buffer_builder.add_data(data);
 420           auto buffer = buffer_builder.Finish();
 421
 422           // Update Buffer Index & Vector
 423           buffer_index = buffer_vec.size();
 424           buffer_vec.emplace_back(buffer);
 425
 426           // save SparsityParameters
 427           auto traversal_order = flatbuffer_builder->CreateVector(traversal_order_vec);
 428
 429           // Create block map
 430           std::vector<int> block_map_vec{};
 431           auto block_map = flatbuffer_builder->CreateVector(block_map_vec);
 432
 433           // Create dimension metadata
 434           const auto &dim_metadata_src = converter.GetDimMetadata();
 435           auto dim_metadata_vec =
 436             make_dim_metadata_vec(flatbuffer_builder.get(), dims_count, traversal_order_vec,
 437                                   format_vec, dim_metadata_src);
 438           auto dim_metadata = flatbuffer_builder->CreateVector(dim_metadata_vec);
 439           sparsity_index = tflite::CreateSparsityParameters(*flatbuffer_builder, traversal_order,
 440                                                             block_map, dim_metadata);
 441         }
 442         else
 443         {
 444           throw std::runtime_error{"NYI: unsupported operand type"};
 445         }
 446       }
 447       else
 448       {
 449         auto data = flatbuffer_builder->CreateVector(data_vec);
 450
 451         // Create Buffer
 452         tflite::BufferBuilder buffer_builder{*flatbuffer_builder};
 453         buffer_builder.add_data(data);
 454         auto buffer = buffer_builder.Finish();
 455
 456         // Update Buffer Index & Vector
 457         buffer_index = buffer_vec.size();
 458         buffer_vec.emplace_back(buffer);
 459       }
 460     }
 461     else
 462     {
 463       // if this is input or output, assign to that buffer_index
 464       int idx = 0;
 465       for (auto it = input_names.begin(); it != input_names.end(); ++it, ++idx)
 466       {
 467         if (*it == operand.name())
 468         {
 469           buffer_index = buffer_start + idx;
 470           break;
 471         }
 472       }
 473       if (buffer_index == 0)
 474       {
 475         idx = 0;
 476         for (auto it = output_names.begin(); it != output_names.end(); ++it, ++idx)
 477         {
 478           if (*it == operand.name())
 479           {
 480             buffer_index = buffer_start + size_input + idx;
 481             break;
 482           }
 483         }
 484       }
 485       if (buffer_index == 0)
 486       {
 487         // we couldn't find the buffer; create an empty buffer for this tensor
 488         buffer_index = buffer_vec.size();
 489
 490         tflite::BufferBuilder buffer_builder{*flatbuffer_builder};
 491         buffer_vec.emplace_back(buffer_builder.Finish());
 492       }
 493     }
 494     assert(buffer_index != 0);
 495
 496     flatbuffers::Offset<tflite::QuantizationParameters> quant_index;
 497
 498     // Create QuantizationParameters if quant is specified
 499     if (operand.has_quant())
 500     {
 501       const auto &quant = operand.quant();
 502
 503       // Create each parameters
 504       // NOTE if some parameters are not given, those will be set to default value
 505       std::vector<float> quant_max_vec(quant.max_size());
 506       std::vector<float> quant_min_vec(quant.min_size());
 507       std::vector<float> quant_scale_vec(quant.scale_size());
 508       std::vector<int64_t> quant_zero_point_vec(quant.zero_point_size());
 509
 510       for (uint32_t i = 0; i < quant.max_size(); ++i)
 511         quant_max_vec.at(i) = quant.max(i);
 512       for (uint32_t i = 0; i < quant.min_size(); ++i)
 513         quant_min_vec.at(i) = quant.min(i);
 514       for (uint32_t i = 0; i < quant.scale_size(); ++i)
 515         quant_scale_vec.at(i) = quant.scale(i);
 516       for (uint32_t i = 0; i < quant.zero_point_size(); ++i)
 517         quant_zero_point_vec.at(i) = quant.zero_point(i);
 518
 519       auto quant_max = flatbuffer_builder->CreateVector(quant_max_vec);
 520       auto quant_min = flatbuffer_builder->CreateVector(quant_min_vec);
 521       auto quant_scale = flatbuffer_builder->CreateVector(quant_scale_vec);
 522       auto quant_zero_point = flatbuffer_builder->CreateVector(quant_zero_point_vec);
 523
 524       // Create QuantizationParameters
 525       tflite::QuantizationParametersBuilder quant_builder{*flatbuffer_builder};
 526       quant_builder.add_max(quant_max);
 527       quant_builder.add_min(quant_min);
 528       quant_builder.add_scale(quant_scale);
 529       quant_builder.add_zero_point(quant_zero_point);
 530       quant_builder.add_quantized_dimension(quant.quantized_dimension());
 531
 532       // Update QuantizationParameters Index
 533       quant_index = quant_builder.Finish();
 534     }
 535
 536     if (operand.has_sparsity())
 537     {
 538       const auto &sparsity = operand.sparsity();
 539
 540       // Create traversal order
 541       std::vector<int> traversal_order_vec{sparsity.traversal_order().dim().begin(),
 542                                            sparsity.traversal_order().dim().end()};
 543       auto traversal_order = flatbuffer_builder->CreateVector(traversal_order_vec);
 544
 545       // Create block map
 546       std::vector<int> block_map_vec{sparsity.block_map().dim().begin(),
 547                                      sparsity.block_map().dim().end()};
 548       auto block_map = flatbuffer_builder->CreateVector(block_map_vec);
 549
 550       // Create dimension metadata
 551       std::vector<flatbuffers::Offset<tflite::DimensionMetadata>> dim_metadata_vec;
 552       auto recipe_dim_metadata = sparsity.dim_metadata();
 553       for (const auto &dm : recipe_dim_metadata)
 554       {
 555         // Create array segments
 556         auto tflite_array_segments =
 557           as_tflite_sparse_index_vec(*flatbuffer_builder, dm.array_segments());
 558
 559         // Create array indices
 560         auto tflite_array_indices =
 561           as_tflite_sparse_index_vec(*flatbuffer_builder, dm.array_indices());
 562
 563         auto tflite_dim_metadata_builder = tflite::DimensionMetadataBuilder{*flatbuffer_builder};
 564         tflite_dim_metadata_builder.add_format(as_tflite_dimensiontype(dm.format()));
 565         tflite_dim_metadata_builder.add_dense_size(dm.dense_size());
 566         tflite_dim_metadata_builder.add_array_segments(tflite_array_segments);
 567         tflite_dim_metadata_builder.add_array_segments_type(
 568           as_tflite_sparse_idx_vec_type(dm.array_segments().type()));
 569         tflite_dim_metadata_builder.add_array_indices(tflite_array_indices);
 570         tflite_dim_metadata_builder.add_array_indices_type(
 571           as_tflite_sparse_idx_vec_type(dm.array_indices().type()));
 572         auto tflite_dim_metadata = tflite_dim_metadata_builder.Finish();
 573         dim_metadata_vec.emplace_back(tflite_dim_metadata);
 574       }
 575       auto dim_metadata = flatbuffer_builder->CreateVector(dim_metadata_vec);
 576
 577       sparsity_index = tflite::CreateSparsityParameters(*flatbuffer_builder, traversal_order,
 578                                                         block_map, dim_metadata);
 579     }
 580
 581     flatbuffers::Offset<flatbuffers::Vector<int32_t>> shape_signature;
 582     if (operand.has_shape_signature())
 583     {
 584       auto signature = as_dims(operand.shape_signature());
 585       shape_signature = flatbuffer_builder->CreateVector(signature);
 586     }
 587
 588     // Create Tensor
 589     tflite::TensorBuilder tensor_builder{*flatbuffer_builder};
 590
 591     tensor_builder.add_shape(shape);
 592     tensor_builder.add_type(as_tflite_tensortype(operand.type()));
 593     tensor_builder.add_buffer(buffer_index);
 594     tensor_builder.add_name(name);
 595     tensor_builder.add_is_variable(operand.is_variable());
 596     if (operand.has_quant())
 597       tensor_builder.add_quantization(quant_index);
 598     tensor_builder.add_sparsity(sparsity_index);
 599     if (operand.has_shape_signature())
 600       tensor_builder.add_shape_signature(shape_signature);
 601
 602     // Append!
 603     tensor_vec.emplace_back(tensor_builder.Finish());
 604
 605     // Update Tensor Name -> Tensor Index Map
 606     int32_t tensor_index = symbol_table.size();
 607     const auto &tensor_name = operand.name();
 608
 609     INFO(l) << "Symbol [" << tensor_name << "] = Tensor " << tensor_index << std::endl;
 610
 611     symbol_table[tensor_name] = tensor_index;
 612   }
 613
 614   // Create Operator
 615   for (const auto &operation : graph.operation())
 616   {
 617     assert(operation.has_type());
 618
 619     auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
 620
 621     // Create 'inputs'
 622     std::vector<int32_t> input_vec = as_dataset(operation.input()).map(lookup).vectorize();
 623     auto inputs = flatbuffer_builder->CreateVector(input_vec);
 624
 625     // Create 'outputs'
 626     std::vector<int32_t> output_vec = as_dataset(operation.output()).map(lookup).vectorize();
 627     auto outputs = flatbuffer_builder->CreateVector(output_vec);
 628
 629     // Create Option
 630     auto options = op_chef->value(*flatbuffer_builder);
 631
 632     // Create Custom option
 633     auto circle_custom_options = op_chef->custom_value(*flatbuffer_builder);
 634
 635     // Create Operator
 636     tflite::OperatorBuilder op_builder{*flatbuffer_builder};
 637
 638     // Note that opcode_index is an index into the operator_codes vector.
 639     // operator_codes consists of buildtin_code and custom_code, which is inserted sequentially.
 640     uint32_t opcode_index = 0;
 641     auto op_it = builtin_code_map.find(op_chef->code());
 642     // builtin operator
 643     if (op_it != builtin_code_map.end())
 644     {
 645       opcode_index = std::distance(builtin_code_map.begin(), op_it);
 646     }
 647     // custom operator
 648     else
 649     {
 650       auto op_it = std::find(custom_code_vec.begin(), custom_code_vec.end(), operation.type());
 651       assert(op_it != custom_code_vec.end());
 652       opcode_index = builtin_code_map.size();
 653       opcode_index += std::distance(custom_code_vec.begin(), op_it);
 654     }
 655
 656     op_builder.add_opcode_index(opcode_index);
 657     op_builder.add_inputs(inputs);
 658     op_builder.add_outputs(outputs);
 659     op_builder.add_builtin_options_type(op_chef->type());
 660     op_builder.add_builtin_options(options);
 661     op_builder.add_custom_options(circle_custom_options);
 662     op_builder.add_custom_options_format(tflite::CustomOptionsFormat_FLEXBUFFERS);
 663     // Append Operator
 664     operator_vec.emplace_back(op_builder.Finish());
 665   }
 666
 667   // Create network input/output vector
 668   std::vector<int32_t> input_vec = as_dataset(graph.input()).map(lookup).vectorize();
 669   std::vector<int32_t> output_vec = as_dataset(graph.output()).map(lookup).vectorize();
 670
 671   // Create "SubGraph" arguments
 672   auto tensors = flatbuffer_builder->CreateVector(tensor_vec);
 673   auto inputs = flatbuffer_builder->CreateVector(input_vec);
 674   auto outputs = flatbuffer_builder->CreateVector(output_vec);
 675   auto operators = flatbuffer_builder->CreateVector(operator_vec);
 676   auto name = flatbuffer_builder->CreateString(graph_name);
 677
 678   tflite::SubGraphBuilder subgraph_builder{*flatbuffer_builder};
 679
 680   subgraph_builder.add_tensors(tensors);
 681   subgraph_builder.add_inputs(inputs);
 682   subgraph_builder.add_outputs(outputs);
 683   subgraph_builder.add_operators(operators);
 684   subgraph_builder.add_name(name);
 685
 686   subgraph_vec.emplace_back(subgraph_builder.Finish());
 687
 688   return symbol_table;
 689 }
 690
 691 } // namespace
 692
 693 namespace tflchef
 694 {
 695
 696 /**
 697  * @brief Generate a (in-memory) TensorFlow Lite model from a given model recipe
 698  */
 699 GeneratedModel cook(const ::tflchef::ModelRecipe &model_recipe)
 700 {
 701 // Initialize Op Chef Registry
 702 #define OP_CHEF(NAME, FACTORY_CLASS) \
 703   op_chef_registry().add(#NAME, std::unique_ptr<FACTORY_CLASS>(new FACTORY_CLASS()));
 704 #include "OpChef.def"
 705 #undef OP_CHEF
 706
 707 // Initialize Data Chef Registry
 708 #define DATA_CHEF(TYPE, NAME, FACTORY_CLASS) \
 709   data_chef_registry(::tflchef::TYPE)        \
 710     .add(#NAME, std::unique_ptr<FACTORY_CLASS>(new FACTORY_CLASS()));
 711 #include "DataChef.def"
 712 #undef DATA_CHEF
 713
 714   //
 715   // Create FlatBufferBuilder
 716   //
 717   auto flatbuffer_builder =
 718     std::unique_ptr<flatbuffers::FlatBufferBuilder>(new flatbuffers::FlatBufferBuilder(1024));
 719
 720   // Operand-related
 721   std::vector<flatbuffers::Offset<::tflite::Buffer>> buffer_vec;
 722
 723   // Operation-related
 724   std::vector<flatbuffers::Offset<::tflite::OperatorCode>> code_vec;
 725
 726   // SignatureDef-related
 727   std::vector<flatbuffers::Offset<::tflite::SignatureDef>> signdef_vec;
 728
 729   // Graphs-related
 730   std::vector<flatbuffers::Offset<::tflite::SubGraph>> subgraph_vec;
 731
 732   // Create OperatorCode with Builtin Operator
 733   auto builtin_code_map = gather_builtincode_map(model_recipe);
 734   for (auto const &opcode : builtin_code_map)
 735   {
 736     tflite::OperatorCodeBuilder code_builder{*flatbuffer_builder};
 737     // TODO support for opcode.first >= 127
 738     assert(opcode.first < 127);
 739     code_builder.add_deprecated_builtin_code(opcode.first);
 740     code_builder.add_version(opcode.second);
 741     code_builder.add_builtin_code(opcode.first);
 742     auto code = code_builder.Finish();
 743     // Update OperatorCode vector
 744     code_vec.emplace_back(code);
 745   }
 746
 747   // Create OperatorCode with Custom Operator
 748   std::set<std::string> custom_code_set = gather_customcode_set(model_recipe);
 749   std::vector<std::string> custom_code_vec{custom_code_set.begin(), custom_code_set.end()};
 750
 751   for (auto opcode : custom_code_vec)
 752   {
 753     auto custom_code = flatbuffer_builder->CreateString(opcode);
 754     tflite::OperatorCodeBuilder code_builder{*flatbuffer_builder};
 755     code_builder.add_deprecated_builtin_code(tflite::BuiltinOperator_CUSTOM);
 756     code_builder.add_custom_code(custom_code);
 757     code_builder.add_builtin_code(tflite::BuiltinOperator_CUSTOM);
 758     auto code = code_builder.Finish();
 759     // Update OperatorCode vector
 760     code_vec.emplace_back(code);
 761   }
 762
 763   // Create an Empty Buffer
 764   //
 765   // Buffer 0 SHOULD be an empty buffer in TensorFlow Lite model file
 766   // (Please refer to the comment for Tensor.buffer field in schema)
 767   {
 768     tflite::BufferBuilder buffer_builder{*flatbuffer_builder};
 769     buffer_vec.emplace_back(buffer_builder.Finish());
 770   }
 771
 772   // symbol_tables stores symbol_table of each sub graph
 773   // this is used to find tensor ID(index) with tensor name
 774   std::vector<std::map<std::string, int32_t>> symbol_tables;
 775
 776   //
 777   // Create Main graph
 778   //
 779   CookParams cp{buffer_vec,       code_vec,        subgraph_vec, flatbuffer_builder,
 780                 builtin_code_map, custom_code_vec, "main"};
 781
 782   auto table = cook_graph<::tflchef::ModelRecipe>(model_recipe, cp);
 783   symbol_tables.push_back(table);
 784
 785   //
 786   // Create subgraphs if exist
 787   //
 788   for (int g = 0; g < model_recipe.graph_size(); ++g)
 789   {
 790     const auto &graph = model_recipe.graph(g);
 791
 792     std::ostringstream stringStream;
 793     stringStream << "sub_" << (g + 1);
 794
 795     CookParams cp{buffer_vec,       code_vec,        subgraph_vec,      flatbuffer_builder,
 796                   builtin_code_map, custom_code_vec, stringStream.str()};
 797
 798     auto table = cook_graph<::tflchef::Graph>(graph, cp);
 799     symbol_tables.push_back(table);
 800   }
 801
 802   // Create Signature-Def
 803   //
 804   for (int s = 0; s < model_recipe.signature_def_size(); ++s)
 805   {
 806     // load from recipe
 807     const auto &rec_signature_def = model_recipe.signature_def(s);
 808
 809     std::vector<flatbuffers::Offset<::tflite::TensorMap>> tensormap_inputs;
 810     std::vector<flatbuffers::Offset<::tflite::TensorMap>> tensormap_outputs;
 811
 812     // which subgraph index to cook
 813     auto subgraph_index = 0;
 814     if (rec_signature_def.has_subgraph_index())
 815     {
 816       subgraph_index = rec_signature_def.subgraph_index();
 817     }
 818     assert(subgraph_index < symbol_tables.size());
 819     auto &symbol_table = symbol_tables[subgraph_index];
 820
 821     // cook for inputs
 822     for (int si = 0; si < rec_signature_def.inputs_size(); ++si)
 823     {
 824       // recipe for input TensorMap
 825       auto rec_tm_input = rec_signature_def.inputs(si);
 826       auto name = flatbuffer_builder->CreateString(rec_tm_input.name());
 827       uint32_t tensor_index = 0;
 828       // either tensor or tensor_index should exist
 829       assert(rec_tm_input.has_tensor() || rec_tm_input.has_tensor_index());
 830       if (rec_tm_input.has_tensor())
 831       {
 832         // we can get tensor_index from symbol_table
 833         auto tensor = rec_tm_input.tensor();
 834         tensor_index = symbol_table[tensor];
 835       }
 836       else
 837       {
 838         // or we can use tensor_index itself
 839         tensor_index = rec_tm_input.tensor_index();
 840       }
 841
 842       ::tflite::TensorMapBuilder tensormap_builder{*flatbuffer_builder};
 843       tensormap_builder.add_name(name);
 844       tensormap_builder.add_tensor_index(tensor_index);
 845       tensormap_inputs.push_back(tensormap_builder.Finish());
 846     }
 847     // cook for outputs, same as inputs
 848     for (int so = 0; so < rec_signature_def.outputs_size(); ++so)
 849     {
 850       auto rec_tm_output = rec_signature_def.outputs(so);
 851       auto name = flatbuffer_builder->CreateString(rec_tm_output.name());
 852       uint32_t tensor_index = 0;
 853       assert(rec_tm_output.has_tensor() || rec_tm_output.has_tensor_index());
 854       if (rec_tm_output.has_tensor())
 855       {
 856         auto tensor = rec_tm_output.tensor();
 857         tensor_index = symbol_table[tensor];
 858       }
 859       else
 860       {
 861         tensor_index = rec_tm_output.tensor_index();
 862       }
 863
 864       ::tflite::TensorMapBuilder tensormap_builder{*flatbuffer_builder};
 865       tensormap_builder.add_name(name);
 866       tensormap_builder.add_tensor_index(tensor_index);
 867       tensormap_outputs.push_back(tensormap_builder.Finish());
 868     }
 869
 870     auto inputs = flatbuffer_builder->CreateVector(tensormap_inputs);
 871     auto outputs = flatbuffer_builder->CreateVector(tensormap_outputs);
 872     auto signature_key = flatbuffer_builder->CreateString(rec_signature_def.signature_key());
 873     // TODO add validation for signature_key
 874
 875     ::tflite::SignatureDefBuilder signature_def_builder{*flatbuffer_builder};
 876     signature_def_builder.add_inputs(inputs);
 877     signature_def_builder.add_outputs(outputs);
 878     signature_def_builder.add_signature_key(signature_key);
 879     signature_def_builder.add_subgraph_index(rec_signature_def.subgraph_index());
 880
 881     signdef_vec.emplace_back(signature_def_builder.Finish());
 882   }
 883
 884   // Create "Model" arguments
 885   auto buffers = flatbuffer_builder->CreateVector(buffer_vec);
 886   auto signdefs = flatbuffer_builder->CreateVector(signdef_vec);
 887   auto operator_codes = flatbuffer_builder->CreateVector(code_vec);
 888   auto subgraphs = flatbuffer_builder->CreateVector(subgraph_vec);
 889   auto description = flatbuffer_builder->CreateString("Generated by tflchef");
 890
 891   // Create "Model"
 892   tflite::ModelBuilder model_builder{*flatbuffer_builder};
 893
 894   model_builder.add_version(3);
 895   model_builder.add_operator_codes(operator_codes);
 896   model_builder.add_signature_defs(signdefs);
 897   model_builder.add_subgraphs(subgraphs);
 898   model_builder.add_description(description);
 899   model_builder.add_buffers(buffers);
 900
 901   auto model = model_builder.Finish();
 902
 903   // Finalize
 904   ::tflite::FinishModelBuffer(*flatbuffer_builder, model);
 905
 906   // Return "GenerateModel"
 907   return GeneratedModel{
 908     std::unique_ptr<GeneratedModelImpl>(new GeneratedModelImpl(std::move(flatbuffer_builder)))};
 909 }
 910
 911 } // namespace tflchef