ONE - On-device Neural Engine
Loading...
Searching...
No Matches
Public Member Functions | Static Public Attributes
onert::backend::train::TensorManager Class Reference

#include <TensorManager.h>

Public Member Functions

 TensorManager (const std::shared_ptr< TensorRegistry > &reg, uint32_t optim_vars_count)
 
virtual ~TensorManager ()=default
 
void allocateNonConstTensors ()
 
void allocateTrainableTensors ()
 
void allocateBackPropTensors ()
 
void allocateGradientTensors ()
 
void allocateDisposableBackPropTensors ()
 
void allocateLayerScopeTensors ()
 
void claimNonConstPlan (const ir::OperandIndex &ind)
 
void releaseNonConstPlan (const ir::OperandIndex &ind)
 
void claimTrainablePlan (const ir::OperandIndex &ind)
 
void releaseTrainablePlan (const ir::OperandIndex &ind)
 
void claimBackPropPlan (const ir::OperandIndex &ind)
 
void releaseBackPropPlan (const ir::OperandIndex &ind)
 
void claimGradientPlan (const ir::OperandIndex &ind)
 
void releaseGradientPlan (const ir::OperandIndex &ind)
 
void claimDisposableBackPropPlan (const DisposableTensorIndex &ind)
 
void releaseDisposableBackPropPlan (const DisposableTensorIndex &ind)
 
void claimLayerScopePlan (const LayerScopeTensorIndex &ind)
 
void releaseLayerScopePlan (const LayerScopeTensorIndex &ind)
 

Static Public Attributes

static constexpr uint64_t _align = 16
 

Detailed Description

Definition at line 34 of file TensorManager.h.

Constructor & Destructor Documentation

◆ TensorManager()

onert::backend::train::TensorManager::TensorManager ( const std::shared_ptr< TensorRegistry > &  reg,
uint32_t  optim_vars_count 
)

Definition at line 56 of file TensorManager.cc.

57 : _nonconst_mgr{new MemoryManager()},
58 _trainable_mgr{new TrainableMemoryManager(optim_vars_count)},
59 _back_prop_mgr{new MemoryManager()}, _gradient_mgr{new MemoryManager()},
60 // TODO Find a suitable planner of disposable tensors to reduce peak memory usage
61 _disposable_back_prop_mgr{new DisposableMemoryManager()},
62 _layer_scope_mgr{new LayerScopeMemoryManager()}, _tensors{reg}
63{
64 // DO NOTHING
65}
onert::backend::train::MemoryManager
backend::basic::MemoryManager MemoryManager
Definition MemoryManager.h:32

◆ ~TensorManager()

virtual onert::backend::train::TensorManager::~TensorManager ( )
virtualdefault

Member Function Documentation

◆ allocateBackPropTensors()

void onert::backend::train::TensorManager::allocateBackPropTensors ( )

Definition at line 92 of file TensorManager.cc.

93{
94 allocateMemory(_back_prop_mgr.get(), _tensors->back_prop_tensors(),
95 std::string{" BACK_PROP TENSOR "});
96}

◆ allocateDisposableBackPropTensors()

void onert::backend::train::TensorManager::allocateDisposableBackPropTensors ( )

Definition at line 104 of file TensorManager.cc.

105{
106 allocateMemory(_disposable_back_prop_mgr.get(), _tensors->disposable_back_prop_tensors(),
107 std::string{"DISPOSABLE BACK_PROP TENSOR "});
108}

◆ allocateGradientTensors()

void onert::backend::train::TensorManager::allocateGradientTensors ( )

Definition at line 98 of file TensorManager.cc.

99{
100 allocateMemory(_gradient_mgr.get(), _tensors->gradient_tensors(),
101 std::string{" GRADIENT TENSOR "});
102}

◆ allocateLayerScopeTensors()

void onert::backend::train::TensorManager::allocateLayerScopeTensors ( )

Definition at line 110 of file TensorManager.cc.

111{
112 allocateMemory(_layer_scope_mgr.get(), _tensors->layerscope_tensors(),
113 std::string{" LAYERSCOPE TENSOR "});
114}

◆ allocateNonConstTensors()

void onert::backend::train::TensorManager::allocateNonConstTensors ( )

Definition at line 67 of file TensorManager.cc.

68{
69 allocateMemory(_nonconst_mgr.get(), _tensors->nonconst_tensors(),
70 std::string{" TENSOR "});
71}

◆ allocateTrainableTensors()

void onert::backend::train::TensorManager::allocateTrainableTensors ( )

Definition at line 73 of file TensorManager.cc.

74{
75 allocateMemory(_trainable_mgr.get(), _tensors->trainable_tensors(),
76 std::string{" TRAINABLE TENSOR "});
77
78 // Set buffers of optimizer variables
79 // Allocating optimizer variables has been done when calling allocate() of TrainableMemoryManager
80 for (const auto &[index, trainable_tensor] : _tensors->trainable_tensors())
81 {
82 for (uint32_t var_pos = 0; var_pos < trainable_tensor->optVars().size(); ++var_pos)
83 {
84 auto *buffer = _trainable_mgr.get()->getOptVarBuffer(index, var_pos);
85 trainable_tensor->setOptVarBuffer(buffer, var_pos);
86 VERBOSE(TensorManager) << std::string{" OPTIM_VAR TENSOR "} << index << " : "
87 << static_cast<void *>(buffer) << std::endl;
88 }
89 }
90}
VERBOSE
#define VERBOSE(name, lv)
Definition Log.h:71
moco::index
loco::GraphInputIndex index(const TFPlaceholder *node)
Definition TFNode.cpp:54
size
int32_t size[5]
Definition Slice.cpp:35

