[exo.tflite] Manage TFLTensorIndex as annotation (#4162)

The current implementation manages TFLTensorIndex via SubGraphContext,
which necessitates some-level of indirection.

This commit makes exo-tflite to annotate each node with its tensor
index.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>

diff --git a/contrib/exo-tflite/src/ExporterUtils.cpp b/contrib/exo-tflite/src/ExporterUtils.cpp
index 7f6d615..fc3234e 100644 (file)
--- a/contrib/exo-tflite/src/ExporterUtils.cpp
+++ b/contrib/exo-tflite/src/ExporterUtils.cpp
@@ -59,3 +59,39 @@ void registerGraphIOName(loco::Graph *graph, SerializedModelData &gd)
     gd._push_to_name[push] = name;
   }
 }
+
+#include <stdex/Memory.h>
+
+#include <cassert>
+
+namespace
+{
+
+class TFLTensorIndexAnnotation final : public loco::NodeAnnotation
+{
+public:
+  TFLTensorIndexAnnotation(const TFLTensorIndex &index) : _index{index}
+  {
+    // DO NOTHING
+  }
+
+public:
+  const TFLTensorIndex &index(void) const { return _index; }
+
+private:
+  TFLTensorIndex _index;
+};
+
+} // namespace
+
+void set_tensor_index(loco::Node *node, const TFLTensorIndex &tensor_id)
+{
+  assert(node->annot<TFLTensorIndexAnnotation>() == nullptr);
+  node->annot(stdex::make_unique<TFLTensorIndexAnnotation>(tensor_id));
+}
+
+TFLTensorIndex get_tensor_index(loco::Node *node)
+{
+  assert(node->annot<TFLTensorIndexAnnotation>() != nullptr);
+  return node->annot<TFLTensorIndexAnnotation>()->index();
+}

diff --git a/contrib/exo-tflite/src/ExporterUtils.h b/contrib/exo-tflite/src/ExporterUtils.h
index 054b560..7ca2430 100644 (file)
--- a/contrib/exo-tflite/src/ExporterUtils.h
+++ b/contrib/exo-tflite/src/ExporterUtils.h
@@ -64,8 +64,6 @@ struct SubGraphContext
   std::vector<int32_t> _inputs;
   /// @brief SubGraph output tensor id
   std::vector<int32_t> _outputs;
-  /// @brief Mapping from loco Node to SubGraph tensor id
-  std::unordered_map<loco::Node *, int32_t> _node_to_tensor_id;
 };
 
 // Prerequisites for tflite::Model object creation
@@ -110,4 +108,9 @@ tflite::Padding getOpPadding(const loco::Pad<2> *pad);
 /// @brief Register graph input and output names to SerializedModelData
 void registerGraphIOName(loco::Graph *graph, SerializedModelData &gd);
 
+using TFLTensorIndex = int32_t;
+
+void set_tensor_index(loco::Node *node, const TFLTensorIndex &tensor_id);
+TFLTensorIndex get_tensor_index(loco::Node *node);
+
 #endif // __EXPORTER_UTILS_H__

