From c5c64c6290a2d88ed49131d07f46a6b1f8d8c39f 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, 12 Nov 2018 19:08:38 +0900
Subject: [PATCH] [coco] Introduce arity and general arg accessor (#2200)

* [coco] Introduce arity and general arg accessor

With this commit, users are able to query the number of arguments over
each coco::Op, and retrive its Nth argument without knowing its details.

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

* Fix a typo
---
 contrib/coco/core/include/coco/IR/Op.h  | 18 ++++++++++++++++++
 contrib/coco/core/include/coco/IR/Ops.h |  6 ++++++
 contrib/coco/core/src/IR/Conv2D.cpp     | 13 +++++++++++++
 contrib/coco/core/src/IR/Load.cpp       | 12 ++++++++++++
 contrib/coco/core/src/IR/Op.cpp         | 25 +++++++++++++++++++++++++
 contrib/coco/core/src/IR/Part.test.cpp  |  4 ++++
 6 files changed, 78 insertions(+)

diff --git a/contrib/coco/core/include/coco/IR/Op.h b/contrib/coco/core/include/coco/IR/Op.h
index 5ea8653..d21ece5 100644
--- a/contrib/coco/core/include/coco/IR/Op.h
+++ b/contrib/coco/core/include/coco/IR/Op.h
@@ -49,6 +49,18 @@ struct Op : public Entity
   virtual ~Op();
 
   /**
+   * @brief Return the number of arguments (# of child Ops)
+   */
+  virtual uint32_t arity(void) const = 0;
+
+  /**
+   * @brief Return N-th argument
+   *
+   * @note The behavior of arg(n) is defined only when n < artiy()
+   */
+  virtual Op *arg(uint32_t n) const = 0;
+
+  /**
    * @brief Return a set of object(s) used during execution
    *
    * NOTE There is no 'def' method as Op is not allowed to define a new object
@@ -188,6 +200,9 @@ public:
   virtual ~UnaryOp() = default;
 
 public:
+  uint32_t arity(void) const final;
+  Op *arg(uint32_t n) const final;
+
   std::set<Object *> uses(void) const final;
 
 public:
@@ -215,6 +230,9 @@ public:
   virtual ~BinaryOp() = default;
 
 public:
+  uint32_t arity(void) const final;
+  Op *arg(uint32_t n) const final;
+
   std::set<Object *> uses(void) const final;
 
 public:
diff --git a/contrib/coco/core/include/coco/IR/Ops.h b/contrib/coco/core/include/coco/IR/Ops.h
index c197a17..d76e8f8 100644
--- a/contrib/coco/core/include/coco/IR/Ops.h
+++ b/contrib/coco/core/include/coco/IR/Ops.h
@@ -44,6 +44,9 @@ public:
   Load(Load &&) = delete;
 
 public:
+  uint32_t arity(void) const final;
+  Op *arg(uint32_t n) const final;
+
   std::set<Object *> uses(void) const override;
 
 public:
@@ -73,6 +76,9 @@ public:
   explicit Conv2D();
 
 public:
+  uint32_t arity(void) const final;
+  Op *arg(uint32_t n) const final;
+
   std::set<Object *> uses(void) const override;
 
 public:
diff --git a/contrib/coco/core/src/IR/Conv2D.cpp b/contrib/coco/core/src/IR/Conv2D.cpp
index 3d0b131..58ffff5 100644
--- a/contrib/coco/core/src/IR/Conv2D.cpp
+++ b/contrib/coco/core/src/IR/Conv2D.cpp
@@ -26,6 +26,19 @@ Conv2D::Conv2D() : _ker{this}, _arg{this}
   // DO NOTHING
 }
 
+uint32_t Conv2D::arity(void) const
+{
+  // Conv2D has one argument (IFM)
+  // NOTE This design is subject to change
+  return 1;
+}
+
+Op *Conv2D::arg(uint32_t n) const
+{
+  assert(n < arity());
+  return arg();
+}
+
 std::set<Object *> Conv2D::uses(void) const
 {
   std::set<Object *> res;
diff --git a/contrib/coco/core/src/IR/Load.cpp b/contrib/coco/core/src/IR/Load.cpp
index ac87483..4985e92 100644
--- a/contrib/coco/core/src/IR/Load.cpp
+++ b/contrib/coco/core/src/IR/Load.cpp
@@ -26,6 +26,18 @@ Load::Load() : _obj{this}
   // DO NOTHING
 }
 
+uint32_t Load::arity(void) const
+{
+  // Load has no child Op
+  return 0;
+}
+
+Op *Load::arg(uint32_t) const
+{
+  assert(!"Load has no argument");
+  return nullptr;
+}
+
 std::set<Object *> Load::uses(void) const
 {
   std::set<Object *> res;
diff --git a/contrib/coco/core/src/IR/Op.cpp b/contrib/coco/core/src/IR/Op.cpp
index 1a783d9..c320608 100644
--- a/contrib/coco/core/src/IR/Op.cpp
+++ b/contrib/coco/core/src/IR/Op.cpp
@@ -67,6 +67,18 @@ UnaryOp::UnaryOp() : _arg{this}
   // DO NOTHING
 }
 
+uint32_t UnaryOp::arity(void) const
+{
+  // There is only one argument
+  return 1;
+}
+
+Op *UnaryOp::arg(uint32_t n) const
+{
+  assert(n < 1);
+  return arg();
+}
+
 std::set<Object *> UnaryOp::uses(void) const
 {
   std::set<Object *> res;
@@ -90,6 +102,19 @@ BinaryOp::BinaryOp() : _left{this}, _right{this}
   // DO NOTHING
 }
 
+uint32_t BinaryOp::arity(void) const
+{
+  // There are two arguments
+  return 2;
+}
+
+Op *BinaryOp::arg(uint32_t n) const
+{
+  assert(n < arity());
+
+  return (n == 0) ? left() : right();
+}
+
 std::set<Object *> BinaryOp::uses(void) const
 {
   std::set<Object *> res;
diff --git a/contrib/coco/core/src/IR/Part.test.cpp b/contrib/coco/core/src/IR/Part.test.cpp
index e328111..3f11d7e 100644
--- a/contrib/coco/core/src/IR/Part.test.cpp
+++ b/contrib/coco/core/src/IR/Part.test.cpp
@@ -28,6 +28,7 @@ namespace
 namespace mock
 {
 
+// TODO Inherit UnaryOp instead of Op
 struct Op final : public coco::Op
 {
 public:
@@ -37,6 +38,9 @@ public:
   }
 
 public:
+  uint32_t arity(void) const final { return 1; }
+  coco::Op *arg(uint32_t n) const final { return arg(); }
+
   std::set<coco::Object *> uses() const override { throw std::runtime_error{"Not supported"}; }
 
 public:
-- 
2.7.4