diff options
-rw-r--r-- | arm_compute/core/ITensorPack.h | 5 | ||||
-rw-r--r-- | src/core/ITensorPack.cpp | 5 | ||||
-rw-r--r-- | src/runtime/gpu/cl/utils/ClAuxTensorHandler.h | 17 |
3 files changed, 26 insertions, 1 deletions
diff --git a/arm_compute/core/ITensorPack.h b/arm_compute/core/ITensorPack.h index 2f41d4d51e..17b7241862 100644 --- a/arm_compute/core/ITensorPack.h +++ b/arm_compute/core/ITensorPack.h @@ -95,6 +95,11 @@ public: * @return The pointer to the tensor if exist and is const else nullptr */ const ITensor *get_const_tensor(int id) const; + /** Remove the tensor stored with the given id + * + * @param[in] id ID of tensor to remove + */ + void remove_tensor(int id); /** Pack size accessor * * @return Number of tensors registered to the pack diff --git a/src/core/ITensorPack.cpp b/src/core/ITensorPack.cpp index 9eaeece271..90f9a45039 100644 --- a/src/core/ITensorPack.cpp +++ b/src/core/ITensorPack.cpp @@ -67,6 +67,11 @@ ITensor *ITensorPack::get_tensor(int id) return it != _pack.end() ? it->second.tensor : nullptr; } +void ITensorPack::remove_tensor(int id) +{ + _pack.erase(id); +} + size_t ITensorPack::size() const { return _pack.size(); diff --git a/src/runtime/gpu/cl/utils/ClAuxTensorHandler.h b/src/runtime/gpu/cl/utils/ClAuxTensorHandler.h index ad893acea5..152e3c6c04 100644 --- a/src/runtime/gpu/cl/utils/ClAuxTensorHandler.h +++ b/src/runtime/gpu/cl/utils/ClAuxTensorHandler.h @@ -50,6 +50,8 @@ public: if(pack_inject) { pack.add_tensor(slot_id, &_tensor); + _injected_tensor_pack = &pack; + _injected_slot_id = slot_id; } } else @@ -68,6 +70,17 @@ public: } } + CLAuxTensorHandler(const CLAuxTensorHandler &) = delete; + CLAuxTensorHandler &operator=(const CLAuxTensorHandler) = delete; + + ~CLAuxTensorHandler() + { + if(_injected_tensor_pack) + { + _injected_tensor_pack->remove_tensor(_injected_slot_id); + } + } + ICLTensor *get() { return &_tensor; @@ -79,7 +92,9 @@ public: } private: - CLTensor _tensor{}; + CLTensor _tensor{}; + ITensorPack *_injected_tensor_pack{ nullptr }; + int _injected_slot_id{ TensorType::ACL_UNKNOWN }; }; } // namespace opencl } // namespace arm_compute |