runtime/onert/core/src/compiler/train/TensorRegistries.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_COMPILER_TRAIN_TENSOR_REGISTRIES_H__
  18 #define __ONERT_COMPILER_TRAIN_TENSOR_REGISTRIES_H__
  19
  20 #include "../../backend/builtin/Config.h"
  21 #include "../../backend/builtin/train/TensorRegistry.h"
  22
  23 #include <backend/train/TrainableBackendContext.h>
  24
  25 #include <memory>
  26 #include <unordered_set>
  27
  28 namespace onert
  29 {
  30 namespace compiler
  31 {
  32 namespace train
  33 {
  34
  35 class TensorRegistries
  36 {
  37 public:
  38   TensorRegistries() = default;
  39
  40   TensorRegistries(const backend::train::TrainableBackendContexts &backend_contexts,
  41                    bool include_builtin)
  42   {
  43     for (const auto &e : backend_contexts)
  44     {
  45       auto tensor_reg = e.second->tensor_registry();
  46       if (e.first->config()->id() == backend::builtin::Config::ID)
  47       {
  48         _builtin_tensor_reg =
  49           std::dynamic_pointer_cast<backend::builtin::train::TensorRegistry>(tensor_reg);
  50         if (include_builtin)
  51           _tensor_regs.insert(tensor_reg);
  52       }
  53       else
  54       {
  55         _tensor_regs.insert(tensor_reg);
  56       }
  57     }
  58   }
  59
  60   std::unordered_set<std::shared_ptr<backend::train::ITensorRegistry>>::const_iterator begin() const
  61   {
  62     return _tensor_regs.cbegin();
  63   }
  64   std::unordered_set<std::shared_ptr<backend::train::ITensorRegistry>>::const_iterator end() const
  65   {
  66     return _tensor_regs.cend();
  67   }
  68
  69   std::shared_ptr<backend::builtin::train::TensorRegistry> getBuiltinTensorRegistry() const
  70   {
  71     return _builtin_tensor_reg;
  72   }
  73
  74   backend::ITensor *getITensor(ir::OperandIndex index) const
  75   {
  76     for (auto &&tensor_reg : _tensor_regs)
  77     {
  78       auto tensor = tensor_reg->getITensor(index);
  79       if (tensor)
  80         return tensor;
  81     }
  82     return nullptr;
  83   }
  84
  85   backend::ITensor *getDerivativeITensor(ir::OperandIndex index) const
  86   {
  87     for (auto &&tensor_reg : _tensor_regs)
  88     {
  89       auto tensor = tensor_reg->getDerivativeITensor(index);
  90       if (tensor)
  91         return tensor;
  92     }
  93     return nullptr;
  94   }
  95
  96 private:
  97   std::unordered_set<std::shared_ptr<backend::train::ITensorRegistry>> _tensor_regs;
  98   std::shared_ptr<backend::builtin::train::TensorRegistry> _builtin_tensor_reg;
  99 };
 100
 101 } // namespace train
 102 } // namespace compiler
 103 } // namespace onert
 104
 105 #endif // __ONERT_COMPILER_TRAIN_TENSOR_REGISTRIES_H__