[enco/tfl] Support concat along feature height/width (#2579)

* [enco/tfl] Support concat along feature height/width

This commit rewrites enco tflite frontend to support feature map
concatenation along height/width dimensions.

This commit also includes related tests.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>
* Use NONE instead of 0

* Fix a typo (compuate -> compute)

diff --git a/contrib/enco/frontend/tflite/src/Op/Concatenation.cpp b/contrib/enco/frontend/tflite/src/Op/Concatenation.cpp
index 4965d1b..b63654d 100644 (file)
--- a/contrib/enco/frontend/tflite/src/Op/Concatenation.cpp
+++ b/contrib/enco/frontend/tflite/src/Op/Concatenation.cpp
@@ -26,9 +26,62 @@
 #include <nncc/core/ADT/tensor/Shape.h>
 #include <schema_generated.h>
 
+#include <array>
+
 using namespace nncc::core::ADT;
 using namespace morph::tflite;
 
+namespace
+{
+
+/**
+ * @brief Convert a numeric tensor axis as a ConcatF FeatureAxis value
+ */
+coco::ConcatF::Axis as_ConcatF_axis(uint32_t axis)
+{
+  // NOTE The feature map (in TensorFlow) is a rank-4 (NHWC) tensor
+  assert(axis < 4);
+
+  coco::ConcatF::Axis res = coco::ConcatF::Axis::Unknown;
+
+  switch (axis)
+  {
+  case 0:
+    res = coco::ConcatF::Axis::Batch;
+    break;
+  case 1:
+    res = coco::ConcatF::Axis::Height;
+    break;
+  case 2:
+    res = coco::ConcatF::Axis::Width;
+    break;
+  case 3:
+    res = coco::ConcatF::Axis::Depth;
+    break;
+  default:
+    break;
+  }
+
+  return res;
+}
+
+/**
+ * @brief Convert a feature shape as an array of 'unit32_t' values
+ */
+std::array<uint32_t, 4> as_dims(const feature::Shape &shape)
+{
+  std::array<uint32_t, 4> res;
+
+  res[0] = 1; /* BATCH */
+  res[1] = shape.height();
+  res[2] = shape.width();
+  res[3] = shape.depth();
+
+  return res;
+}
+
+} // namespace
+
 namespace tflimport
 {
 
@@ -68,8 +121,7 @@ void ConcatenationGraphBuilder::build(const tflite::Operator *op,
     assert(concat_axis >= 0);
     assert(concat_axis < rank);
   }
-  // TODO handle other axis
-  assert(concat_axis == 3);
+  assert(as_ConcatF_axis(concat_axis) != coco::ConcatF::Axis::Unknown);
   assert(activation == tflite::ActivationFunctionType_NONE);
 
   // Construct a vector of input objects
@@ -115,13 +167,33 @@ void ConcatenationGraphBuilder::build(const tflite::Operator *op,
     assert(left_feature->layout()->batch() == 1);
     assert(right_feature->layout()->batch() == 1);
 
-    // Height and Width SHOULD BE IDENTICAL for depth concat
-    assert(left_shape.height() == right_shape.height());
-    assert(left_shape.width() == right_shape.width());
+    // Compute output dimensionalities
+    auto compute_out_dims = [&left_shape, &right_shape, concat_axis](void) {
+      std::array<uint32_t, 4> out_dims;
+
+      const auto left_dims = as_dims(left_shape);
+      const auto right_dims = as_dims(right_shape);
+
+      for (uint32_t axis = 0; axis < 4 /* FEATURE MAP RANK */; ++axis)
+      {
+        // The dimensionality of all the axises except 'concat' axis SHOULD BE INDETICAL
+        assert((concat_axis == axis) || (left_dims[axis] == right_dims[axis]));
+
+        out_dims[axis] = left_dims[axis];
+        if (axis == concat_axis)
+        {
+          out_dims[axis] += right_dims[axis];
+        }
+      }
+
+      return out_dims;
+    };
+
+    const auto out_dims = compute_out_dims();
 
-    const uint32_t C = left_shape.depth() + right_shape.depth();
-    const uint32_t H = left_shape.height();
-    const uint32_t W = left_shape.width();
+    const uint32_t C = out_dims[3 /* DEPTH */];
+    const uint32_t H = out_dims[1 /* HEIGHT */];
+    const uint32_t W = out_dims[2 /* WIDTH */];
 
     const nncc::core::ADT::feature::Shape out_shape{C, H, W};
 
@@ -136,7 +208,7 @@ void ConcatenationGraphBuilder::build(const tflite::Operator *op,
 
     auto concat_f = m->entity()->op()->create<coco::ConcatF>();
 
-    concat_f->axis(coco::ConcatF::Axis::Depth);
+    concat_f->axis(as_ConcatF_axis(concat_axis));
     concat_f->left(left_load);
     concat_f->right(right_load);
 

diff --git a/contrib/enco/test/tflite/Concat_001/INFERENCE b/contrib/enco/test/tflite/Concat_001/INFERENCE
new file mode 100644 (file)
index 0000000..e69de29

diff --git a/contrib/enco/test/tflite/Concat_001/test.recipe b/contrib/enco/test/tflite/Concat_001/test.recipe
new file mode 100644 (file)
index 0000000..7adaf16
--- /dev/null
+++ b/contrib/enco/test/tflite/Concat_001/test.recipe
@@ -0,0 +1,29 @@
+# Concatenate two feature maps along "width" dimension
+operand {
+  name: "ifm1"
+  type: FLOAT32
+  shape { dim: 1 dim: 1 dim: 1 dim: 1 }
+}
+operand {
+  name: "ifm2"
+  type: FLOAT32
+  shape { dim: 1 dim: 1 dim: 2 dim: 1 }
+}
+operand {
+  name: "ofm"
+  type: FLOAT32
+  shape { dim: 1 dim: 1 dim: 3 dim: 1 }
+}
+operation {
+  type: "Concatenation"
+  concatenation_options {
+    axis: 2
+    activation: NONE
+  }
+  input: "ifm1"
+  input: "ifm2"
+  output: "ofm"
+}
+input: "ifm1"
+input: "ifm2"
+output: "ofm"

diff --git a/contrib/enco/test/tflite/Concat_002/INFERENCE b/contrib/enco/test/tflite/Concat_002/INFERENCE
new file mode 100644 (file)
index 0000000..e69de29

diff --git a/contrib/enco/test/tflite/Concat_002/test.recipe b/contrib/enco/test/tflite/Concat_002/test.recipe
new file mode 100644 (file)
index 0000000..918cb13
--- /dev/null
+++ b/contrib/enco/test/tflite/Concat_002/test.recipe
@@ -0,0 +1,29 @@
+# Concatenate two feature maps along "height" dimension
+operand {
+  name: "ifm1"
+  type: FLOAT32
+  shape { dim: 1 dim: 1 dim: 1 dim: 1 }
+}
+operand {
+  name: "ifm2"
+  type: FLOAT32
+  shape { dim: 1 dim: 2 dim: 1 dim: 1 }
+}
+operand {
+  name: "ofm"
+  type: FLOAT32
+  shape { dim: 1 dim: 3 dim: 1 dim: 1 }
+}
+operation {
+  type: "Concatenation"
+  concatenation_options {
+    axis: 1
+    activation: NONE
+  }
+  input: "ifm1"
+  input: "ifm2"
+  output: "ofm"
+}
+input: "ifm1"
+input: "ifm2"
+output: "ofm"