From 1e49d4238aa162286dd331dbc21c35b7a5ca2c8d 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, 30 Aug 2018 14:28:43 +0900
Subject: [PATCH] [enco] Identify input/output of ANN subnet (#1243)

This commit revises Split pass to identify input/ouput of android NN
sub-network based on how block and bags are related.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>
---
 contrib/enco/core/src/Transforms/Split.cpp | 56 +++++++++++++++++++++++++++++-
 1 file changed, 55 insertions(+), 1 deletion(-)

diff --git a/contrib/enco/core/src/Transforms/Split.cpp b/contrib/enco/core/src/Transforms/Split.cpp
index b46d8c3..0009609 100644
--- a/contrib/enco/core/src/Transforms/Split.cpp
+++ b/contrib/enco/core/src/Transforms/Split.cpp
@@ -132,6 +132,8 @@ void ANNGroupBuilder::build(void) const
 
 } // namespace
 
+#include "Usage.h"
+
 #include <stdexcept>
 
 #include <nncc/core/ADT/kernel/NHWCLayout.h>
@@ -245,6 +247,10 @@ public:
     // DO NOTHING
   }
 
+private:
+  std::set<coco::Bag *> inputs(coco::Block *blk) const;
+  std::set<coco::Bag *> outputs(coco::Block *blk) const;
+
 public:
   void build(void) const;
 
@@ -252,6 +258,52 @@ private:
   enco::Code *_code;
 };
 
+std::set<coco::Bag *> ANNModuleBuilder::inputs(coco::Block *blk) const
+{
+  std::set<coco::Bag *> res;
+
+  for (auto bag : enco::reads(blk))
+  {
+    auto u = enco::updaters(bag);
+    u.erase(blk);
+
+    /**
+     * A bag is the input of this block if
+     *  1. it is an input of the whole network, or
+     *  2. it is updated by preceding blocks during execution
+     */
+    if (bag->isInput() || (u.size() > 0))
+    {
+      res.insert(bag);
+    }
+  }
+
+  return res;
+}
+
+std::set<coco::Bag *> ANNModuleBuilder::outputs(coco::Block *blk) const
+{
+  std::set<coco::Bag *> res;
+
+  for (auto bag : enco::updates(blk))
+  {
+    auto r = enco::readers(bag);
+    r.erase(blk);
+
+    /**
+     * A bag is the output of this block if
+     *  1. it is an output of the whole network, or
+     *  2. it is read by following blocks during execution
+     */
+    if (bag->isOutput() || (r.size() > 0))
+    {
+      res.insert(bag);
+    }
+  }
+
+  return res;
+}
+
 void ANNModuleBuilder::build(void) const
 {
   for (uint32_t n = 0; n < _code->ann()->count(); ++n)
@@ -266,7 +318,9 @@ void ANNModuleBuilder::build(void) const
       ins->accept(op_builder);
     }
 
-    // TODO identify input/output bags
+    // Let's identify input/output bags
+    binder->identifyInputs(inputs(binder->block()));
+    binder->identifyOutputs(outputs(binder->block()));
   }
 }
 
-- 
2.7.4