From 5eabbcf51ce734bfb47d638de0d55b8bde7213cc 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: Mon, 8 Oct 2018 14:37:07 +0900
Subject: [PATCH] [coco] Tree-like Conv2D op (#1773)

* [coco] Tree-like Conv2D op

This commit introduces 'arg' methods to Conv2D class, which allows
frontend to construct a tree-like structure with Conv2D Op.

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

* Fix a typo
---
 contrib/coco/core/include/coco/IR/Conv2D.h |  9 +++++++++
 contrib/coco/core/src/IR/Conv2D.cpp        | 10 +++++++++-
 contrib/coco/core/src/IR/Conv2D.test.cpp   |  2 ++
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/contrib/coco/core/include/coco/IR/Conv2D.h b/contrib/coco/core/include/coco/IR/Conv2D.h
index 44334a4..1b6d867 100644
--- a/contrib/coco/core/include/coco/IR/Conv2D.h
+++ b/contrib/coco/core/include/coco/IR/Conv2D.h
@@ -18,6 +18,7 @@
 #define __COCO_IR_CONV2D_H__
 
 #include "coco/IR/Op.h"
+#include "coco/IR/Part.h"
 #include "coco/IR/KernelObject.h"
 #include "coco/IR/Padding2D.h"
 #include "coco/IR/Stride2D.h"
@@ -51,6 +52,10 @@ private:
   Use _ker;
 
 public:
+  Op *arg(void) const { return _arg.child(); }
+  void arg(Op *arg) { _arg.child(arg); }
+
+public:
   KernelObject *ker(void) const;
   void ker(KernelObject *ker);
 
@@ -84,6 +89,10 @@ private:
 
   Padding2D _pad;
   Stride2D _stride;
+
+private:
+  // @brief Link to an argument of Conv2D operation (= IFM)
+  Part _arg;
 };
 
 } // namespace coco
diff --git a/contrib/coco/core/src/IR/Conv2D.cpp b/contrib/coco/core/src/IR/Conv2D.cpp
index 2912cab..b835dcd 100644
--- a/contrib/coco/core/src/IR/Conv2D.cpp
+++ b/contrib/coco/core/src/IR/Conv2D.cpp
@@ -21,7 +21,7 @@
 namespace coco
 {
 
-Conv2D::Conv2D() : _ker{this}
+Conv2D::Conv2D() : _ker{this}, _arg{this}
 {
   // DO NOTHING
 }
@@ -35,6 +35,14 @@ std::set<Object *> Conv2D::uses(void) const
     res.insert(ker());
   }
 
+  if (auto ifm = arg())
+  {
+    for (auto obj : ifm->uses())
+    {
+      res.insert(obj);
+    }
+  }
+
   return res;
 }
 
diff --git a/contrib/coco/core/src/IR/Conv2D.test.cpp b/contrib/coco/core/src/IR/Conv2D.test.cpp
index dfc82a9..920e723 100644
--- a/contrib/coco/core/src/IR/Conv2D.test.cpp
+++ b/contrib/coco/core/src/IR/Conv2D.test.cpp
@@ -56,6 +56,8 @@ TEST_F(Conv2DTest, ctor)
 {
   auto op = allocate();
 
+  // arg() should be initialized as nullptr on construction
+  ASSERT_EQ(op->arg(), nullptr);
   // ker() should be initialized as nullptr on construction
   ASSERT_EQ(op->ker(), nullptr);
 
-- 
2.7.4