[coco] Introduce OpManager::destroy_all (#2498)

With this commit, OpManager now allows users to destroy a Op tree via
destroy_all method.

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

diff --git a/contrib/coco/core/include/coco/IR/OpManager.h b/contrib/coco/core/include/coco/IR/OpManager.h
index 95ac849..95769f9 100644 (file)
--- a/contrib/coco/core/include/coco/IR/OpManager.h
+++ b/contrib/coco/core/include/coco/IR/OpManager.h
@@ -44,6 +44,11 @@ public:
   //
   // NOTE destroy(op) WILL NOT update op->parent(). Client SHOULD detach op before destroy(op) call
   void destroy(Op *);
+
+  // @brief Destroy a Op tree
+  //
+  // @require op->parent() == nullptr && op->up() == nullptr
+  void destroy_all(Op *);
 };
 
 } // namespace coco

diff --git a/contrib/coco/core/src/IR/OpManager.cpp b/contrib/coco/core/src/IR/OpManager.cpp
index 5f8965c..ce94b2d 100644 (file)
--- a/contrib/coco/core/src/IR/OpManager.cpp
+++ b/contrib/coco/core/src/IR/OpManager.cpp
@@ -19,6 +19,7 @@
 #include <nncc/foundation/Memory.h>
 
 #include <cassert>
+#include <queue>
 
 using nncc::foundation::make_unique;
 
@@ -44,4 +45,29 @@ void OpManager::destroy(Op *op)
   release(op);
 }
 
+void OpManager::destroy_all(Op *op)
+{
+  assert(op->parent() == nullptr);
+  assert(op->up() == nullptr);
+
+  std::queue<coco::Op *> q;
+
+  q.emplace(op);
+
+  while (q.size() > 0)
+  {
+    auto cur = q.front();
+    q.pop();
+
+    // Insert child op nodes
+    for (uint32_t n = 0; n < cur->arity(); ++n)
+    {
+      q.emplace(cur->arg(n));
+    }
+
+    // Destroy the current op node
+    destroy(cur);
+  }
+}
+
 } // namespace coco

diff --git a/contrib/coco/core/src/IR/OpManager.test.cpp b/contrib/coco/core/src/IR/OpManager.test.cpp
index daaea4c..3afc2da 100644 (file)
--- a/contrib/coco/core/src/IR/OpManager.test.cpp
+++ b/contrib/coco/core/src/IR/OpManager.test.cpp
@@ -74,3 +74,16 @@ TEST_F(OpManagerTest, destroy)
   mgr.destroy(op);
   ASSERT_EQ(mgr.size(), 0);
 }
+
+TEST_F(OpManagerTest, destroy_all)
+{
+  // Create a Op tree
+  auto load_op = mgr.create<coco::Load>();
+  auto conv_op = mgr.create<coco::Conv2D>();
+
+  conv_op->arg(load_op);
+
+  mgr.destroy_all(conv_op);
+
+  ASSERT_EQ(mgr.size(), 0);
+}