Add infra to generate extra tensor

runtime/onert/backend/train/BackendContext.cc

@@ -16,12 +16,14 @@
 
 #include "BackendContext.h"
 
+#include "ExtraTensorGenerator.h"
 #include "TensorBuilder.h"
 #include "TensorPlanner.h"
 #include "KernelGenerator.h"
 #include "ops/BackPropInitializer.h"
 
 #include <backend/basic/train/TrainableBackendContextHelpers.h>
+#include <ir/train/ITrainableOperation.h>
 #include <misc/polymorphic_downcast.h>
 
 #include <cassert>
@@ -179,6 +181,32 @@ FunctionMap BackendContext::gen()
   //   fn_seq->iterate([&](exec::IFunction &ifunc) { ifunc.prepare(); });
   // }
 
+  ExtraTensorGenerator extra_tensor_gen(trainable_graph(), _tensor_builder, _tensor_registry);
+
+  const auto &ops = trainable_graph()->operations();
+
+  for (auto &pair : fn_map)
+  {
+    auto &op_idx = pair.first;
+    auto &fn_seq = pair.second;
+
+    const ir::IOperation *op = &ops.at(op_idx);
+    const auto trainable_op = dynamic_cast<const ir::train::TrainableOperation *>(op);
+    assert(trainable_op != nullptr);
+
+    if (not trainable_op->isRequiredForBackward())
+      continue;
+
+    VERBOSE(ExtraTensor) << "register tensor for " << trainable_op->name() << std::endl;
+
+    fn_seq->iterate([&](exec::train::ITrainableFunction &fn) {
+      extra_tensor_gen.register_tensors(op_idx, (&fn)->registerExtraTensors());
+    });
+  }
+  extra_tensor_gen.plan();
+  extra_tensor_gen.allocate();
+
+
   return fn_map;
 }
 

runtime/onert/backend/train/ExtraTensorGenerator.cc

