[moco-tf] Introduce AvgPool Canonicalizer (#5819)

This will introduce AvgPool node Canonicalizier

Signed-off-by: SaeHie Park <saehie.park@samsung.com>

diff --git a/compiler/moco-tf/src/Canonicalization/AvgPoolCanonicalizer.cpp b/compiler/moco-tf/src/Canonicalization/AvgPoolCanonicalizer.cpp
new file mode 100644 (file)
index 0000000..2d45260
--- /dev/null
+++ b/compiler/moco-tf/src/Canonicalization/AvgPoolCanonicalizer.cpp
@@ -0,0 +1,187 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "AvgPoolCanonicalizer.h"
+
+#include "Annotations/PadData.h"
+#include "Annotations/StrideData.h"
+#include "Annotations/ShapeInferenceData.h"
+#include "Annotations/WindowData.h"
+
+#include "Knob.h"
+
+#include "Dialect/TFDialect.h"
+#include "Dialect/TFNodes.h"
+#include "Dialect/TFNodeVisitor.h"
+#include "Dialect/TFNodeImpl.h"
+
+#include <moco/Log.h>
+
+#include <stdex/Memory.h>
+
+namespace
+{
+
+void set_feature_enc(loco::FeatureEncode *feature_enc, moco::tf::DataLayout data_layout)
+{
+  auto enc = stdex::make_unique<loco::PermutingEncoder<loco::Domain::Feature>>();
+
+  if (data_layout == moco::tf::DataLayout::NHWC)
+  {
+    enc->perm()->axis(loco::FeatureAxis::Count) = 0;
+    enc->perm()->axis(loco::FeatureAxis::Height) = 1;
+    enc->perm()->axis(loco::FeatureAxis::Width) = 2;
+    enc->perm()->axis(loco::FeatureAxis::Depth) = 3;
+  }
+  else if (data_layout == moco::tf::DataLayout::NCHW)
+  {
+    enc->perm()->axis(loco::FeatureAxis::Count) = 0;
+    enc->perm()->axis(loco::FeatureAxis::Depth) = 1;
+    enc->perm()->axis(loco::FeatureAxis::Height) = 2;
+    enc->perm()->axis(loco::FeatureAxis::Width) = 3;
+  }
+
+  feature_enc->encoder(std::move(enc));
+}
+
+void set_feature_dec(loco::FeatureDecode *feature_dec, moco::tf::DataLayout data_layout)
+{
+  auto dec = stdex::make_unique<loco::PermutingDecoder<loco::Domain::Feature>>();
+
+  if (data_layout == moco::tf::DataLayout::NHWC)
+  {
+    dec->perm()->axis(loco::FeatureAxis::Count) = 0;
+    dec->perm()->axis(loco::FeatureAxis::Height) = 1;
+    dec->perm()->axis(loco::FeatureAxis::Width) = 2;
+    dec->perm()->axis(loco::FeatureAxis::Depth) = 3;
+  }
+  else if (data_layout == moco::tf::DataLayout::NCHW)
+  {
+    dec->perm()->axis(loco::FeatureAxis::Count) = 0;
+    dec->perm()->axis(loco::FeatureAxis::Depth) = 1;
+    dec->perm()->axis(loco::FeatureAxis::Height) = 2;
+    dec->perm()->axis(loco::FeatureAxis::Width) = 3;
+  }
+
+  feature_dec->decoder(std::move(dec));
+}
+
+bool canonicalize_avgpool2d(loco::Graph *graph, moco::tf::TFAvgPool *node)
+{
+  LOGGER(l);
+
+  /**
+   * @note This will replace TFAvgPool node with Canonical FeatureEncode +
+   *       AvgPool2D + FeatureDecode
+   *
+   *       Before
+   *                 A -- TFAvgPool -- C
+   *
+   *       After
+   *                    +- TFAvgPool --
+   *                    |
+   *                 A -+- FeatureEncode -- AvgPool2D -- FeatureDecode -- C
+   *
+   *       Where
+   *                 A : value of TFAvgPool
+   *                 C : a node that uses TFAvgPool as an input
+   *                 TFAvgPool is disconnected from other nodes
+   */
+
+  auto data_layout = moco::tf::as_DataLayout(node->data_layout());
+
+  auto feature_enc = graph->nodes()->create<loco::FeatureEncode>();
+  auto avgPool2d_node = graph->nodes()->create<loco::AvgPool2D>();
+  auto feature_dec = graph->nodes()->create<loco::FeatureDecode>();
+
+  set_feature_enc(feature_enc, data_layout);
+  set_feature_dec(feature_dec, data_layout);
+
+  avgPool2d_node->convention(loco::AvgPool2D::Convention::Valid);
+
+  // paddata to pad
+  auto pad_data = node->annot<moco::tf::PadData>();
+  assert(pad_data != nullptr);
+
+  avgPool2d_node->pad()->top(pad_data->pad()->top());
+  avgPool2d_node->pad()->bottom(pad_data->pad()->bottom());
+  avgPool2d_node->pad()->left(pad_data->pad()->left());
+  avgPool2d_node->pad()->right(pad_data->pad()->right());
+
+  // windowdata to window (ksize to window)
+  auto window_data = node->annot<moco::tf::WindowData>();
+  assert(window_data != nullptr);
+
+  auto window = avgPool2d_node->window();
+  window->vertical(window_data->window()->vertical());
+  window->horizontal(window_data->window()->horizontal());
+
+  // stridedata to stride (strides to stride)
+  auto stride_data = node->annot<moco::tf::StrideData>();
+  assert(stride_data != nullptr);
+
+  auto stride = avgPool2d_node->stride();
+  stride->vertical(stride_data->stride()->vertical());
+  stride->horizontal(stride_data->stride()->horizontal());
+
+  INFO(l) << "Canonicalize TFAvgPool pad = T " << avgPool2d_node->pad()->top() << ", L "
+          << avgPool2d_node->pad()->left() << ", B " << avgPool2d_node->pad()->bottom() << ", R "
+          << avgPool2d_node->pad()->right() << std::endl;
+
+  // update graph
+  auto node_A = node->value();
+
+  // update connections
+  feature_enc->input(node_A);
+  avgPool2d_node->ifm(feature_enc);
+  feature_dec->input(avgPool2d_node);
+
+  // replace node
+  replace(node).with(feature_dec);
+
+  return true;
+}
+
+} // namespace
+
+namespace moco
+{
+namespace tf
+{
+
+bool AvgPoolCanonicalizer::run(loco::Graph *graph)
+{
+  auto active_nodes = loco::active_nodes(loco::output_nodes(graph));
+  bool changed = false;
+
+  for (auto node : active_nodes)
+  {
+    if (node->dialect() == TFDialect::get())
+    {
+      auto tf_node = dynamic_cast<moco::tf::TFAvgPool *>(node);
+      if (tf_node != nullptr)
+      {
+        if (canonicalize_avgpool2d(graph, tf_node))
+          changed = true;
+      }
+    }
+  }
+
+  return changed;
+}
+
+} // namespace tf
+} // namespace moco

diff --git a/compiler/moco-tf/src/Canonicalization/AvgPoolCanonicalizer.h b/compiler/moco-tf/src/Canonicalization/AvgPoolCanonicalizer.h
new file mode 100644 (file)
index 0000000..7d7e6a8
--- /dev/null
+++ b/compiler/moco-tf/src/Canonicalization/AvgPoolCanonicalizer.h
@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef __MOCO_TF_AVGPOOL_CANONICALIZER_H__
+#define __MOCO_TF_AVGPOOL_CANONICALIZER_H__
+
+#include "Transform.h"
+
+#include <loco.h>
+
+namespace moco
+{
+namespace tf
+{
+
+/**
+ * @brief  Convert TFAvgPool to Canonical AvgPool2D
+ */
+class AvgPoolCanonicalizer : public Transform
+{
+public:
+  const char *name(void) const final { return "AvgPoolCanonicalizer"; }
+
+public:
+  bool run(loco::Graph *graph) override;
+};
+
+} // namespace tf
+} // namespace moco
+
+#endif // __MOCO_TF_AVGPOOL_CANONICALIZER_H__