From b70d6eb667f82917cfb68d5b22afac0fd54651d6 Mon Sep 17 00:00:00 2001
From: SeungHui Youn <61981457+zetwhite@users.noreply.github.com>
Date: Mon, 14 Oct 2024 11:14:02 +0900
Subject: [PATCH] [onert/backend] Add LayerScopeTensor interface to
 TensorBuilder (#14198)

This PR adds LayerScopeTensor related interface to TensorBuilder.

ONE-DCO-1.0-Signed-off-by: seunghui youn <sseung.youn@samsung.com>
---
 runtime/onert/backend/train/TensorBuilder.cc | 56 ++++++++++++++++++++
 runtime/onert/backend/train/TensorBuilder.h  | 14 +++++
 2 files changed, 70 insertions(+)

diff --git a/runtime/onert/backend/train/TensorBuilder.cc b/runtime/onert/backend/train/TensorBuilder.cc
index ee737222be2..55472f3dd92 100644
--- a/runtime/onert/backend/train/TensorBuilder.cc
+++ b/runtime/onert/backend/train/TensorBuilder.cc
@@ -95,6 +95,27 @@ void TensorBuilder::registerDisposableBackwardTensorInfo(const DisposableTensorI
   _disposable_backprops.add(index);
 }
 
+void TensorBuilder::registerLayerScopeTensor(const LayerScopeTensorIndex &index,
+                                             std::shared_ptr<LayerScopeTensor> &tensor)
+{
+  const auto op_idx = index.op_index();
+
+  const auto pair = _operation_to_layerscope.find(op_idx);
+  if (pair == _operation_to_layerscope.end())
+  {
+    util::Set<LayerScopeTensorIndex> tensor_indices;
+    tensor_indices.add(index);
+    _operation_to_layerscope[op_idx] = tensor_indices;
+  }
+  else
+  {
+    assert(!pair->second.contains(index));
+    pair->second.add(index);
+  }
+
+  _tensor_reg->setLayerScopeTensor(index, tensor);
+}
+
 void TensorBuilder::notifyFirstUse(const ir::OperandIndex &index)
 {
   // TODO Support momory plan
@@ -155,6 +176,16 @@ void TensorBuilder::notifyDisposableBackPropLastUse(const DisposableTensorIndex
   _tensor_mgr->releaseDisposableBackPropPlan(index);
 }
 
+void TensorBuilder::notifyLayerScopeFirstUse(const LayerScopeTensorIndex &index)
+{
+  _tensor_mgr->claimLayerScopePlan(index);
+}
+
+void TensorBuilder::notifyLayerScopeLastUse(const LayerScopeTensorIndex &index)
+{
+  _tensor_mgr->releaseLayerScopePlan(index);
+}
+
 bool TensorBuilder::isRegistered(const ir::OperandIndex &index) const
 {
   return _tensor_info_map.find(index) != _tensor_info_map.end();
@@ -170,6 +201,29 @@ bool TensorBuilder::isRegisteredDisposableBackwardTensor(const DisposableTensorI
   return _disposable_backprops.contains(index);
 }
 
+bool TensorBuilder::isRegisteredLayerScopeTensor(const ir::OperationIndex &index) const
+{
+  const auto pair = _operation_to_layerscope.find(index);
+  return (pair != _operation_to_layerscope.end());
+}
+
+const util::Set<LayerScopeTensorIndex> &
+TensorBuilder::getRegisteredLayerScopeTensorIndices(const ir::OperationIndex &index) const
+{
+  const auto pair = _operation_to_layerscope.find(index);
+  assert(pair != _operation_to_layerscope.end());
+
+  return pair->second;
+}
+
+LayerScopeTensorLifeTime
+TensorBuilder::getLayerScopeTensorLifeTime(const LayerScopeTensorIndex &index) const
+{
+  const auto &ls_tensors = _tensor_reg->layerscope_tensors();
+  const auto &tensor = ls_tensors.at(index);
+  return tensor->lifetime();
+}
+
 void TensorBuilder::allocate(void)
 {
   _tensor_mgr->allocateNonConstTensors();
@@ -183,6 +237,8 @@ void TensorBuilder::allocateBackward(void)
   _tensor_mgr->allocateDisposableBackPropTensors();
 }
 
+void TensorBuilder::allocateLayerScope(void) { _tensor_mgr->allocateLayerScopeTensors(); }
+
 } // namespace train
 } // namespace backend
 } // namespace onert
diff --git a/runtime/onert/backend/train/TensorBuilder.h b/runtime/onert/backend/train/TensorBuilder.h
index 1fa46855142..c15d2393209 100644
--- a/runtime/onert/backend/train/TensorBuilder.h
+++ b/runtime/onert/backend/train/TensorBuilder.h
@@ -18,10 +18,12 @@
 #define __ONERT_BACKEND_TRAIN_TENSOR_BUILDER_H__
 
 #include "DisposableTensorIndex.h"
+#include "LayerScopeTensorIndex.h"
 #include "TensorManager.h"
 #include "TensorRegistry.h"
 #include "util/Set.h"
 
+#include <ir/OperationIndexMap.h>
 #include <exec/train/optimizer/Optimizer.h>
 
 namespace onert
@@ -55,6 +57,9 @@ class TensorBuilder
   void registerDisposableBackwardTensorInfo(const DisposableTensorIndex &index,
                                             const ir::OperandInfo &info);
 
+  void registerLayerScopeTensor(const LayerScopeTensorIndex &index,
+                                std::shared_ptr<LayerScopeTensor> &info);
+
   // TODO Support memory plan of all tensors
   void notifyFirstUse(const ir::OperandIndex &);
   void notifyLastUse(const ir::OperandIndex &);
@@ -62,13 +67,21 @@ class TensorBuilder
   void notifyBackwardLastUse(const ir::OperandIndex &);
   void notifyDisposableBackPropFirstUse(const DisposableTensorIndex &);
   void notifyDisposableBackPropLastUse(const DisposableTensorIndex &);
+  void notifyLayerScopeFirstUse(const LayerScopeTensorIndex &);
+  void notifyLayerScopeLastUse(const LayerScopeTensorIndex &);
 
   bool isRegistered(const ir::OperandIndex &) const;
   bool isRegisteredBackward(const ir::OperandIndex &) const;
   bool isRegisteredDisposableBackwardTensor(const DisposableTensorIndex &index) const;
+  bool isRegisteredLayerScopeTensor(const ir::OperationIndex &) const;
+
+  const util::Set<LayerScopeTensorIndex> &
+  getRegisteredLayerScopeTensorIndices(const ir::OperationIndex &) const;
+  LayerScopeTensorLifeTime getLayerScopeTensorLifeTime(const LayerScopeTensorIndex &) const;
 
   void allocate(void);
   void allocateBackward(void);
+  void allocateLayerScope(void);
 
 private:
   const std::shared_ptr<TensorRegistry> _tensor_reg;
@@ -77,6 +90,7 @@ class TensorBuilder
   ir::OperandIndexMap<ir::OperandInfo> _backward_tensor_info_map;
   ir::OperandIndexMap<bool> _as_constants;
   util::Set<DisposableTensorIndex> _disposable_backprops;
+  ir::OperationIndexMap<util::Set<LayerScopeTensorIndex>> _operation_to_layerscope;
   const exec::train::optimizer::Optimizer *_optimizer;
 };