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   std::shared_ptr<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   std::shared_ptr<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   std::shared_ptr<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   std::shared_ptr<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   std::shared_ptr<Tensor> getNativeOwnTensor(const ir::OperandIndex &ind)
  84   {
  85     return _base_reg->getNativeTensor(ind);
  86   }
  87
  88   std::shared_ptr<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;
  93     return nullptr;
  94   }
  95
  96   bool setMigrantTensor(const ir::OperandIndex &ind,
  97                         const std::shared_ptr<IPortableTensor> &tensor) override
  98   {
  99     assert(tensor);
 100     assert(!getITensor(ind)); // For the ind, tensor is not registered yet
 101     _base_reg->setMigrantTensor(ind, tensor);
 102     return true;
 103   }
 104
 105   void setNativeOwnTensor(ir::OperandIndex ind, const std::shared_ptr<Tensor> &tensor)
 106   {
 107     assert(tensor);
 108     assert(!getITensor(ind)); // For the ind, tensor is not registered yet
 109     _base_reg->setNativeTensor(ind, tensor);
 110   }
 111
 112   void setNativeUserTensor(ir::OperandIndex ind, const std::shared_ptr<UserTensor> &tensor)
 113   {
 114     assert(tensor);
 115     assert(!getITensor(ind)); // For the ind, tensor is not registered yet
 116     _native_user_tensors[ind] = tensor;
 117   }
 118
 119   const ir::OperandIndexMap<std::shared_ptr<UserTensor>> &native_user_tensors()
 120   {
 121     return _native_user_tensors;
 122   }
 123   std::shared_ptr<cpu_common::TensorRegistry> base_reg() { return _base_reg; }
 124
 125 private:
 126   std::shared_ptr<cpu_common::TensorRegistry> _base_reg;
 127   ir::OperandIndexMap<std::shared_ptr<UserTensor>> _native_user_tensors;
 128 };
 129
 130 } // namespace controlflow
 131 } // namespace backend
 132 } // namespace onert
 133
 134 #endif // ifndef __ONERT_BACKEND_CONTROLFLOW_TENSOR_REGISTRY_H__