2 * Copyright (c) 2023 Samsung Electronics Co., Ltd. All Rights Reserved
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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.
17 #ifndef __ONERT_BACKEND_BUILTIN_TRAIN_TENSOR_REGISTRY_H__
18 #define __ONERT_BACKEND_BUILTIN_TRAIN_TENSOR_REGISTRY_H__
20 #include <backend/train/ITensorRegistry.h>
22 #include "../IOTensor.h"
23 #include "../Tensor.h"
35 using BaseTensorRegistry =
36 backend::train::PortableTensorRegistryTemplate<Tensor, TrainableTensor, DerivativeTensor,
39 class TensorRegistry : public backend::train::ITensorRegistry
42 TensorRegistry() : _base_reg{new BaseTensorRegistry} {}
44 ITensor *getITensor(const ir::OperandIndex &index) override
46 auto base_tensor = _base_reg->getITensor(index);
49 return getNativeIOTensor(index);
52 ITensor *getNativeITensor(const ir::OperandIndex &index) override
54 auto base_tensor = _base_reg->getNativeITensor(index);
57 return getNativeIOTensor(index);
60 IPortableTensor *getPortableTensor(const ir::OperandIndex &index)
62 auto base_tensor = _base_reg->getPortableTensor(index);
65 return getNativeIOTensor(index);
68 IOTensor *getNativeIOTensor(const ir::OperandIndex &index)
70 auto tensor = _native_io_tensors.find(index);
71 if (tensor != _native_io_tensors.end())
72 return tensor->second.get();
76 ITensor *getDerivativeITensor(const ir::OperandIndex &index) override
78 return _base_reg->getDerivativeTensor(index);
81 ITensor *getGradientITensor(const ir::OperandIndex &index) override
83 return _base_reg->getGradientTensor(index);
86 DerivativeTensor *getDerivativeTensor(const ir::OperandIndex &index)
88 return _base_reg->getDerivativeTensor(index);
91 bool setMigrantTensor(const ir::OperandIndex &index, IPortableTensor *tensor) override
94 assert(!getITensor(index)); // For the index, tensor is not registered yet
95 _base_reg->setMigrantTensor(index, tensor);
99 void setDerivativeTensor(const ir::OperandIndex &index, std::unique_ptr<DerivativeTensor> tensor)
101 _base_reg->setDerivativeTensor(index, std::move(tensor));
104 void setGradientTensor(const ir::OperandIndex &index, std::unique_ptr<GradientTensor> tensor)
106 _base_reg->setGradientTensor(index, std::move(tensor));
109 void setNativeIOTensor(ir::OperandIndex index, std::unique_ptr<IOTensor> &&tensor)
112 assert(!getITensor(index)); // For the index, tensor is not registered yet
113 _native_io_tensors[index] = std::move(tensor);
116 const ir::OperandIndexMap<std::unique_ptr<IOTensor>> &native_io_tensors()
118 return _native_io_tensors;
120 std::shared_ptr<BaseTensorRegistry> base_reg() { return _base_reg; }
123 std::shared_ptr<BaseTensorRegistry> _base_reg;
124 ir::OperandIndexMap<std::unique_ptr<IOTensor>> _native_io_tensors;
128 } // namespace builtin
129 } // namespace backend
132 #endif // __ONERT_BACKEND_BUILTIN_TRAIN_TENSOR_REGISTRY_H__