runtime/onert/backend/train/TensorManager.cc

   1 /*
   2  * Copyright (c) 2023 Samsung Electronics Co., Ltd. All Rights Reserved
   3  *
   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
   7  *
   8  *      http://www.apache.org/licenses/LICENSE-2.0
   9  *
  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.
  15  */
  16
  17 #include "TensorManager.h"
  18
  19 #include <util/logging.h>
  20
  21 namespace
  22 {
  23
  24 using namespace onert;
  25
  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)
  30 {
  31   mgr->allocate();
  32
  33   for (auto &&pair : tensors)
  34   {
  35     const auto &index = pair.first;
  36     auto tensor = pair.second.get();
  37     assert(!tensor->is_dynamic());
  38
  39     auto *buffer = mgr->getBuffer(index);
  40     tensor->setBuffer(buffer);
  41     VERBOSE(TensorManager) << tensor_type << index << " : " << static_cast<void *>(buffer)
  42                            << std::endl;
  43   }
  44 }
  45
  46 } // namespace
  47
  48 namespace onert
  49 {
  50 namespace backend
  51 {
  52 namespace train
  53 {
  54
  55 TensorManager::TensorManager(const std::shared_ptr<TensorRegistry> &reg,
  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}
  60 {
  61   // DO NOTHING
  62 }
  63
  64 void TensorManager::allocateNonConstTensors()
  65 {
  66   allocateMemory(_nonconst_mgr.get(), _tensors->nonconst_tensors(),
  67                  std::string{"          TENSOR "});
  68 }
  69
  70 void TensorManager::allocateTrainableTensors()
  71 {
  72   allocateMemory(_trainable_mgr.get(), _tensors->trainable_tensors(),
  73                  std::string{"TRAINABLE TENSOR "});
  74 }
  75
  76 void TensorManager::allocateDerivativeTensors()
  77 {
  78   allocateMemory(_derivative_mgr.get(), _tensors->derivative_tensors(),
  79                  std::string{"DERIVATIVE TENSOR "});
  80 }
  81
  82 void TensorManager::allocateGradientTensors()
  83 {
  84   allocateMemory(_gradient_mgr.get(), _tensors->gradient_tensors(),
  85                  std::string{"GRADIENT TENSOR "});
  86 }
  87
  88 void TensorManager::claimNonConstPlan(const ir::OperandIndex &index)
  89 {
  90   auto tensor = _tensors->getNonConstTensor(index);
  91   assert(tensor && !tensor->is_dynamic());
  92
  93   auto size = tensor->total_size();
  94   _nonconst_mgr->claimPlan(index, size);
  95 }
  96
  97 void TensorManager::releaseNonConstPlan(const ir::OperandIndex &index)
  98 {
  99   assert(_tensors->getNonConstTensor(index) && !_tensors->getNonConstTensor(index)->is_dynamic());
 100
 101   _nonconst_mgr->releasePlan(index);
 102 }
 103
 104 void TensorManager::claimTrainablePlan(const ir::OperandIndex &index)
 105 {
 106   auto tensor = _tensors->getTrainableTensor(index);
 107   assert(tensor && !tensor->is_dynamic());
 108
 109   auto size = tensor->total_size();
 110   _trainable_mgr->claimPlan(index, size);
 111 }
 112
 113 void TensorManager::releaseTrainablePlan(const ir::OperandIndex &index)
 114 {
 115   assert(_tensors->getTrainableTensor(index) && !_tensors->getTrainableTensor(index)->is_dynamic());
 116
 117   _trainable_mgr->releasePlan(index);
 118 }
 119
 120 void TensorManager::claimDerivativePlan(const ir::OperandIndex &index)
 121 {
 122   auto tensor = _tensors->getDerivativeTensor(index);
 123   assert(tensor && !tensor->is_dynamic());
 124
 125   auto size = tensor->total_size();
 126   _derivative_mgr->claimPlan(index, size);
 127 }
 128
 129 void TensorManager::releaseDerivativePlan(const ir::OperandIndex &index)
 130 {
 131   assert(_tensors->getDerivativeTensor(index) &&
 132          !_tensors->getDerivativeTensor(index)->is_dynamic());
 133
 134   _derivative_mgr->releasePlan(index);
 135 }
 136
 137 void TensorManager::claimGradientPlan(const ir::OperandIndex &index)
 138 {
 139   auto tensor = _tensors->getGradientTensor(index);
 140   assert(tensor && !tensor->is_dynamic());
 141
 142   auto size = tensor->total_size();
 143   _gradient_mgr->claimPlan(index, size);
 144 }
 145
 146 void TensorManager::releaseGradientPlan(const ir::OperandIndex &index)
 147 {
 148   assert(_tensors->getGradientTensor(index) && !_tensors->getGradientTensor(index)->is_dynamic());
 149
 150   _gradient_mgr->releasePlan(index);
 151 }
 152
 153 } // namespace train
 154 } // namespace backend
 155 } // namespace onert