From 3e505d24a388e4d2795108f522aad218b0935be6 Mon Sep 17 00:00:00 2001
From: =?utf8?q?=EB=B0=95=EC=A2=85=ED=98=84/=EB=8F=99=EC=9E=91=EC=A0=9C?=
 =?utf8?q?=EC=96=B4Lab=28SR=29/Staff=20Engineer/=EC=82=BC=EC=84=B1?=
 =?utf8?q?=EC=A0=84=EC=9E=90?= <jh1302.park@samsung.com>
Date: Fri, 7 Dec 2018 09:34:15 +0900
Subject: [PATCH] [enco] Support Layout Conversion over ANNDepthConcatF (#2550)

Now, Data Layout Conversion pass will insert copy instructions for
layout conversion before/after ANNDepthConcatF instruction if they
are necessary.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>
---
 .../core/src/Transforms/DataLayoutConversion.cpp   | 84 ++++++++++++++++++++++
 1 file changed, 84 insertions(+)

diff --git a/contrib/enco/core/src/Transforms/DataLayoutConversion.cpp b/contrib/enco/core/src/Transforms/DataLayoutConversion.cpp
index c53c11e..e556e86 100644
--- a/contrib/enco/core/src/Transforms/DataLayoutConversion.cpp
+++ b/contrib/enco/core/src/Transforms/DataLayoutConversion.cpp
@@ -16,6 +16,9 @@
 
 #include "DataLayoutConversion.h"
 #include "Session.h"
+#include "IRUtils.h"
+
+#include "coex/IR.h"
 
 #include <coco/IR/FeatureLayouts.h>
 #include <coco/IR/KernelLayouts.h>
@@ -124,6 +127,63 @@ void insert_copy_after_eval(coco::Eval *eval)
 }
 
 /**
+ * @brief Insert copy (for data layout change) before/after ANNDepthConcatF if necessary
+ */
+void convert_data_layout(ANNDepthConcatF *concat)
+{
+  if (auto out = concat->out())
+  {
+    if (auto ofm = out->asFeature())
+    {
+      if (ofm->layout()->id() != coco::FeatureLayouts::BHWC::uid())
+      {
+        auto oldobj = ofm;
+        auto newobj = clone_feature(oldobj);
+
+        concat->out(newobj);
+
+        auto copy = make_copy(newobj, oldobj);
+        copy->insertAfter(concat);
+      }
+    }
+  }
+
+  if (auto obj = concat->fst())
+  {
+    if (auto ifm = obj->asFeature())
+    {
+      if (ifm->layout()->id() != coco::FeatureLayouts::BHWC::uid())
+      {
+        auto oldobj = ifm;
+        auto newobj = clone_feature(oldobj);
+
+        concat->fst(newobj);
+
+        auto copy = make_copy(oldobj, newobj);
+        copy->insertBefore(concat);
+      }
+    }
+  }
+
+  if (auto obj = concat->snd())
+  {
+    if (auto ifm = obj->asFeature())
+    {
+      if (ifm->layout()->id() != coco::FeatureLayouts::BHWC::uid())
+      {
+        auto oldobj = ifm;
+        auto newobj = clone_feature(oldobj);
+
+        concat->snd(newobj);
+
+        auto copy = make_copy(oldobj, newobj);
+        copy->insertBefore(concat);
+      }
+    }
+  }
+}
+
+/**
  * @brief Update convolution kernel data layout
  */
 void change_conv2d_kernel_layout(coco::Conv2D *conv)
@@ -253,6 +313,24 @@ std::set<coco::Conv2D *> convs(coco::Module *m)
   return res;
 }
 
+/**
+ * @brief Return the set of "bounded" ANNDepthConcatF instructions
+ */
+std::set<ANNDepthConcatF *> depth_concats(coco::Module *m)
+{
+  std::set<ANNDepthConcatF *> res;
+
+  for (auto ins : enco::instr_sequence(m))
+  {
+    if (auto depth_concat_f = coco::safe_cast<ANNDepthConcatF>(ins))
+    {
+      res.insert(depth_concat_f);
+    }
+  }
+
+  return res;
+}
+
 class NormalizePass
 {
 private:
@@ -281,6 +359,12 @@ void NormalizePass::runOnModule(coco::Module *m) const
   {
     change_conv2d_kernel_layout(conv);
   }
+
+  // Insert Copy (for Layout Conversion) before/after ANNDepthConcatF instructions (if necessary)
+  for (auto depth_concat : depth_concats(m))
+  {
+    convert_data_layout(depth_concat);
+  }
 }
 
 void NormalizePass::runOnCode(enco::Code *code) const { runOnModule(code->module()); }
-- 
2.7.4