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 &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 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 & +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 #include 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 &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 & + getRegisteredLayerScopeTensorIndices(const ir::OperationIndex &) const; + LayerScopeTensorLifeTime getLayerScopeTensorLifeTime(const LayerScopeTensorIndex &) const; void allocate(void); void allocateBackward(void); + void allocateLayerScope(void); private: const std::shared_ptr _tensor_reg; @@ -77,6 +90,7 @@ class TensorBuilder ir::OperandIndexMap _backward_tensor_info_map; ir::OperandIndexMap _as_constants; util::Set _disposable_backprops; + ir::OperationIndexMap> _operation_to_layerscope; const exec::train::optimizer::Optimizer *_optimizer; };