tmp

runtime/onert/backend/train/Backend.h

@@ -60,7 +60,7 @@ class Backend : public ::onert::backend::Backend, public backend::train::ITraina
                                                            std::move(optimizer));
 
     context->kernel_gen = std::make_shared<train::KernelGenerator>(
-      tgraph, tr, context->external_context(), context->optimizer());
+      tgraph, tr, tb, context->external_context(), context->optimizer());
     return context;
   }
 

runtime/onert/backend/train/BackendContext.cc

@@ -152,8 +152,8 @@ backend::train::ITensorRegistry *BackendContext::genTrainingTensors()
   }
 
   // generate and plan extra tensors
-  ExtraTensorGenerator extra_tensor_generator(tgraph, _tensor_builder, _tensor_registry);
-  extra_tensor_generator.generate();
+  // ExtraTensorGenerator extra_tensor_generator(tgraph, _tensor_builder, _tensor_registry);
+  // extra_tensor_generator.generate();
 
   planDisposableBackPropTensors();
 

runtime/onert/backend/train/ExtraTensorGenerator.cc

@@ -64,27 +64,18 @@ bool is_bwd_or_fwd2bwd(const ExtraTensorLifeTime lifetime)
 
 ExtraTensorGenerator::ExtraTensorGenerator(const ir::train::TrainableGraph &tgraph,
                                            std::shared_ptr<TensorBuilder> &tensor_builder,
-                                           std::shared_ptr<ITensorRegistry> &tensor_registry)
-  : _tgraph(tgraph), _tensor_builder(tensor_builder)
-{
-  _tensor_reg = std::dynamic_pointer_cast<TensorRegistry>(tensor_registry);
-
-  for (const auto &index : _tgraph.topolSortOperations())
-  {
-    const auto &node = _tgraph.operation(index);
-    _node_to_idx[&node] = index;
-  }
-};
+                                           std::shared_ptr<TensorRegistry> &tensor_registry)
+  : _tgraph(tgraph), _tensor_builder(tensor_builder), _tensor_reg(tensor_registry){};
 
