[moco/tf] TFConv2D to Conv2D Canonicalizer (#4202)

This will implement TFConv2D to Conv2D Canonicalizer

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

diff --git a/contrib/moco-tf/src/Canonicalization/Conv2DCanonicalizer.cpp b/contrib/moco-tf/src/Canonicalization/Conv2DCanonicalizer.cpp
index 61da9f9..13f9c56 100644 (file)
--- a/contrib/moco-tf/src/Canonicalization/Conv2DCanonicalizer.cpp
+++ b/contrib/moco-tf/src/Canonicalization/Conv2DCanonicalizer.cpp
@@ -16,19 +16,151 @@
 
 #include "Conv2DCanonicalizer.h"
 
+#include "Annotations/PadData.h"
+#include "Annotations/StrideData.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_filter_enc(loco::FilterEncode *filter_enc, moco::tf::DataLayout data_layout)
+{
+  auto enc = stdex::make_unique<loco::PermutingEncoder<loco::Domain::Filter>>();
+
+  // In TensorFlow, conv2d filter is a 4-D tensor of following shape:
+  // [filter_height, filter_width, in_channels, out_channels] -> HWIO (HWCN)
+  enc->perm()->axis(loco::FilterAxis::Height) = 0;
+  enc->perm()->axis(loco::FilterAxis::Width) = 1;
+  enc->perm()->axis(loco::FilterAxis::Depth) = 2;
+  enc->perm()->axis(loco::FilterAxis::Count) = 3;
+
+  filter_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_conv2d(loco::Graph *graph, moco::tf::TFConv2D *node)
 {
-  std::runtime_error("NYI canonicalize_conv2d");
+  if (!moco::tf::get<moco::tf::Knob::CanonicalizeConv2D>())
+    return false;
+
+  LOGGER(l);
+
+  /**
+   * @note This will replace TFCon2D node with Canonical FeatureEncode +
+   *       FilterEncode + Conv2D + FeatureDecode
+   *
+   *       Before
+   *                 A -- TFConv2D - C
+   *                 B -/
+   *
+   *       After
+   *                    - TFConv2D -
+   *                 A -- FeatureEncode - Conv2D - FeatureDecode - C
+   *                 B -- FilterEncode -/
+   *
+   *       Where
+   *                 A : ifm of TFConv2D
+   *                 B : ker of TFConv2D
+   *                 C : a node that uses TFConv2D as an input
+   *                 TFConv2D is disconnected from other nodes
+   */
+
+  auto data_layout = moco::tf::as_DataLayout(node->data_layout());
+
+  auto feature_enc = graph->nodes()->create<loco::FeatureEncode>();
+  auto filter_enc = graph->nodes()->create<loco::FilterEncode>();
+  auto conv2d = graph->nodes()->create<loco::Conv2D>();
+  auto feature_dec = graph->nodes()->create<loco::FeatureDecode>();
+
+  set_feature_enc(feature_enc, data_layout);
+  set_filter_enc(filter_enc, data_layout);
+  set_feature_dec(feature_dec, data_layout);
+
+  // Set Conv2D attributes from TFConv2D
+  auto pad_data = node->annot<moco::tf::PadData>();
+  assert(pad_data != nullptr);
+
+  conv2d->pad()->top(pad_data->pad()->top());
+  conv2d->pad()->bottom(pad_data->pad()->bottom());
+  conv2d->pad()->left(pad_data->pad()->left());
+  conv2d->pad()->right(pad_data->pad()->right());
+
+  auto stride_data = node->annot<moco::tf::StrideData>();
+  assert(stride_data != nullptr);
+
+  conv2d->stride()->vertical(stride_data->stride()->vertical());
+  conv2d->stride()->horizontal(stride_data->stride()->horizontal());
+
+  // update graph
+  auto node_A = node->ifm();
+  auto node_B = node->ker();
+
+  // update connections
+  feature_enc->input(node_A);
+  filter_enc->input(node_B);
+  conv2d->ifm(feature_enc);
+  conv2d->ker(filter_enc);
+  feature_dec->input(conv2d);
+
+  // replace and disconnect old node
+  replace(node).with(feature_dec);
+  node->ifm(nullptr);
+  node->ker(nullptr);
 
-  return false;
+  return true;
 }
 
 } // namespace