From 05f0b40c261f9e5ecafab4b7ddc75a1115359a83 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: Thu, 11 Oct 2018 09:43:44 +0900
Subject: [PATCH] [coco] AvgPool2D as a Op tree node (#1816)

With this commit, AvgPool2D may serve as a node of Op tree.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>
---
 contrib/coco/core/include/coco/IR/AvgPool2D.h |  9 +++++++++
 contrib/coco/core/src/IR/AvgPool2D.cpp        | 11 ++++++++++-
 contrib/coco/core/src/IR/AvgPool2D.test.cpp   |  3 +++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/contrib/coco/core/include/coco/IR/AvgPool2D.h b/contrib/coco/core/include/coco/IR/AvgPool2D.h
index 7295cc1..2978edb 100644
--- a/contrib/coco/core/include/coco/IR/AvgPool2D.h
+++ b/contrib/coco/core/include/coco/IR/AvgPool2D.h
@@ -18,6 +18,7 @@
 #define __COCO_IR_AVG_POOL_2D_H__
 
 #include "coco/IR/Op.h"
+#include "coco/IR/Part.h"
 #include "coco/IR/Window2D.h"
 #include "coco/IR/Padding2D.h"
 #include "coco/IR/Stride2D.h"
@@ -55,6 +56,10 @@ public:
   const AvgPool2D *asAvgPool2D(void) const override { return this; }
 
 public:
+  Op *arg(void) const { return _arg.child(); }
+  void arg(Op *arg) { _arg.child(arg); }
+
+public:
   Divisor divisor(void) const { return _divisor; }
   void divisor(const Divisor &divisor) { _divisor = divisor; }
 
@@ -71,6 +76,10 @@ public:
   const Stride2D *stride(void) const { return &_stride; }
 
 private:
+  // @brief An argument of AvgPool2D operation (= IFM)
+  Part _arg;
+
+private:
   Divisor _divisor = Divisor::Unknown;
 
   Window2D _window;
diff --git a/contrib/coco/core/src/IR/AvgPool2D.cpp b/contrib/coco/core/src/IR/AvgPool2D.cpp
index 1dcf09b..2b70a12 100644
--- a/contrib/coco/core/src/IR/AvgPool2D.cpp
+++ b/contrib/coco/core/src/IR/AvgPool2D.cpp
@@ -19,7 +19,7 @@
 namespace coco
 {
 
-AvgPool2D::AvgPool2D()
+AvgPool2D::AvgPool2D() : _arg{this}
 {
   // DO NOTHING
 }
@@ -28,6 +28,15 @@ std::set<Object *> AvgPool2D::uses(void) const
 {
   // NOTE AvgPool2D accesses no object except IFM/OFM
   std::set<Object *> res;
+
+  if (auto ifm = arg())
+  {
+    for (auto obj : ifm->uses())
+    {
+      res.insert(obj);
+    }
+  }
+
   return res;
 }
 
diff --git a/contrib/coco/core/src/IR/AvgPool2D.test.cpp b/contrib/coco/core/src/IR/AvgPool2D.test.cpp
index 5fdda62..c94f01f 100644
--- a/contrib/coco/core/src/IR/AvgPool2D.test.cpp
+++ b/contrib/coco/core/src/IR/AvgPool2D.test.cpp
@@ -61,6 +61,9 @@ TEST_F(AvgPool2DTest, initialization)
   // parent() should be nullptr on construction
   ASSERT_EQ(op->parent(), nullptr);
 
+  // arg() should be nullptr on construction
+  ASSERT_EQ(immutable_ptr->arg(), nullptr);
+
   // divisor() SHOULD be unknow on construction
   ASSERT_EQ(immutable_ptr->divisor(), coco::AvgPool2D::Divisor::Unknown);
 
-- 
2.7.4