@@ -0,0 +1,131 @@
+/*
+ * Copyright (c) 2024 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "ExtraTensorGenerator.h"
+
+#include "ExtraTensorIndex.h"
+
+#include <ir/Operations.h>
+#include <util/logging.h>
+#include <memory>
+
+namespace onert
+{
+namespace backend
+{
+namespace train
+{
+
+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);
+}
+
+// Move to BackendContext.cc
+void ExtraTensorGenerator::register_tensors(ir::OperationIndex op_idx,
+                                            std::optional<ExtraTensors> &&tensors)
+{
+  if (not tensors.has_value())
+    return;
+
+  auto extra_tensors = tensors.value();
+
+  auto &operations = _tgraph->operations();
+
+  for (size_t i = 0; i < extra_tensors.size(); i++)
+  {
+    // register tensor
+    ExtraTensorIndex tensor_idx(op_idx, i);
+    _tensor_builder->registerExtraTensor(tensor_idx, extra_tensors[i]);
+
+    std::stringstream op_info;
+    op_info << op_idx << "_" << operations.at(op_idx).name();
+    VERBOSE(ExtraTensorGenerator) << "register (idx:" << tensor_idx << ") requested from "
+                                  << op_info.str() << std::endl;
+  }
+  return;
+}
+
+ExtraTensors ExtraTensorGenerator::getExtraTensors(const ir::OperationIndex &op_index)
+{
+  ExtraTensors tensors;
+
+  int sub_index = 0;
+
+  ExtraTensorIndex index(op_index, sub_index);
+  auto tensor = _tensor_reg->getExtraTensor(index);
+
+  while (tensor != nullptr)
+  {
+    sub_index++;
+    tensors.push_back(tensor);
+
+    ExtraTensorIndex index(op_index, sub_index);
+    tensor = _tensor_reg->getExtraTensor(index);
+
+    VERBOSE(ExtraTensorGenerator) << "HERE " << op_index << "+" << sub_index << std::endl;
+  }
+
+  return tensors;
+}
+
+// Move to TensorPlanner
+void ExtraTensorGenerator::plan()
+{
+  // forwarding order
+  const auto f_order = _tgraph->topolSortOperations();
+  for (const auto &op_index : f_order)
+  {
+    auto tensors = getExtraTensors(op_index);
+    for (auto i = 0u; i < tensors.size(); ++i)
+    {
+      const auto &lt = tensors[i]->lifetime();
+      if (lt == ExtraTensorLifeTime::FORWARD_TO_BACKWARD)
+        _tensor_builder->notifyFirstUse(ExtraTensorIndex(op_index, i));
+    }
+  }
+
+  // backwarding order
+  const auto b_order = _tgraph->essentialBackwardOrder();
+  for (const auto &op_index : b_order)
+  {
+
+    auto tensors = getExtraTensors(op_index);
+    for (auto i = 0u; i < tensors.size(); ++i)
+    {
+      const auto &lt = tensors[i]->lifetime();
+      if (lt == ExtraTensorLifeTime::BACKWARD)
+        _tensor_builder->notifyFirstUse(ExtraTensorIndex(op_index, i));
+    }
+
+    for (auto i = 0u; i < tensors.size(); ++i)
+    {
+      const auto &lt = tensors[i]->lifetime();
+      if (lt == ExtraTensorLifeTime::FORWARD_TO_BACKWARD || lt == ExtraTensorLifeTime::BACKWARD)
+        _tensor_builder->notifyLastUse(ExtraTensorIndex(op_index, i));
+    }
+  }
+}
+
+// Move to allocateBackward()
+void ExtraTensorGenerator::allocate() { _tensor_builder->allocateExtra(); }
+
+} // namespace train
+} // namespace backend
+} // namespace onert

runtime/onert/backend/train/ExtraTensorGenerator.h

@@ -0,0 +1,60 @@
+/*
+ * Copyright (c) 2024 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef __ONERT_BACKEND_EXTRA_TENSOR_GENERATOR_H__
+#define __ONERT_BACKEND_EXTRA_TENSOR_GENERATOR_H__
+
+#include <ir/train/TrainableGraph.h>
+#include <ir/Index.h>
+#include <backend/train/ExtraTensor.h>
+
+#include "TensorBuilder.h"
+
+namespace onert
+{
+namespace backend
+{
+namespace train
+{
+
+class ExtraTensorGenerator
+{
+public:
+  ExtraTensorGenerator() = delete;
+
+  ExtraTensorGenerator(const ir::train::TrainableGraph *tgraph,
+                       std::shared_ptr<TensorBuilder> &tensor_builder,
+                       std::shared_ptr<ITensorRegistry> &tensor_registry);
+
+public:
+  // Since register is reserved keyword, use 'register_tensors' intead of 'register'
+  void register_tensors(ir::OperationIndex idx, std::optional<ExtraTensors> &&tensors);
+  void plan();
+  void allocate();
+
+private:
+  ExtraTensors getExtraTensors(const ir::OperationIndex &op_index);
+
+  const ir::train::TrainableGraph *_tgraph;
+  std::shared_ptr<TensorBuilder> _tensor_builder;
+  std::shared_ptr<TensorRegistry> _tensor_reg;
+};
+
+} // namespace train
+} // namespace backend
+} // namespace onert
+
+#endif // __ONERT_BACKEND_EXTRA_TENSOR_GENERATOR_H__

runtime/onert/backend/train/MemoryManager.cc

@@ -17,6 +17,7 @@
 #include "MemoryManager.h"
 
 #include "MemoryPlannerFactory.h"
+#include "ExtraTensorIndex.h"
 
 #include <util/ConfigSource.h>
 
@@ -53,46 +54,53 @@ uint8_t *TrainableMemoryManager::getOptVarBuffer(const ir::OperandIndex &ind,
   return _var_mem_alloc->base() + var_offset + mem_blk.offset;
 }
 
-DisposableMemoryManager::DisposableMemoryManager() : _mem_planner{createMemoryPlanner()}
+template <typename Index>
+TrainMemoryManager<Index>::TrainMemoryManager() : _mem_planner{createMemoryPlanner()}
 {
   // DO NOTHING
 }
 
-basic::IMemoryPlanner<DisposableTensorIndex> *DisposableMemoryManager::createMemoryPlanner()
+template <typename Index>
+basic::IMemoryPlanner<Index> *TrainMemoryManager<Index>::createMemoryPlanner()
 {
   auto planner_id = util::getConfigString(util::config::CPU_MEMORY_PLANNER);
-  return MemoryPlannerFactory::get().create(planner_id);
+  return MemoryPlannerFactory<Index>::get().create(planner_id);
 }
 
-basic::IMemoryPlanner<DisposableTensorIndex> *
-DisposableMemoryManager::createMemoryPlanner(const std::string planner_id)
+template <typename Index>
+basic::IMemoryPlanner<Index> *
+TrainMemoryManager<Index>::createMemoryPlanner(const std::string planner_id)
 {
-  return MemoryPlannerFactory::get().create(planner_id);
+  return MemoryPlannerFactory<Index>::get().create(planner_id);
 }
 
-void DisposableMemoryManager::claimPlan(const DisposableTensorIndex &ind, uint32_t size)
+template <typename Index> void TrainMemoryManager<Index>::claimPlan(const Index &ind, uint32_t size)
 {
   _mem_planner->claim(ind, size);
 }
 
-void DisposableMemoryManager::releasePlan(const DisposableTensorIndex &ind)
+template <typename Index> void TrainMemoryManager<Index>::releasePlan(const Index &ind)
 {
   _mem_planner->release(ind);
 }
 
-void DisposableMemoryManager::allocate(void)
+template <typename Index> void TrainMemoryManager<Index>::allocate(void)
 {
   _mem_alloc = std::make_shared<basic::Allocator>(_mem_planner->capacity());
   assert(_mem_alloc->base());
 }
 
-uint8_t *DisposableMemoryManager::getBuffer(const DisposableTensorIndex &ind) const
+template <typename Index> uint8_t *TrainMemoryManager<Index>::getBuffer(const Index &ind) const
 {
   assert(_mem_planner->memory_plans().find(ind) != _mem_planner->memory_plans().end());
   const auto &mem_blk = _mem_planner->memory_plans().at(ind);
   return _mem_alloc->base() + mem_blk.offset;
 }
 
+// Instatiation
+template class TrainMemoryManager<DisposableTensorIndex>;
+template class TrainMemoryManager<ExtraTensorIndex>;
+
 } // namespace train
 } // namespace backend
 } // namespace onert

runtime/onert/backend/train/MemoryManager.h

@@ -20,6 +20,7 @@
 #include <backend/basic/MemoryManager.h>
 
 #include "DisposableTensorIndex.h"
+#include "ExtraTensorIndex.h"
 
 namespace onert
 {
@@ -44,29 +45,33 @@ class TrainableMemoryManager : public MemoryManager
   uint32_t _optim_vars_count;
 };
 
-class DisposableMemoryManager
+// TODO: Find a better name
+template <typename Index> class TrainMemoryManager
 {
 public:
-  DisposableMemoryManager();
+  TrainMemoryManager();
 
   void allocate(void);
-  uint8_t *getBuffer(const DisposableTensorIndex &ind) const;
+  uint8_t *getBuffer(const Index &ind) const;
   void deallocate(void) { _mem_alloc->release(); }
 
-  void claimPlan(const DisposableTensorIndex &ind, uint32_t size);
-  void releasePlan(const DisposableTensorIndex &ind);
+  void claimPlan(const Index &ind, uint32_t size);
+  void releasePlan(const Index &ind);
 
   std::shared_ptr<basic::Allocator> getMemAlloc() { return _mem_alloc; }
 
 private:
-  basic::IMemoryPlanner<DisposableTensorIndex> *createMemoryPlanner();
-  basic::IMemoryPlanner<DisposableTensorIndex> *createMemoryPlanner(const std::string planner_id);
+  basic::IMemoryPlanner<Index> *createMemoryPlanner();
+  basic::IMemoryPlanner<Index> *createMemoryPlanner(const std::string planner_id);
 
 private:
-  std::shared_ptr<basic::IMemoryPlanner<DisposableTensorIndex>> _mem_planner;
+  std::shared_ptr<basic::IMemoryPlanner<Index>> _mem_planner;
   std::shared_ptr<basic::Allocator> _mem_alloc;
 };
 
+using DisposableMemoryManager = TrainMemoryManager<DisposableTensorIndex>;
+using ExtraMemoryManager = TrainMemoryManager<ExtraTensorIndex>;
+
 } // namespace train
 } // namespace backend
 } // namespace onert