1 // 2 // Copyright © 2017 Arm Ltd. All rights reserved. 3 // SPDX-License-Identifier: MIT 4 // 5 6 #pragma once 7 8 #include <armnn/backends/ITensorHandleFactory.hpp> 9 10 #include <memory> 11 #include <vector> 12 13 namespace armnn 14 { 15 16 //Forward 17 class IMemoryManager; 18 19 /// 20 class TensorHandleFactoryRegistry 21 { 22 public: 23 TensorHandleFactoryRegistry() = default; 24 25 TensorHandleFactoryRegistry(const TensorHandleFactoryRegistry& other) = delete; 26 TensorHandleFactoryRegistry(TensorHandleFactoryRegistry&& other) = delete; 27 28 /// Register a TensorHandleFactory and transfer ownership 29 void RegisterFactory(std::unique_ptr<ITensorHandleFactory> allocator); 30 31 /// Register a memory manager with shared ownership 32 void RegisterMemoryManager(std::shared_ptr<IMemoryManager> memoryManger); 33 34 /// Find a TensorHandleFactory by Id 35 /// Returns nullptr if not found 36 ITensorHandleFactory* GetFactory(ITensorHandleFactory::FactoryId id) const; 37 38 /// Aquire memory required for inference 39 void AquireMemory(); 40 41 /// Release memory required for inference 42 void ReleaseMemory(); 43 44 private: 45 std::vector<std::unique_ptr<ITensorHandleFactory>> m_Factories; 46 std::vector<std::shared_ptr<IMemoryManager>> m_MemoryManagers; 47 }; 48 49 } // namespace armnn 50