References VERBOSE.

◆ claimBackPropPlan()

void onert::backend::train::TensorManager::claimBackPropPlan ( const ir::OperandIndex &  ind)

Definition at line 148 of file TensorManager.cc.

149{
150 auto tensor = _tensors->getBackPropTensor(index);
151 assert(tensor && !tensor->is_dynamic());
152
153 auto size = alignedSize(tensor->total_size(), _align);
154 _back_prop_mgr->claimPlan(index, size);
155}
onert::backend::train::TensorManager::_align
static constexpr uint64_t _align
Definition TensorManager.h:41

References _align, and size.

◆ claimDisposableBackPropPlan()

void onert::backend::train::TensorManager::claimDisposableBackPropPlan ( const DisposableTensorIndex &  ind)

Definition at line 180 of file TensorManager.cc.

181{
182 const auto tensor = _tensors->getDisposableBackPropTensor(index);
183 assert(tensor && !tensor->is_dynamic());
184
185 auto size = alignedSize(tensor->total_size(), _align);
186 _disposable_back_prop_mgr->claimPlan(index, size);
187}

References _align, and size.

◆ claimGradientPlan()

void onert::backend::train::TensorManager::claimGradientPlan ( const ir::OperandIndex &  ind)

Definition at line 164 of file TensorManager.cc.

165{
166 auto tensor = _tensors->getGradientTensor(index);
167 assert(tensor && !tensor->is_dynamic());
168
169 auto size = alignedSize(tensor->total_size(), _align);
170 _gradient_mgr->claimPlan(index, size);
171}

References _align, and size.

◆ claimLayerScopePlan()

void onert::backend::train::TensorManager::claimLayerScopePlan ( const LayerScopeTensorIndex &  ind)

Definition at line 197 of file TensorManager.cc.

198{
199 const auto tensor = _tensors->getLayerScopeTensor(index);
200
201 auto size = alignedSize(tensor->total_size(), _align);
202 _layer_scope_mgr->claimPlan(index, size);
203}

References _align, and size.

◆ claimNonConstPlan()

void onert::backend::train::TensorManager::claimNonConstPlan ( const ir::OperandIndex &  ind)

Definition at line 116 of file TensorManager.cc.

117{
118 auto tensor = _tensors->getNonConstTensor(index);
119 assert(tensor && !tensor->is_dynamic());
120
121 auto size = alignedSize(tensor->total_size(), _align);
122 _nonconst_mgr->claimPlan(index, size);
123}

References _align, and size.

◆ claimTrainablePlan()

void onert::backend::train::TensorManager::claimTrainablePlan ( const ir::OperandIndex &  ind)

Definition at line 132 of file TensorManager.cc.

133{
134 auto tensor = _tensors->getTrainableTensor(index);
135 assert(tensor && !tensor->is_dynamic());
136
137 auto size = alignedSize(tensor->total_size(), _align);
138 _trainable_mgr->claimPlan(index, size);
139}

References _align, and size.

◆ releaseBackPropPlan()

void onert::backend::train::TensorManager::releaseBackPropPlan ( const ir::OperandIndex &  ind)

Definition at line 157 of file TensorManager.cc.

158{
159 assert(_tensors->getBackPropTensor(index) && !_tensors->getBackPropTensor(index)->is_dynamic());
160
161 _back_prop_mgr->releasePlan(index);
162}

◆ releaseDisposableBackPropPlan()

void onert::backend::train::TensorManager::releaseDisposableBackPropPlan ( const DisposableTensorIndex &  ind)

Definition at line 189 of file TensorManager.cc.

190{
191 assert(_tensors->getDisposableBackPropTensor(index) &&
192 !_tensors->getDisposableBackPropTensor(index)->is_dynamic());
193
194 _disposable_back_prop_mgr->releasePlan(index);
195}

◆ releaseGradientPlan()

void onert::backend::train::TensorManager::releaseGradientPlan ( const ir::OperandIndex &  ind)

Definition at line 173 of file TensorManager.cc.

174{
175 assert(_tensors->getGradientTensor(index) && !_tensors->getGradientTensor(index)->is_dynamic());
176
177 _gradient_mgr->releasePlan(index);
178}

◆ releaseLayerScopePlan()

void onert::backend::train::TensorManager::releaseLayerScopePlan ( const LayerScopeTensorIndex &  ind)

Definition at line 205 of file TensorManager.cc.

206{
207 assert(_tensors->getLayerScopeTensor(index));
208 _layer_scope_mgr->releasePlan(index);
209}

◆ releaseNonConstPlan()

void onert::backend::train::TensorManager::releaseNonConstPlan ( const ir::OperandIndex &  ind)

Definition at line 125 of file TensorManager.cc.

126{
127 assert(_tensors->getNonConstTensor(index) && !_tensors->getNonConstTensor(index)->is_dynamic());
128
129 _nonconst_mgr->releasePlan(index);
130}

◆ releaseTrainablePlan()

void onert::backend::train::TensorManager::releaseTrainablePlan ( const ir::OperandIndex &  ind)

Definition at line 141 of file TensorManager.cc.

142{
143 assert(_tensors->getTrainableTensor(index) && !_tensors->getTrainableTensor(index)->is_dynamic());
144
145 _trainable_mgr->releasePlan(index);
146}

Field Documentation

◆ _align

constexpr uint64_t onert::backend::train::TensorManager::_align = 16
staticconstexpr

Definition at line 41 of file TensorManager.h.

Referenced by claimBackPropPlan(), claimDisposableBackPropPlan(), claimGradientPlan(), claimLayerScopePlan(), claimNonConstPlan(), and claimTrainablePlan().

The documentation for this class was generated from the following files: