Apply IStage to IStageGenerator and PlanBuilder (#4414)

Apply IStage(AtomicStage) to IStageGenerator and PlanBuilder.
StageSequence will be applied for Subgraph soon.

Signed-off-by: Yongseop Kim <yons.kim@samsung.com>

diff --git a/runtimes/neurun/src/backend/interface/IStageGenerator.h b/runtimes/neurun/src/backend/interface/IStageGenerator.h
index 5966923..23f6f4c 100644 (file)
--- a/runtimes/neurun/src/backend/interface/IStageGenerator.h
+++ b/runtimes/neurun/src/backend/interface/IStageGenerator.h
@@ -23,7 +23,14 @@
 #include "exec/interface/IFunction.h"
 
 #include "backend/interface/ITensorBuilder.h"
+#include "backend/interface/IStage.h"
 #include "model/operation/NodeVisitor.h"
+#include "cpp14/memory.h"
+
+// TODO Remove dependencies for below two headers
+#include "backend/AtomicStage.h"
+// for Subgraph
+//#include "backend/StageSequence.h"
 
 struct IExecutionBuilder
 {
@@ -32,9 +39,6 @@ struct IExecutionBuilder
   virtual void append(std::unique_ptr<::neurun::exec::IFunction> &&f) = 0;
 };
 
-// TODO: Replace as IStage
-using Stage = std::function<void(IExecutionBuilder &)>;
-
 namespace neurun
 {
 namespace backend
@@ -56,17 +60,24 @@ protected:
   virtual void visit(const model::operation::Subgraph &) override {}
 
 protected:
-  void returnStage(const Stage &stage) { _return = stage; }
+  void returnStage(const StageFn fn)
+  {
+    assert(_return);
+    *_return << fn;
+  }
 
 public:
-  Stage generate(const model::operation::Node &node)
+  std::unique_ptr<IStage> generate(const model::operation::Node &node)
   {
+    // TODO Consider Subgraph and
+    // remove directly dependency for classes not interface
+    _return = nnfw::cpp14::make_unique<AtomicStage>();
     node.accept(std::move(*this));
-    return _return;
+    return std::move(_return);
   }
 
 private:
-  Stage _return = nullptr;
+  std::unique_ptr<IStage> _return = nullptr;
 };
 
 } // namespace backend

diff --git a/runtimes/neurun/src/compiler/PlanBuilder.cc b/runtimes/neurun/src/compiler/PlanBuilder.cc
index 54d8d80..da35bb0 100644 (file)
--- a/runtimes/neurun/src/compiler/PlanBuilder.cc
+++ b/runtimes/neurun/src/compiler/PlanBuilder.cc
@@ -23,7 +23,10 @@ namespace neurun
 namespace compiler
 {
 
-void PlanBuilder::addStage(const Stage &stage) { _stages.emplace_back(stage); }
+void PlanBuilder::addStage(std::unique_ptr<backend::IStage> stage)
+{
+  _stages.emplace_back(std::move(stage));
+}
 
 void PlanBuilder::finalize(const backend::TensorBuilderSet &tensor_builders)
 {
@@ -44,7 +47,7 @@ void PlanBuilder::finalize(const backend::TensorBuilderSet &tensor_builders)
 
   for (const auto &stage : _stages)
   {
-    stage(execution_builder);
+    (*stage)(execution_builder);
   }
 
   // Allocate Tensor Memory for cl_tensors

diff --git a/runtimes/neurun/src/compiler/PlanBuilder.h b/runtimes/neurun/src/compiler/PlanBuilder.h
index fcf62eb..d9bdf5f 100644 (file)
--- a/runtimes/neurun/src/compiler/PlanBuilder.h
+++ b/runtimes/neurun/src/compiler/PlanBuilder.h
@@ -21,6 +21,7 @@
 #include "compiler/operation/Sequence.h"
 #include "backend/interface/IStageGenerator.h"
 #include "backend/interface/ITensorBuilder.h"
+#include "backend/interface/IStage.h"
 
 namespace neurun
 {
@@ -55,7 +56,7 @@ public:
   }
 
 public:
-  void addStage(const Stage &stage);
+  void addStage(std::unique_ptr<backend::IStage> stage);
 
 public:
   // TODO Remove the argument `tensor_builders`
@@ -66,7 +67,7 @@ private:
   operation::Sequence &_operations;
 
 private:
-  std::vector<Stage> _stages;
+  std::vector<std::unique_ptr<backend::IStage>> _stages;
 };
 
 } // namepsace compiler