-void ExtraTensorGenerator::generate()
+/* void ExtraTensorGenerator::generate()
 {
   _tgraph.operations().iterate([&](const ir::OperationIndex &, const ir::IOperation &op) {
     const auto trainable_op = dynamic_cast<const ir::train::TrainableOperation *>(&op);
     trainable_op->accept(*this);
   });
 
   handle_requests();
-}
+}*/
 
 void ExtraTensorGenerator::notify_first_use(ir::OperationIndex op_idx,
                                             const ExtraTensorRequests &reqs,
@@ -110,6 +101,24 @@ void ExtraTensorGenerator::notify_last_use(ir::OperationIndex op_idx,
   return;
 }
 
+std::vector<ExtraTensor *> ExtraTensorGenerator::register_requests(ir::OperationIndex op_idx,
+                                                                   const ExtraTensorRequests &reqs)
+{
+  std::vector<ExtraTensor *> registered_tensors;
+
+  // save requests
+  _idx_to_requests[op_idx] = reqs;
+
+  // register tensor
+  for (size_t i = 0; i < reqs.size(); i++)
+  {
+    ExtraTensorIndex tensor_idx(op_idx, i);
+    _tensor_builder->registerExtraTensorInfo(tensor_idx, reqs[i].info, reqs[i].layout);
+    registered_tensors.push_back(_tensor_reg->getExtraTensor(tensor_idx));
+  }
+  return registered_tensors;
+}
+/*
 void ExtraTensorGenerator::handle_requests()
 {
   // register tensor
@@ -186,7 +195,7 @@ void ExtraTensorGenerator::visit(const ir::train::operation::Conv2D &node)
   _idx_to_requests.emplace(op_idx, requests);
 
   return;
-}
+}*/
 
 /*
 void ExtraTensorGenerator::visit(const ir::train::operation::Pool2D &node)

runtime/onert/backend/train/ExtraTensorGenerator.h

@@ -35,42 +35,29 @@ namespace backend
 namespace train
 {
 
-class ExtraTensorGenerator : public ir::train::TrainableOperationVisitor
+class ExtraTensorGenerator
 {
 public:
   ExtraTensorGenerator() = delete;
   ExtraTensorGenerator(const ir::train::TrainableGraph &tgraph,
                        std::shared_ptr<TensorBuilder> &tensor_builder,
-                       std::shared_ptr<ITensorRegistry> &tensor_registry);
-
-public:
-  void visit(const ir::train::operation::FullyConnected &node) override;
-  // void visit(const ir::train::operation::BinaryArithmetic &node) override;
-  void visit(const ir::train::operation::Conv2D &node) override;
-  // void visit(const train::operation::DepthwiseConv2D &node) override;
-  // void visit(const train::operation::ElementwiseActivation &node) override;
-  // void visit(const train::operation::Loss &node) override;
-  // void visit(const train::operation::Pad &node) override;
-  // void visit(const ir::train::operation::Pool2D &node) override;
-  // void visit(const train::operation::Reduce &node) override;
-  // void visit(const train::operation::Reshape &node) override;
-  // void visit(const train::operation::Softmax &node) override;
-
-public:
-  void generate();
+                       std::shared_ptr<TensorRegistry> &tensor_registry);
 
 private:
   void notify_first_use(ir::OperationIndex, const ExtraTensorRequests &,
                         bool (*cond)(const ExtraTensorLifeTime));
   void notify_last_use(ir::OperationIndex, const ExtraTensorRequests &,
                        bool (*cond)(const ExtraTensorLifeTime));
-  void handle_requests();
+
+public:
+  std::vector<ExtraTensor *> register_requests(ir::OperationIndex idx,
+                                               const ExtraTensorRequests &requests);
 
 private:
   const ir::train::TrainableGraph &_tgraph;
   std::shared_ptr<TensorBuilder> _tensor_builder;
   std::shared_ptr<TensorRegistry> _tensor_reg;
-  std::unordered_map<const ir::train::ITrainableOperation *, ir::OperationIndex> _node_to_idx;
+  // std::unordered_map<const ir::train::ITrainableOperation *, ir::OperationIndex> _node_to_idx;
   std::unordered_map<ir::OperationIndex, ExtraTensorRequests> _idx_to_requests;
 };
 

runtime/onert/backend/train/KernelGenerator.cc

@@ -109,6 +109,14 @@ generateGradientApplier(const exec::train::optimizer::Optimizer *optimizer,
 }
 } // namespace
 
+void KernelGenerator::genExtraTensors(ir::OperationIndex idx,
+                                      onert::exec::train::ITrainableFunction *layer)
+{
+  auto reqs = layer->requestExtraTensors();
+  auto extra_tensors = _extra_tensor_generator->register_requests(idx, reqs);
+  layer->ConfigureExtraTensors(extra_tensors);
+}
+
 std::unique_ptr<exec::train::TrainableFnSequence> KernelGenerator::generate(ir::OperationIndex idx)
 {
   auto ret = std::make_unique<exec::train::TrainableFnSequence>();
@@ -121,6 +129,8 @@ std::unique_ptr<exec::train::TrainableFnSequence> KernelGenerator::generate(ir::
 
   op.accept(*this);
   assert(_return_fn);
+  genExtraTensors(idx, _return_fn.get());
+
   ret->append(std::move(_return_fn));
 
   for (auto &&update_fn : _update_funcs)
@@ -144,12 +154,14 @@ std::unique_ptr<exec::train::TrainableFnSequence> KernelGenerator::generate(ir::
 }
 
 KernelGenerator::KernelGenerator(const ir::train::TrainableGraph &tgraph,
-                                 const std::shared_ptr<TensorRegistry> &tensor_reg,
+                                 std::shared_ptr<TensorRegistry> &tensor_reg,
+                                 std::shared_ptr<TensorBuilder> &tensor_builder,
                                  const std::shared_ptr<ExternalContext> &external_context,
                                  const exec::train::optimizer::Optimizer *optimizer)
   : backend::train::KernelGeneratorBase{tgraph}, _current_layout{tgraph.layout()},
     _tensor_reg{tensor_reg}, _external_context(external_context), _optimizer{optimizer},
-    _update_funcs{}, _node_to_idx{}
+    _update_funcs{}, _node_to_idx{},
+    _extra_tensor_generator{new ExtraTensorGenerator(tgraph, tensor_builder, tensor_reg)}
 {
   tgraph.operations().iterate(
     [&](const onert::ir::OperationIndex &idx, const onert::ir::IOperation &op) {
@@ -240,9 +252,6 @@ void KernelGenerator::visit(const ir::train::operation::Conv2D &node)
     fn->configureBackward(ker_tensor, in_back_prop_tensor, ker_grad_tensor, bias_grad_tensor,
                           out_back_prop_tensor, activation);
 
-    auto extra_tensors = getExtraTensors(node);
-    fn->configureExtraTensors(extra_tensors);
-
     // Generate GradientApplier
     if (bias_tensor)
       _update_funcs.emplace_back(

runtime/onert/backend/train/KernelGenerator.h

@@ -18,6 +18,7 @@
 #define __ONERT_BACKEND_TRAIN_KERNEL_GENERATOR_H__
 
 #include "ExternalContext.h"
+#include "ExtraTensorGenerator.h"
 #include "backend/basic/TensorRegistry.h"
 #include "TensorBuilder.h"
 #include "Tensor.h"
@@ -40,7 +41,8 @@ class KernelGenerator : public backend::train::KernelGeneratorBase
 {
 public:
   KernelGenerator(const ir::train::TrainableGraph &tgraph,
-                  const std::shared_ptr<TensorRegistry> &tensor_reg,
+                  std::shared_ptr<TensorRegistry> &tensor_reg,
+                  std::shared_ptr<TensorBuilder> &tensor_builder,
                   const std::shared_ptr<ExternalContext> &external_context,
                   const exec::train::optimizer::Optimizer *optimizer);
 
@@ -63,7 +65,7 @@ class KernelGenerator : public backend::train::KernelGeneratorBase
                                  const ir::OperandIndex &operand_index);
   IPortableTensor *getBackPropOut(const ir::OperandIndex &index);
 
-  std::vector<ExtraTensor *> getExtraTensors(const ir::Operation &);
+  void genExtraTensors(ir::OperationIndex idx, onert::exec::train::ITrainableFunction *layer);
 
 private:
   ir::Layout _current_layout;
@@ -72,6 +74,8 @@ class KernelGenerator : public backend::train::KernelGeneratorBase
   const exec::train::optimizer::Optimizer *_optimizer;
   std::vector<std::unique_ptr<exec::train::IGradientApplier>> _update_funcs;
   std::unordered_map<const ir::IOperation *, ir::OperationIndex> _node_to_idx;
+
+  std::unique_ptr<ExtraTensorGenerator> _extra_tensor_generator;
 };
 
 } // namespace train

runtime/onert/backend/train/ops/BinaryArithmeticLayer.cc

@@ -41,39 +41,37 @@ BinaryArithmeticLayer::BinaryArithmeticLayer()
   // DO NOTHING
 }
 
-ExtraTensorRequests
-BinaryArithmeticLayer::requestExtraTensor(const IPortableTensor *back_prop_output,
-                                          ir::Activation activation)
+void BinaryArithmeticLayer::configureBackward(IPortableTensor *back_prop_lhs,
+                                              IPortableTensor *back_prop_rhs,
+                                              const IPortableTensor *back_prop_output,
+                                              const ir::Activation activation,
+                                              const ArithmeticType arithmetic_type)
+{
+  _back_prop_lhs = back_prop_lhs;
+  _back_prop_rhs = back_prop_rhs;
+  _back_prop_output = back_prop_output;
+  _arithmetic_type = arithmetic_type;
+  _activation = activation;
+}
+
+ExtraTensorRequests BinaryArithmeticLayer::requestExtraTensor()
 {
   ExtraTensorRequests req;
 
-  if (activation != ir::Activation::NONE)
-    req.push_back(ExtraTensorRequest::createRequestLike(back_prop_output));
+  if (_activation != ir::Activation::NONE)
+    req.push_back(ExtraTensorRequest::createRequestLike(_back_prop_output));
 
   return req;
 }
 
 void BinaryArithmeticLayer::configureExtraTensors(std::vector<ExtraTensor *> extra_tensors)
 {
-  if (extra_tensors.size() == 1)
+  if (_activation != ir::Activation::NONE)
     _act_back_prop_output = extra_tensors[0];
 
   return;
 }
 
-void BinaryArithmeticLayer::configureBackward(IPortableTensor *back_prop_lhs,
-                                              IPortableTensor *back_prop_rhs,
-                                              const IPortableTensor *back_prop_output,
-                                              const ir::Activation activation,
-                                              const ArithmeticType arithmetic_type)
-{
-  _back_prop_lhs = back_prop_lhs;
-  _back_prop_rhs = back_prop_rhs;
-  _back_prop_output = back_prop_output;
-  _arithmetic_type = arithmetic_type;
-  _activation = activation;
-}
-
 void BinaryArithmeticLayer::forward(bool) { cpu::ops::BinaryArithmeticLayer::run(); }
 
 void BinaryArithmeticLayer::backward()

runtime/onert/backend/train/ops/BinaryArithmeticLayer.h

@@ -52,8 +52,8 @@ class BinaryArithmeticLayer : public ::onert::exec::train::ITrainableFunction,
                          const IPortableTensor *back_prop_output, const ir::Activation activation,
                          const ArithmeticType arithmetic_type);
 
-  static ExtraTensorRequests requestExtraTensor(const IPortableTensor *back_prop_output,
-                                                ir::Activation activation);
+  ExtraTensorRequests requestExtraTensor(const IPortableTensor *back_prop_output,
+                                         ir::Activation activation);
   void configureExtraTensors(std::vector<ExtraTensor *> extra_tensors);
 
   void forward(bool training) override;

runtime/onert/backend/train/ops/ConvolutionLayer.h

@@ -43,9 +43,9 @@ class ConvolutionLayer : public ::onert::exec::train::ITrainableFunction,
                          IPortableTensor *grad_weights, IPortableTensor *grad_bias,
                          const IPortableTensor *back_prop_output, const ir::Activation activation);
 
-  static ExtraTensorRequests requestExtraTensors(const IPortableTensor *weights,
-                                                 const IPortableTensor *back_prop_out,
-                                                 const ir::Activation activation);
+  ExtraTensorRequests requestExtraTensors(const IPortableTensor *weights,
+                                          const IPortableTensor *back_prop_out,
+                                          const ir::Activation activation);
   void configureExtraTensors(std::vector<ExtraTensor *> extra_tensors);
 
   void forward(bool training) override;

runtime/onert/core/include/exec/train/ITrainableFunction.h

@@ -18,6 +18,8 @@
 #define __ONERT_EXEC_TRAIN_I_TRAINABLE_FUNCTION_H__
 
 #include <cstdint>
+#include "backend/train/ExtraTensorRequest.h"
+#include "backend/train/Tensor.h"
 
 namespace onert
 {
@@ -32,6 +34,9 @@ class ITrainableFunction
   virtual ~ITrainableFunction() = default;
   virtual void forward(bool training) = 0;
   virtual void backward() = 0;
+
+  virtual backend::train::ExtraTensorRequests requestExtraTensors();
+  virtual void ConfigureExtraTensors(std::vector<backend::train::ExtraTensor *> extra_tensors);
 };
 
 } // namespace train

runtime/onert/core/src/compiler/ExecutorFactory.cc

@@ -873,6 +873,8 @@ exec::IExecutor *ExecutorFactory::createTrainableExecutor(
       code_map.insert(
         {op_ind, train::TrainableCodeAndInfo{op_ind, &op, lower_info, std::move(tn_seq)}});
     }
+
+    // allocate extra tensors
   }
 
   // [idea2] : Generate Extra Tensors by traversing code_map