diff --git a/contrib/exo-tflite/src/OperationExporter.cpp b/contrib/exo-tflite/src/OperationExporter.cpp
index 0309edb..ada47f3 100644 (file)
--- a/contrib/exo-tflite/src/OperationExporter.cpp
+++ b/contrib/exo-tflite/src/OperationExporter.cpp
@@ -26,8 +26,8 @@ namespace
 void exportRelu(loco::ReLU *node, FlatBufferBuilder &builder, SerializedModelData &gd)
 {
   uint32_t op_idx = gd.registerBuiltinOpcode(tflite::BuiltinOperator_RELU);
-  std::vector<int32_t> inputs_vec{gd._node_to_tensor_id[node->input()]};
-  std::vector<int32_t> outputs_vec{gd._node_to_tensor_id[static_cast<loco::Node *>(node)]};
+  std::vector<int32_t> inputs_vec{get_tensor_index(node->input())};
+  std::vector<int32_t> outputs_vec{get_tensor_index(static_cast<loco::Node *>(node))};
   auto inputs = builder.CreateVector(inputs_vec);
   auto outputs = builder.CreateVector(outputs_vec);
   auto op_offset = CreateOperator(builder, op_idx, inputs, outputs);
@@ -37,8 +37,8 @@ void exportRelu(loco::ReLU *node, FlatBufferBuilder &builder, SerializedModelDat
 void exportMaxPool2D(loco::MaxPool2D *node, FlatBufferBuilder &builder, SerializedModelData &gd)
 {
   uint32_t op_idx = gd.registerBuiltinOpcode(tflite::BuiltinOperator_MAX_POOL_2D);
-  std::vector<int32_t> inputs_vec{gd._node_to_tensor_id[node->ifm()]};
-  std::vector<int32_t> outputs_vec{gd._node_to_tensor_id[static_cast<loco::Node *>(node)]};
+  std::vector<int32_t> inputs_vec{get_tensor_index(node->ifm())};
+  std::vector<int32_t> outputs_vec{get_tensor_index(static_cast<loco::Node *>(node))};
   auto inputs = builder.CreateVector(inputs_vec);
   auto outputs = builder.CreateVector(outputs_vec);
   tflite::Padding padding = getOpPadding(node->pad());
@@ -56,8 +56,8 @@ void exportAvgPool2D(loco::AvgPool2D *node, FlatBufferBuilder &builder, Serializ
   assert(node->convention() == loco::AvgPool2D::Convention::Valid);
 
   uint32_t op_idx = gd.registerBuiltinOpcode(tflite::BuiltinOperator_AVERAGE_POOL_2D);
-  std::vector<int32_t> inputs_vec{gd._node_to_tensor_id[node->ifm()]};
-  std::vector<int32_t> outputs_vec{gd._node_to_tensor_id[static_cast<loco::Node *>(node)]};
+  std::vector<int32_t> inputs_vec{get_tensor_index(node->ifm())};
+  std::vector<int32_t> outputs_vec{get_tensor_index(static_cast<loco::Node *>(node))};
   auto inputs = builder.CreateVector(inputs_vec);
   auto outputs = builder.CreateVector(outputs_vec);
   tflite::Padding padding = getOpPadding(node->pad());
@@ -101,9 +101,9 @@ void exportConv2D(loco::Conv2D *node, FlatBufferBuilder &builder, SerializedMode
   gd._tensors.push_back(bias_tensor_offset);
 
   // Make input, output and options for operator
-  std::vector<int32_t> inputs_vec{gd._node_to_tensor_id[node->ifm()],
-                                  gd._node_to_tensor_id[node->ker()], bias_tensor_id};
-  std::vector<int32_t> outputs_vec{gd._node_to_tensor_id[static_cast<loco::Node *>(node)]};
+  std::vector<int32_t> inputs_vec{get_tensor_index(node->ifm()), get_tensor_index(node->ker()),
+                                  bias_tensor_id};
+  std::vector<int32_t> outputs_vec{get_tensor_index(static_cast<loco::Node *>(node))};
   auto inputs = builder.CreateVector(inputs_vec);
   auto outputs = builder.CreateVector(outputs_vec);
   tflite::Padding padding = getOpPadding(node->pad());
@@ -121,8 +121,8 @@ template <typename NodeT>
 void exportIdentity(NodeT *node, FlatBufferBuilder &builder, SerializedModelData &gd)
 {
   uint32_t op_idx = gd.registerBuiltinOpcode(tflite::BuiltinOperator_CONCATENATION);
-  std::vector<int32_t> inputs_vec{gd._node_to_tensor_id[node->arg(0)]};
-  std::vector<int32_t> outputs_vec{gd._node_to_tensor_id[static_cast<loco::Node *>(node)]};
+  std::vector<int32_t> inputs_vec{get_tensor_index(node->arg(0))};
+  std::vector<int32_t> outputs_vec{get_tensor_index(static_cast<loco::Node *>(node))};
   auto inputs = builder.CreateVector(inputs_vec);
   auto outputs = builder.CreateVector(outputs_vec);
   auto options = CreateConcatenationOptions(builder); // use dummy 0 axis and NONE activation
@@ -163,8 +163,8 @@ void exportAsTranspose(loco::Node *node, FlatBufferBuilder &builder,
 
   // Create permutation node
 
-  std::vector<int32_t> inputs_vec{gd._node_to_tensor_id[node->arg(0)], perm_tensor_id};
-  std::vector<int32_t> outputs_vec{gd._node_to_tensor_id[node]};
+  std::vector<int32_t> inputs_vec{get_tensor_index(node->arg(0)), perm_tensor_id};
+  std::vector<int32_t> outputs_vec{get_tensor_index(node)};
 
   auto inputs = builder.CreateVector(inputs_vec);
   auto outputs = builder.CreateVector(outputs_vec);
@@ -250,9 +250,8 @@ void exportBiasAdd(loco::BiasAdd<loco::Domain::Tensor> *node, FlatBufferBuilder
                    SerializedModelData &gd)
 {
   uint32_t op_idx = gd.registerBuiltinOpcode(tflite::BuiltinOperator_ADD);
-  std::vector<int32_t> inputs_vec{gd._node_to_tensor_id[node->value()],
-                                  gd._node_to_tensor_id[node->bias()]};
-  std::vector<int32_t> outputs_vec{gd._node_to_tensor_id[static_cast<loco::Node *>(node)]};
+  std::vector<int32_t> inputs_vec{get_tensor_index(node->value()), get_tensor_index(node->bias())};
+  std::vector<int32_t> outputs_vec{get_tensor_index(static_cast<loco::Node *>(node))};
   auto inputs = builder.CreateVector(inputs_vec);
   auto outputs = builder.CreateVector(outputs_vec);
   auto options = CreateAddOptions(builder); // dummy option
@@ -265,9 +264,8 @@ void exportBiasAdd(loco::BiasAdd<loco::Domain::Tensor> *node, FlatBufferBuilder
 void exportConcat(loco::TensorConcat *node, FlatBufferBuilder &builder, SerializedModelData &gd)
 {
   uint32_t op_idx = gd.registerBuiltinOpcode(tflite::BuiltinOperator_CONCATENATION);
-  std::vector<int32_t> inputs_vec{gd._node_to_tensor_id[node->lhs()],
-                                  gd._node_to_tensor_id[node->rhs()]};
-  std::vector<int32_t> outputs_vec{gd._node_to_tensor_id[static_cast<loco::Node *>(node)]};
+  std::vector<int32_t> inputs_vec{get_tensor_index(node->lhs()), get_tensor_index(node->rhs())};
+  std::vector<int32_t> outputs_vec{get_tensor_index(static_cast<loco::Node *>(node))};
   auto inputs = builder.CreateVector(inputs_vec);
   auto outputs = builder.CreateVector(outputs_vec);
   auto options = CreateConcatenationOptions(builder, node->axis());
@@ -286,11 +284,11 @@ void exportNode(loco::Node *node, flatbuffers::FlatBufferBuilder &builder,
   }
   else if (dynamic_cast<loco::Pull *>(node))
   {
-    data._inputs.push_back(data._node_to_tensor_id[node]);
+    data._inputs.push_back(get_tensor_index(node));
   }
   else if (dynamic_cast<loco::Push *>(node))
   {
-    data._outputs.push_back(data._node_to_tensor_id[node->arg(0)]);
+    data._outputs.push_back(get_tensor_index(node->arg(0)));
   }
   else if (auto *encode = dynamic_cast<loco::FeatureEncode *>(node))
   {

diff --git a/contrib/exo-tflite/src/TensorExporter.cpp b/contrib/exo-tflite/src/TensorExporter.cpp
index 9ad5afc..76f05b0 100644 (file)
--- a/contrib/exo-tflite/src/TensorExporter.cpp
+++ b/contrib/exo-tflite/src/TensorExporter.cpp
@@ -94,7 +94,7 @@ void exportOpDefinedTensor(NodeT *node, FlatBufferBuilder &builder, SerializedMo
   auto name_offset = builder.CreateString(name);
   auto tensor_offset = CreateTensor(builder, shape_offset, tensor_type, buffer_id, name_offset,
                                     /*quantization*/ 0, /*is_variable*/ false);
-  gd._node_to_tensor_id[node] = tensor_id;
+  set_tensor_index(node, tensor_id);
   gd._tensors.push_back(tensor_offset);
 }