register etensors

Samsung · Jul 29, 2024 · 5d258f4 · 5d258f4
1 parent dde2bbf
commit 5d258f4
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 2 deletions.
diff --git a/runtime/onert/backend/train/BackendContext.cc b/runtime/onert/backend/train/BackendContext.cc
@@ -231,6 +231,17 @@ FunctionMap BackendContext::genKernels()
   //   fn_seq->iterate([&](exec::IFunction &ifunc) { ifunc.prepare(); });
   // }
 
+  // register estra tensors
+  // from c++ 17, structured binding is supported
+  for (auto &[operation_index, fn_seq] : ret)
+  {
+    fn_seq->iterate([&](exec::train::ITrainableFunction &fn) {
+      extra_tensor_gen->register_tensors(operation_index, fn.requestExtraTensors());
+    });
+  }
+
+  // TODO : plan extra tensors
+
   return ret;
 }
 

diff --git a/runtime/onert/backend/train/ExtraTensorGenerator.cc b/runtime/onert/backend/train/ExtraTensorGenerator.cc
@@ -30,7 +30,8 @@ ExtraTensorGenerator::ExtraTensorGenerator(const ir::train::TrainableGraph &tgra
                                            std::shared_ptr<TensorRegistry> &tensor_registry)
   : _tgraph(tgraph), _tensor_builder(tensor_builder), _tensor_reg(tensor_registry){};
 
-void ExtraTensorGenerator::generate(ir::OperationIndex op_idx, const ExtraTensorRequests &reqs)
+void ExtraTensorGenerator::register_tensors(ir::OperationIndex op_idx,
+                                            const ExtraTensorRequests &reqs)
 {
   // save request, _idx_to_reuqests used for memory planning
   _idx_to_requests[op_idx] = reqs;

diff --git a/runtime/onert/backend/train/ExtraTensorGenerator.h b/runtime/onert/backend/train/ExtraTensorGenerator.h
@@ -39,7 +39,7 @@ class ExtraTensorGenerator
                        std::shared_ptr<TensorRegistry> &tensor_registry);
 
 public:
-  void generate(ir::OperationIndex idx, const ExtraTensorRequests &requests);
+  void register_tensors(ir::OperationIndex idx, const ExtraTensorRequests &requests);
 
 private:
   const ir::train::TrainableGraph &_tgraph;