2 * Copyright (c) 2023 Samsung Electronics Co., Ltd. All Rights Reserved
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
8 * http://www.apache.org/licenses/LICENSE-2.0
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
17 #include "TensorManager.h"
19 #include <util/logging.h>
24 using namespace onert;
26 template <typename Tensor>
27 void allocateMemory(backend::train::MemoryManager *mgr,
28 const ir::OperandIndexMap<std::unique_ptr<Tensor>> &tensors,
29 const std::string tensor_type)
33 for (auto &&pair : tensors)
35 const auto &index = pair.first;
36 auto tensor = pair.second.get();
37 assert(!tensor->is_dynamic());
39 auto *buffer = mgr->getBuffer(index);
40 tensor->setBuffer(buffer);
41 VERBOSE(TensorManager) << tensor_type << index << " : " << static_cast<void *>(buffer)
55 TensorManager::TensorManager(const std::shared_ptr<TensorRegistry> ®,
56 const std::string planner_id)
57 : _nonconst_mgr{new MemoryManager(planner_id)}, _trainable_mgr{new MemoryManager(planner_id)},
58 _derivative_mgr{new MemoryManager(planner_id)},
59 _gradient_mgr{new MemoryManager(planner_id)}, _tensors{reg}
64 void TensorManager::allocateNonConstTensors()
66 allocateMemory(_nonconst_mgr.get(), _tensors->nonconst_tensors(),
67 std::string{" TENSOR "});
70 void TensorManager::allocateTrainableTensors()
72 allocateMemory(_trainable_mgr.get(), _tensors->trainable_tensors(),
73 std::string{"TRAINABLE TENSOR "});
76 void TensorManager::allocateDerivativeTensors()
78 allocateMemory(_derivative_mgr.get(), _tensors->derivative_tensors(),
79 std::string{"DERIVATIVE TENSOR "});
82 void TensorManager::allocateGradientTensors()
84 allocateMemory(_gradient_mgr.get(), _tensors->gradient_tensors(),
85 std::string{"GRADIENT TENSOR "});
88 void TensorManager::claimNonConstPlan(const ir::OperandIndex &index)
90 auto tensor = _tensors->getNonConstTensor(index);
91 assert(tensor && !tensor->is_dynamic());
93 auto size = tensor->total_size();
94 _nonconst_mgr->claimPlan(index, size);
97 void TensorManager::releaseNonConstPlan(const ir::OperandIndex &index)
99 assert(_tensors->getNonConstTensor(index) && !_tensors->getNonConstTensor(index)->is_dynamic());
101 _nonconst_mgr->releasePlan(index);
104 void TensorManager::claimTrainablePlan(const ir::OperandIndex &index)
106 auto tensor = _tensors->getTrainableTensor(index);
107 assert(tensor && !tensor->is_dynamic());
109 auto size = tensor->total_size();
110 _trainable_mgr->claimPlan(index, size);
113 void TensorManager::releaseTrainablePlan(const ir::OperandIndex &index)
115 assert(_tensors->getTrainableTensor(index) && !_tensors->getTrainableTensor(index)->is_dynamic());
117 _trainable_mgr->releasePlan(index);
120 void TensorManager::claimDerivativePlan(const ir::OperandIndex &index)
122 auto tensor = _tensors->getDerivativeTensor(index);
123 assert(tensor && !tensor->is_dynamic());
125 auto size = tensor->total_size();
126 _derivative_mgr->claimPlan(index, size);
129 void TensorManager::releaseDerivativePlan(const ir::OperandIndex &index)
131 assert(_tensors->getDerivativeTensor(index) &&
132 !_tensors->getDerivativeTensor(index)->is_dynamic());
134 _derivative_mgr->releasePlan(index);
137 void TensorManager::claimGradientPlan(const ir::OperandIndex &index)
139 auto tensor = _tensors->getGradientTensor(index);
140 assert(tensor && !tensor->is_dynamic());
142 auto size = tensor->total_size();
143 _gradient_mgr->claimPlan(index, size);
146 void TensorManager::releaseGradientPlan(const ir::OperandIndex &index)
148 assert(_tensors->getGradientTensor(index) && !_tensors->getGradientTensor(index)->is_dynamic());
150 _gradient_mgr->releasePlan(index);
154 } // namespace backend