[coco] Support Bag replace (#1467)

This commit introduces replaceWith method in Bag, which allows users to
replace all the occurences of a bag in IR with another bag.

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

diff --git a/contrib/coco/core/include/coco/IR/Bag.h b/contrib/coco/core/include/coco/IR/Bag.h
index d9a54a8..b3e4bf8 100644 (file)
--- a/contrib/coco/core/include/coco/IR/Bag.h
+++ b/contrib/coco/core/include/coco/IR/Bag.h
@@ -65,6 +65,12 @@ public:
   // @brief Return the set of Update links that point to this bag
   const UpdateSet *updates(void) const;
 
+public:
+  // @brief Replace all the occurence of a bag (except those in Input/Output) with another bag
+  //
+  // NOTE reaplceWith(b) works correctly only when b is neither Input nor Output
+  void replaceWith(Bag *b);
+
 private:
   std::unique_ptr<BagInfo> _info;
 };

diff --git a/contrib/coco/core/src/IR.test.cpp b/contrib/coco/core/src/IR.test.cpp
index c368b5c..5371a42 100644 (file)
--- a/contrib/coco/core/src/IR.test.cpp
+++ b/contrib/coco/core/src/IR.test.cpp
@@ -229,3 +229,33 @@ TEST(IR, caffe_conv)
     }
   }
 }
+
+//
+// This test demonstrates how to use 'replaceWith' method
+//
+TEST(IR, bag_replaceWith)
+{
+  auto m = coco::Module::create();
+
+  auto bag_1 = m->entity()->bag()->create(1);
+  auto bag_2 = m->entity()->bag()->create(1);
+
+  auto obj = m->entity()->object()->create(nncc::core::ADT::feature::Shape{1, 1, 1});
+  obj->bag(bag_1);
+
+  auto shuffle_1 = m->entity()->instr()->create<coco::Shuffle>();
+  shuffle_1->into(bag_1);
+
+  auto shuffle_2 = m->entity()->instr()->create<coco::Shuffle>();
+  shuffle_2->from(bag_1);
+
+  ASSERT_EQ(obj->bag(), bag_1);
+  ASSERT_EQ(shuffle_1->into(), bag_1);
+  ASSERT_EQ(shuffle_2->from(), bag_1);
+
+  bag_1->replaceWith(bag_2);
+
+  ASSERT_EQ(obj->bag(), bag_2);
+  ASSERT_EQ(shuffle_1->into(), bag_2);
+  ASSERT_EQ(shuffle_2->from(), bag_2);
+}

diff --git a/contrib/coco/core/src/IR/Bag.cpp b/contrib/coco/core/src/IR/Bag.cpp
index 1e2e4e4..4038028 100644 (file)
--- a/contrib/coco/core/src/IR/Bag.cpp
+++ b/contrib/coco/core/src/IR/Bag.cpp
@@ -34,6 +34,39 @@ const DepSet *Bag::deps(void) const { return _info->deps(); }
 const ReadSet *Bag::reads(void) const { return _info->reads(); }
 const UpdateSet *Bag::updates(void) const { return _info->updates(); }
 
+void Bag::replaceWith(Bag *b)
+{
+  assert(!isInput() && !isOutput());
+
+  // Replace all the occurence inside Dep
+  while (!(deps()->empty()))
+  {
+    auto dep = *(deps()->begin());
+    assert(dep->bag() == this);
+    dep->bag(b);
+  }
+
+  // Replace all the occurence inside Read
+  while (!(reads()->empty()))
+  {
+    auto read = *(reads()->begin());
+    assert(read->bag() == this);
+    read->bag(b);
+  }
+
+  // Replace all the occurence insider Update
+  while (!(updates()->empty()))
+  {
+    auto update = *(updates()->begin());
+    assert(update->bag() == this);
+    update->bag(b);
+  }
+
+  assert(deps()->empty());
+  assert(reads()->empty());
+  assert(updates()->empty());
+}
+
 ObjectSet dependent_objects(const Bag *b)
 {
   ObjectSet res;