[neurun] Revise IStageGenerator (#3536)

* [neurun] Revise IStageGenerator

Revise IStageGenerator to be a NodeVisitor. This will make it possible
to unify Planner implementation.

Signed-off-by: Hanjoung Lee <hanjoung.lee@samsung.com>
* Pass Stage as const reference

diff --git a/runtimes/neurun/src/backend/acl_cl/StageGenerator.cc b/runtimes/neurun/src/backend/acl_cl/StageGenerator.cc
index b173d5a..bdd5a11 100644 (file)
--- a/runtimes/neurun/src/backend/acl_cl/StageGenerator.cc
+++ b/runtimes/neurun/src/backend/acl_cl/StageGenerator.cc
@@ -120,7 +120,7 @@ StageGenerator::StageGenerator(const neurun::graph::operand::Set &ctx,
   // DO NOTHING
 }
 
-Stage StageGenerator::generate(const graph::operation::Conv2DNode &node)
+void StageGenerator::visit(const graph::operation::Conv2DNode &node)
 {
   using graph::operation::Conv2DNode;
 
@@ -181,7 +181,7 @@ Stage StageGenerator::generate(const graph::operation::Conv2DNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto ofm_alloc = tensors->at(param.ofm_index).get();
     auto ifm_alloc = tensors->at(param.ifm_index).get();
     auto ker_alloc = tensors->at(param.ker_index).get();
@@ -196,10 +196,10 @@ Stage StageGenerator::generate(const graph::operation::Conv2DNode &node)
     builder.append(std::move(fn));
 
     ActivationBuilder{builder}.append(param.activation, ofm_alloc);
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::MaxPool2DNode &node)
+void StageGenerator::visit(const graph::operation::MaxPool2DNode &node)
 {
   const auto ofm_index{node.getOutputs().at(0)};
   const auto ifm_index{node.getInputs().at(graph::operation::MaxPool2DNode::Input::INPUT)};
@@ -269,7 +269,7 @@ Stage StageGenerator::generate(const graph::operation::MaxPool2DNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto ofm_alloc = tensors->at(param.ofm_index).get();
     auto ifm_alloc = tensors->at(param.ifm_index).get();
 
@@ -282,10 +282,10 @@ Stage StageGenerator::generate(const graph::operation::MaxPool2DNode &node)
     fn->configure(ifm_alloc, ofm_alloc, info);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::AvgPool2DNode &node)
+void StageGenerator::visit(const graph::operation::AvgPool2DNode &node)
 {
   const auto ofm_index{node.getOutputs().at(0)};
   const auto ifm_index{node.getInputs().at(graph::operation::AvgPool2DNode::Input::INPUT)};
@@ -359,7 +359,7 @@ Stage StageGenerator::generate(const graph::operation::AvgPool2DNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto ofm_alloc = tensors->at(param.ofm_index).get();
     auto ifm_alloc = tensors->at(param.ifm_index).get();
 
@@ -372,10 +372,10 @@ Stage StageGenerator::generate(const graph::operation::AvgPool2DNode &node)
     fn->configure(ifm_alloc, ofm_alloc, info);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::ConcatNode &node)
+void StageGenerator::visit(const graph::operation::ConcatNode &node)
 {
   const auto ofm_index{node.getOutputs().at(0)};
   const auto axis_index{node.param().axis_index};
@@ -399,7 +399,7 @@ Stage StageGenerator::generate(const graph::operation::ConcatNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     // If tensor allocator allocate as subtensor
     bool canEliminate = true;
     for (auto ifm_ind : param.input_indexes)
@@ -430,10 +430,10 @@ Stage StageGenerator::generate(const graph::operation::ConcatNode &node)
     fn->configure(input_allocs, param.axis, output_alloc);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::FullyConnectedNode &node)
+void StageGenerator::visit(const graph::operation::FullyConnectedNode &node)
 {
   using graph::operation::FullyConnectedNode;
 
@@ -466,7 +466,7 @@ Stage StageGenerator::generate(const graph::operation::FullyConnectedNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto output_alloc = tensors->at(param.output_index).get();
     auto input_alloc = tensors->at(param.input_index).get();
     auto weight_alloc = tensors->at(param.weight_index).get();
@@ -479,10 +479,10 @@ Stage StageGenerator::generate(const graph::operation::FullyConnectedNode &node)
     builder.append(std::move(fn));
 
     ActivationBuilder{builder}.append(param.activation, output_alloc);
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::ReshapeNode &node)
+void StageGenerator::visit(const graph::operation::ReshapeNode &node)
 {
   const auto output_index{node.getOutputs().at(0)};
   const auto input_index{node.getInputs().at(graph::operation::ReshapeNode::Input::INPUT)};
@@ -500,7 +500,7 @@ Stage StageGenerator::generate(const graph::operation::ReshapeNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto output_alloc = tensors->at(param.output_index).get();
     auto input_alloc = tensors->at(param.input_index).get();
 
@@ -509,10 +509,10 @@ Stage StageGenerator::generate(const graph::operation::ReshapeNode &node)
     fn->configure(input_alloc, output_alloc);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::SoftmaxNode &node)
+void StageGenerator::visit(const graph::operation::SoftmaxNode &node)
 {
   const auto output_index{node.getOutputs().at(0)};
   const auto input_index{node.getInputs().at(graph::operation::SoftmaxNode::Input::INPUT)};
@@ -535,7 +535,7 @@ Stage StageGenerator::generate(const graph::operation::SoftmaxNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto output_alloc = tensors->at(param.output_index).get();
     auto input_alloc = tensors->at(param.input_index).get();
 
@@ -544,21 +544,21 @@ Stage StageGenerator::generate(const graph::operation::SoftmaxNode &node)
     fn->configure(input_alloc, output_alloc, param.scale);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::NOPNode & /* node */)
+void StageGenerator::visit(const graph::operation::NOPNode & /* node */)
 {
   // DO NOTHING
-  return nullptr;
+  returnStage(nullptr);
 }
 
-Stage StageGenerator::generate(const graph::operation::PermuteNode & /* node */)
+void StageGenerator::visit(const graph::operation::PermuteNode & /* node */)
 {
   throw "Unsupported";
 }
 
-Stage StageGenerator::generate(const graph::operation::AddNode &)
+void StageGenerator::visit(const graph::operation::AddNode &)
 {
   VERBOSE(Add) << "generate CPU Add" << std::endl;
 

diff --git a/runtimes/neurun/src/backend/acl_cl/StageGenerator.h b/runtimes/neurun/src/backend/acl_cl/StageGenerator.h
index 4a48086..40bd43c 100644 (file)
--- a/runtimes/neurun/src/backend/acl_cl/StageGenerator.h
+++ b/runtimes/neurun/src/backend/acl_cl/StageGenerator.h
@@ -37,16 +37,16 @@ public:
 
   virtual std::shared_ptr<ITensorBuilder> tensor_builder() override { return _tensor_builder; }
 
-  virtual Stage generate(const graph::operation::Conv2DNode &node) override;
-  virtual Stage generate(const graph::operation::MaxPool2DNode &node) override;
-  virtual Stage generate(const graph::operation::AvgPool2DNode &node) override;
-  virtual Stage generate(const graph::operation::ConcatNode &node) override;
-  virtual Stage generate(const graph::operation::FullyConnectedNode &node) override;
-  virtual Stage generate(const graph::operation::ReshapeNode &node) override;
-  virtual Stage generate(const graph::operation::SoftmaxNode &node) override;
-  virtual Stage generate(const graph::operation::NOPNode &node) override;
-  virtual Stage generate(const graph::operation::PermuteNode &node) override;
-  virtual Stage generate(const graph::operation::AddNode &node) override;
+  virtual void visit(const graph::operation::Conv2DNode &node) override;
+  virtual void visit(const graph::operation::MaxPool2DNode &node) override;
+  virtual void visit(const graph::operation::AvgPool2DNode &node) override;
+  virtual void visit(const graph::operation::ConcatNode &node) override;
+  virtual void visit(const graph::operation::FullyConnectedNode &node) override;
+  virtual void visit(const graph::operation::ReshapeNode &node) override;
+  virtual void visit(const graph::operation::SoftmaxNode &node) override;
+  virtual void visit(const graph::operation::NOPNode &node) override;
+  virtual void visit(const graph::operation::PermuteNode &node) override;
+  virtual void visit(const graph::operation::AddNode &node) override;
 
 private:
   const neurun::graph::operand::Set &_ctx;

diff --git a/runtimes/neurun/src/backend/cpu/StageGenerator.cc b/runtimes/neurun/src/backend/cpu/StageGenerator.cc
index 105324d..d2d3148 100644 (file)
--- a/runtimes/neurun/src/backend/cpu/StageGenerator.cc
+++ b/runtimes/neurun/src/backend/cpu/StageGenerator.cc
@@ -50,7 +50,7 @@ StageGenerator::StageGenerator(const neurun::graph::operand::Set &operand_ctx,
   // DO NOTHING
 }
 
-Stage StageGenerator::generate(const graph::operation::Conv2DNode &node)
+void StageGenerator::visit(const graph::operation::Conv2DNode &node)
 {
   using graph::operation::Conv2DNode;
 
@@ -119,7 +119,7 @@ Stage StageGenerator::generate(const graph::operation::Conv2DNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto ofm_alloc = tensors->at(param.ofm_index);
     auto ifm_alloc = tensors->at(param.ifm_index);
     auto ker_alloc = tensors->at(param.ker_index);
@@ -134,10 +134,10 @@ Stage StageGenerator::generate(const graph::operation::Conv2DNode &node)
                   param.stride.vertical, param.activation, ofm_alloc->buffer(), param.ofm_shape);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::MaxPool2DNode &node)
+void StageGenerator::visit(const graph::operation::MaxPool2DNode &node)
 {
   VERBOSE(MaxPool2D) << "generate CPU MaxPool2D" << std::endl;
 
@@ -217,7 +217,7 @@ Stage StageGenerator::generate(const graph::operation::MaxPool2DNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto ofm_alloc = tensors->at(param.ofm_index).get();
     auto ifm_alloc = tensors->at(param.ifm_index).get();
 
@@ -230,10 +230,10 @@ Stage StageGenerator::generate(const graph::operation::MaxPool2DNode &node)
                   param.ofm_shape);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::AvgPool2DNode &node)
+void StageGenerator::visit(const graph::operation::AvgPool2DNode &node)
 {
   VERBOSE(AvgPool2D) << "generate CPU AvgPool2D" << std::endl;
 
@@ -317,7 +317,7 @@ Stage StageGenerator::generate(const graph::operation::AvgPool2DNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto ofm_alloc = tensors->at(param.ofm_index).get();
     auto ifm_alloc = tensors->at(param.ifm_index).get();
 
@@ -330,10 +330,10 @@ Stage StageGenerator::generate(const graph::operation::AvgPool2DNode &node)
                   param.ofm_shape);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::ConcatNode &node)
+void StageGenerator::visit(const graph::operation::ConcatNode &node)
 {
   VERBOSE(Concat) << "generate CPU Concat" << std::endl;
 
@@ -369,7 +369,7 @@ Stage StageGenerator::generate(const graph::operation::ConcatNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto output_alloc = tensors->at(param.output_index).get();
 
     std::vector<const uint8_t *> input_buffers;
@@ -384,10 +384,10 @@ Stage StageGenerator::generate(const graph::operation::ConcatNode &node)
                   param.ofm_shape);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::FullyConnectedNode &node)
+void StageGenerator::visit(const graph::operation::FullyConnectedNode &node)
 {
   VERBOSE(FullyConnected) << "generate CPU FullyConnected" << std::endl;
 
@@ -431,7 +431,7 @@ Stage StageGenerator::generate(const graph::operation::FullyConnectedNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto output_alloc = tensors->at(param.output_index).get();
     auto input_alloc = tensors->at(param.input_index).get();
     auto weight_alloc = tensors->at(param.weight_index).get();
@@ -445,10 +445,10 @@ Stage StageGenerator::generate(const graph::operation::FullyConnectedNode &node)
                   output_alloc->buffer(), param.ofm_shape);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::ReshapeNode &node)
+void StageGenerator::visit(const graph::operation::ReshapeNode &node)
 {
   const auto output_index{node.getOutputs().at(0)};
   const auto input_index{node.getInputs().at(graph::operation::ReshapeNode::Input::INPUT)};
@@ -472,7 +472,7 @@ Stage StageGenerator::generate(const graph::operation::ReshapeNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto output_alloc = tensors->at(param.output_index).get();
     auto input_alloc = tensors->at(param.input_index).get();
 
@@ -482,10 +482,10 @@ Stage StageGenerator::generate(const graph::operation::ReshapeNode &node)
     fn->configure(input_alloc->buffer(), param.ifm_shape, output_alloc->buffer(), param.ofm_shape);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::SoftmaxNode &node)
+void StageGenerator::visit(const graph::operation::SoftmaxNode &node)
 {
   VERBOSE(Softmax) << "generate CPU Softmax" << std::endl;
 
@@ -516,7 +516,7 @@ Stage StageGenerator::generate(const graph::operation::SoftmaxNode &node)
 
   auto tensors = _tensor_builder;
 
-  return [tensors, param](IExecutionBuilder &builder) {
+  returnStage([tensors, param](IExecutionBuilder &builder) {
     auto output_alloc = tensors->at(param.output_index).get();
     auto input_alloc = tensors->at(param.input_index).get();
 
@@ -527,16 +527,16 @@ Stage StageGenerator::generate(const graph::operation::SoftmaxNode &node)
                   param.ofm_shape);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::NOPNode & /* node */)
+void StageGenerator::visit(const graph::operation::NOPNode & /* node */)
 {
   // DO NOTHING
-  return nullptr;
+  returnStage(nullptr);
 }
 
-Stage StageGenerator::generate(const graph::operation::PermuteNode &node)
+void StageGenerator::visit(const graph::operation::PermuteNode &node)
 {
   VERBOSE(Permute) << "generate CPU Permute" << std::endl;
 
@@ -587,7 +587,7 @@ Stage StageGenerator::generate(const graph::operation::PermuteNode &node)
   const auto input_tensors = input_backend->tensor_builder();
   const auto output_tensors = output_backend->tensor_builder();
 
-  return [input_tensors, output_tensors, param](IExecutionBuilder &builder) {
+  returnStage([input_tensors, output_tensors, param](IExecutionBuilder &builder) {
     auto output_alloc = output_tensors->tensorAt(param.output_index).get();
     auto input_alloc = input_tensors->tensorAt(param.input_index).get();
 
@@ -596,10 +596,10 @@ Stage StageGenerator::generate(const graph::operation::PermuteNode &node)
     fn->configure(input_alloc, output_alloc, param.shape, param.type);
 
     builder.append(std::move(fn));
-  };
+  });
 }
 
-Stage StageGenerator::generate(const graph::operation::AddNode &)
+void StageGenerator::visit(const graph::operation::AddNode &)
 {
   VERBOSE(Add) << "generate CPU Add" << std::endl;
 

diff --git a/runtimes/neurun/src/backend/cpu/StageGenerator.h b/runtimes/neurun/src/backend/cpu/StageGenerator.h
index cb99b1a..9006e90 100644 (file)
--- a/runtimes/neurun/src/backend/cpu/StageGenerator.h
+++ b/runtimes/neurun/src/backend/cpu/StageGenerator.h
@@ -38,16 +38,16 @@ public:
 
   virtual std::shared_ptr<ITensorBuilder> tensor_builder() override { return _tensor_builder; }
 
-  virtual Stage generate(const graph::operation::Conv2DNode &node) override;
-  virtual Stage generate(const graph::operation::MaxPool2DNode &node) override;
-  virtual Stage generate(const graph::operation::AvgPool2DNode &node) override;
-  virtual Stage generate(const graph::operation::ConcatNode &node) override;
-  virtual Stage generate(const graph::operation::FullyConnectedNode &node) override;
-  virtual Stage generate(const graph::operation::ReshapeNode &node) override;
-  virtual Stage generate(const graph::operation::SoftmaxNode &node) override;
-  virtual Stage generate(const graph::operation::NOPNode &node) override;
-  virtual Stage generate(const graph::operation::PermuteNode &node) override;
-  virtual Stage generate(const graph::operation::AddNode &node) override;
+  virtual void visit(const graph::operation::Conv2DNode &node) override;
+  virtual void visit(const graph::operation::MaxPool2DNode &node) override;
+  virtual void visit(const graph::operation::AvgPool2DNode &node) override;
+  virtual void visit(const graph::operation::ConcatNode &node) override;
+  virtual void visit(const graph::operation::FullyConnectedNode &node) override;
+  virtual void visit(const graph::operation::ReshapeNode &node) override;
+  virtual void visit(const graph::operation::SoftmaxNode &node) override;
+  virtual void visit(const graph::operation::NOPNode &node) override;
+  virtual void visit(const graph::operation::PermuteNode &node) override;
+  virtual void visit(const graph::operation::AddNode &node) override;
 
 private:
   const neurun::graph::operand::Set &_ctx;

diff --git a/runtimes/neurun/src/backend/interface/IStageGenerator.h b/runtimes/neurun/src/backend/interface/IStageGenerator.h
index 58e0272..a9008de 100644 (file)
--- a/runtimes/neurun/src/backend/interface/IStageGenerator.h
+++ b/runtimes/neurun/src/backend/interface/IStageGenerator.h
@@ -23,16 +23,7 @@
 #include <arm_compute/runtime/IFunction.h>
 
 #include "backend/interface/ITensorBuilder.h"
-#include "graph/operation/Conv2D.h"
-#include "graph/operation/MaxPool2D.h"
-#include "graph/operation/AvgPool2D.h"
-#include "graph/operation/Concat.h"
-#include "graph/operation/FullyConnected.h"
-#include "graph/operation/Reshape.h"
-#include "graph/operation/Softmax.h"
-#include "graph/operation/NOP.h"
-#include "graph/operation/Permute.h"
-#include "graph/operation/Add.h"
+#include "graph/operation/NodeVisitor.h"
 
 struct IExecutionBuilder
 {
@@ -48,22 +39,37 @@ namespace neurun
 namespace backend
 {
 
-struct IStageGenerator
+class IStageGenerator : graph::operation::NodeVisitor
 {
+public:
   virtual ~IStageGenerator() = default;
 
   virtual std::shared_ptr<ITensorBuilder> tensor_builder() = 0;
 
-  virtual Stage generate(const graph::operation::Conv2DNode &node) = 0;
-  virtual Stage generate(const graph::operation::MaxPool2DNode &node) = 0;
-  virtual Stage generate(const graph::operation::AvgPool2DNode &node) = 0;
-  virtual Stage generate(const graph::operation::ConcatNode &node) = 0;
-  virtual Stage generate(const graph::operation::FullyConnectedNode &node) = 0;
-  virtual Stage generate(const graph::operation::ReshapeNode &node) = 0;
-  virtual Stage generate(const graph::operation::SoftmaxNode &node) = 0;
-  virtual Stage generate(const graph::operation::NOPNode &node) = 0;
-  virtual Stage generate(const graph::operation::PermuteNode &node) = 0;
-  virtual Stage generate(const graph::operation::AddNode &node) = 0;
+protected:
+  virtual void visit(const graph::operation::Conv2DNode &) override {}
+  virtual void visit(const graph::operation::MaxPool2DNode &) override {}
+  virtual void visit(const graph::operation::AvgPool2DNode &) override {}
+  virtual void visit(const graph::operation::ConcatNode &) override {}
+  virtual void visit(const graph::operation::FullyConnectedNode &) override {}
+  virtual void visit(const graph::operation::ReshapeNode &) override {}
+  virtual void visit(const graph::operation::SoftmaxNode &) override {}
+  virtual void visit(const graph::operation::NOPNode &) override {}
+  virtual void visit(const graph::operation::PermuteNode &) override {}
+  virtual void visit(const graph::operation::AddNode &) override {}
+
+protected:
+  void returnStage(const Stage &stage) { _return = stage; }
+
+public:
+  Stage generate(const graph::operation::Node &node)
+  {
+    node.accept(std::move(*this));
+    return _return;
+  }
+
+private:
+  Stage _return = nullptr;
 };
 
 } // namespace backend