From a07ae97b979a53107e65db1431988ab90bc59a59 Mon Sep 17 00:00:00 2001
From: sseung <zetwhite@naver.com>
Date: Fri, 13 Sep 2024 17:34:18 +0900
Subject: [PATCH] [onert/backend] Add LayerScopeTensors into TensorRegistry

This PR adds LayerScopeTensors into TensorRegistry.

ONE-DCO-1.0-Signed-off-by: seunghui youn <sseung.youn@samsung.com>

draft : https://github.com/Samsung/ONE/pull/13486
for : https://github.com/Samsung/ONE/issues/13282
---
 runtime/onert/backend/train/TensorRegistry.h | 29 ++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/runtime/onert/backend/train/TensorRegistry.h b/runtime/onert/backend/train/TensorRegistry.h
index 13932199a9d..2678ce2374d 100644
--- a/runtime/onert/backend/train/TensorRegistry.h
+++ b/runtime/onert/backend/train/TensorRegistry.h
@@ -18,8 +18,10 @@
 #define __ONERT_BACKEND_TRAIN_TENSOR_REGISTRY__
 
 #include <backend/train/ITensorRegistry.h>
+#include <backend/train/LayerScopeTensor.h>
 
 #include "DisposableTensorIndex.h"
+#include "LayerScopeTensorIndex.h"
 #include "Tensor.h"
 
 namespace onert
@@ -60,9 +62,36 @@ class TensorRegistry
     return _disposable_back_prop;
   }
 
+  std::shared_ptr<LayerScopeTensor> getLayerScopeTensor(const LayerScopeTensorIndex &index)
+  {
+    auto itr = _layer_scope.find(index);
+    if (itr != _layer_scope.end())
+      return itr->second;
+
+    return nullptr;
+  }
+
+  void setLayerScopeTensor(const LayerScopeTensorIndex &index, std::shared_ptr<LayerScopeTensor> &tensor)
+  {
+    assert(tensor != nullptr);
+    auto itr = _layer_scope.find(index);
+    if (itr != _layer_scope.end())
+      throw std::runtime_error{
+        "Tried to set a layer scope tensor but another layer scope tensor already exists."};
+
+    _layer_scope[index] = tensor;
+  }
+
+  const std::unordered_map<LayerScopeTensorIndex, std::shared_ptr<LayerScopeTensor>> &
+  layerscope_tensors()
+  {
+    return _layer_scope;
+  }
+
 private:
   // Disposable Tensors to be accumulated to BackPropTensor
   std::unordered_map<DisposableTensorIndex, std::unique_ptr<BackPropTensor>> _disposable_back_prop;
+  std::unordered_map<LayerScopeTensorIndex, std::shared_ptr<LayerScopeTensor>> _layer_scope;
 };
 
 } // namespace train