runtime/onert/core/src/backend/controlflow/TensorRegistry.h

   1 /*
   2  * Copyright (c) 2020 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_CONTROLFLOW_TENSOR_REGISTRY_H__
  18 #define __ONERT_BACKEND_CONTROLFLOW_TENSOR_REGISTRY_H__
  19
  20 #include "backend/cpu_common/TensorRegistry.h"
  21 #include "backend/ITensorRegistry.h"
  22 #include "Tensor.h"
  23 #include "UserTensor.h"
  24 #include <assert.h>
  25
  26 namespace onert
  27 {
  28 namespace backend
  29 {
  30 namespace controlflow
  31 {
  32
  33 /**
  34  * @brief Tensor registry class for controlflow backend
  35  *
  36  * This class contains three types of tensors. Two native tensors(tensors that are managed by this
  37  * backend) and the other is migrant tensor.
  38  *
  39  * - NativeUserTensor - @c UserTensor managed by this backend, buffer is user-given
  40  * - NativeOwnTensor  - @c cpu_common::Tensor managed by this backend ( in @c _base_reg )
  41  * - MigrantTensor    - @c IPortableTensor managed by other backends ( in @c _base_reg )
  42  *
  43  * @note @c _base_reg is used in implementation to reuse @c cpu_common::StaticTensorManager
  44  *
  45  */
  46 class TensorRegistry : public ITensorRegistry
  47 {
  48 public:
  49   TensorRegistry() : _base_reg{new cpu_common::TensorRegistry} {}
  50
  51   ITensor *getITensor(const ir::OperandIndex &ind) override
  52   {
  53     auto base_tensor = _base_reg->getITensor(ind);
  54     if (base_tensor)
  55       return base_tensor;
  56     return getNativeUserTensor(ind);
  57   }
  58
  59   ITensor *getNativeITensor(const ir::OperandIndex &ind) override
  60   {
  61     auto base_tensor = _base_reg->getNativeITensor(ind);
  62     if (base_tensor)
  63       return base_tensor;
  64     return getNativeUserTensor(ind);
  65   }
  66
  67   IPortableTensor *getPortableTensor(const ir::OperandIndex &ind)
  68   {
  69     auto base_tensor = _base_reg->getPortableTensor(ind);
  70     if (base_tensor)
  71       return base_tensor;
  72     return getNativeUserTensor(ind);
  73   }
  74
  75   IPortableTensor *getNativeTensor(const ir::OperandIndex &ind)
  76   {
  77     auto base_tensor = _base_reg->getNativeTensor(ind);
  78     if (base_tensor)
  79       return base_tensor;
  80     return getNativeUserTensor(ind);
  81   }
  82
  83   Tensor *getNativeOwnTensor(const ir::OperandIndex &ind)
  84   {
  85     return _base_reg->getNativeTensor(ind);
  86   }
  87
  88   UserTensor *getNativeUserTensor(const ir::OperandIndex &ind)
  89   {
  90     auto tensor = _native_user_tensors.find(ind);
  91     if (tensor != _native_user_tensors.end())
  92       return tensor->second.get();
  93     return nullptr;
  94   }
  95
  96   bool setMigrantTensor(const ir::OperandIndex &ind, IPortableTensor *tensor) override
  97   {
  98     assert(tensor);
  99     assert(!getITensor(ind)); // For the ind, tensor is not registered yet
 100     _base_reg->setMigrantTensor(ind, tensor);
 101     return true;
 102   }
 103
 104   void setNativeOwnTensor(ir::OperandIndex ind, std::unique_ptr<Tensor> &&tensor)
 105   {
 106     assert(tensor);
 107     assert(!getITensor(ind)); // For the ind, tensor is not registered yet
 108     _base_reg->setNativeTensor(ind, std::move(tensor));
 109   }
 110
 111   void setNativeUserTensor(ir::OperandIndex ind, std::unique_ptr<UserTensor> &&tensor)
 112   {
 113     assert(tensor);
 114     assert(!getITensor(ind)); // For the ind, tensor is not registered yet
 115     _native_user_tensors[ind] = std::move(tensor);
 116   }
 117
 118   const ir::OperandIndexMap<std::unique_ptr<UserTensor>> &native_user_tensors()
 119   {
 120     return _native_user_tensors;
 121   }
 122   std::shared_ptr<cpu_common::TensorRegistry> base_reg() { return _base_reg; }
 123
 124 private:
 125   std::shared_ptr<cpu_common::TensorRegistry> _base_reg;
 126   ir::OperandIndexMap<std::unique_ptr<UserTensor>> _native_user_tensors;
 127 };
 128
 129 } // namespace controlflow
 130 } // namespace backend
 131 } // namespace onert
 132
 133 #endif // ifndef __ONERT_BACKEND_CONTROLFLOW_TENSOR_REGISTRY_H__