runtime/onert/backend/train/TensorBuilder.h

   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 #ifndef __ONERT_BACKEND_TRAIN_TENSOR_BUILDER_H__
  18 #define __ONERT_BACKEND_TRAIN_TENSOR_BUILDER_H__
  19
  20 #include "TensorManager.h"
  21 #include "TensorRegistry.h"
  22
  23 namespace onert
  24 {
  25 namespace backend
  26 {
  27 namespace train
  28 {
  29
  30 // TODO Support dynamic tensors
  31 class TensorBuilder
  32 {
  33 public:
  34   TensorBuilder(const std::shared_ptr<TensorRegistry> &tensor_reg, const std::string planner_id);
  35
  36   /**
  37    * @brief     Register tensor information to allocate on train backend
  38    * @param[in] ind    Operand index
  39    * @param[in] info   Operand information
  40    * @param[in] layout Operand data layout
  41    */
  42   void registerTensorInfo(const ir::OperandIndex &ind, const ir::OperandInfo &info,
  43                           ir::Layout backend_layout);
  44
  45   /**
  46    * @brief     Register informations of tensor used only in backward to allocate on train backend
  47    * @param[in] ind    Operand index
  48    * @param[in] info   Operand information
  49    * @param[in] layout Operand data layout
  50    */
  51   void registerBackwardTensorInfo(const ir::OperandIndex &ind, const ir::OperandInfo &info,
  52                                   ir::Layout backend_layout);
  53
  54   // TODO Support memory plan of all tensors
  55   void notifyFirstUse(const ir::OperandIndex &);
  56   void notifyLastUse(const ir::OperandIndex &);
  57   void notifyBackwardFirstUse(const ir::OperandIndex &);
  58
  59   bool isRegistered(const ir::OperandIndex &) const;
  60   bool isRegisteredBackward(const ir::OperandIndex &) const;
  61
  62   void allocate(void);
  63   void allocateBackward(void);
  64
  65 private:
  66   const std::shared_ptr<TensorRegistry> _tensor_reg;
  67   std::unique_ptr<TensorManager> _tensor_mgr;
  68   ir::OperandIndexMap<ir::OperandInfo> _tensor_info_map;
  69   ir::OperandIndexMap<ir::OperandInfo> _backward_tensor_info_map;
  70   ir::OperandIndexMap<bool> _as_constants;
  71 };
  72
  73 } // namespace train
  74 } // namespace backend
  75 } // namespace onert
  76
  77 #endif // __ONERT_BACKEND_TRAIN_TENSOR_BUILDER_H__