1 // 2 // Copyright © 2017,2022 Arm Ltd and Contributors. All rights reserved. 3 // SPDX-License-Identifier: MIT 4 // 5 #pragma once 6 7 #include <armnn/BackendId.hpp> 8 #include <armnn/Optional.hpp> 9 #include <memory> 10 #include <unordered_map> 11 #include <functional> 12 #include <stddef.h> 13 #include <string> 14 15 namespace arm 16 { 17 namespace pipe 18 { 19 20 class IProfilingService; 21 22 } // namespace arm 23 } // namespace pipe 24 25 namespace armnn 26 { 27 28 class IBackendInternal; 29 class ICustomAllocator; 30 class IMemoryOptimizerStrategy; 31 32 using IBackendInternalUniquePtr = std::unique_ptr<IBackendInternal>; 33 using MemoryOptimizerStrategiesMapRef = std::unordered_map<BackendId, std::shared_ptr<IMemoryOptimizerStrategy>>; 34 35 class BackendRegistry 36 { 37 public: 38 using PointerType = IBackendInternalUniquePtr; 39 using FactoryFunction = std::function<PointerType()>; 40 41 void Register(const BackendId& id, FactoryFunction factory); 42 bool IsBackendRegistered(const BackendId& id) const; 43 FactoryFunction GetFactory(const BackendId& id) const; 44 size_t Size() const; 45 BackendIdSet GetBackendIds() const; 46 std::string GetBackendIdsAsString() const; 47 void SetProfilingService(armnn::Optional<arm::pipe::IProfilingService&> profilingService); 48 void RegisterAllocator(const BackendId& id, std::shared_ptr<ICustomAllocator> alloc); 49 std::unordered_map<BackendId, std::shared_ptr<ICustomAllocator>> GetAllocators(); 50 void RegisterMemoryOptimizerStrategy(const BackendId& id, std::shared_ptr<IMemoryOptimizerStrategy> strategy); 51 MemoryOptimizerStrategiesMapRef GetMemoryOptimizerStrategies(); 52 BackendRegistry()53 BackendRegistry() {} ~BackendRegistry()54 virtual ~BackendRegistry() {} 55 56 struct StaticRegistryInitializer 57 { StaticRegistryInitializerarmnn::BackendRegistry::StaticRegistryInitializer58 StaticRegistryInitializer(BackendRegistry& instance, 59 const BackendId& id, 60 FactoryFunction factory) 61 { 62 instance.Register(id, factory); 63 } 64 }; 65 66 void Deregister(const BackendId& id); 67 void DeregisterAllocator(const BackendId &id); 68 void DeregisterMemoryOptimizerStrategy(const BackendId &id); 69 70 protected: 71 using FactoryStorage = std::unordered_map<BackendId, FactoryFunction>; 72 73 /// For testing only 74 static void Swap(BackendRegistry& instance, FactoryStorage& other); 75 76 private: 77 BackendRegistry(const BackendRegistry&) = delete; 78 BackendRegistry& operator=(const BackendRegistry&) = delete; 79 80 FactoryStorage m_Factories; 81 armnn::Optional<arm::pipe::IProfilingService&> m_ProfilingService; 82 std::unordered_map<BackendId, std::shared_ptr<ICustomAllocator>> m_CustomMemoryAllocatorMap; 83 std::unordered_map<BackendId, std::shared_ptr<IMemoryOptimizerStrategy>> m_MemoryOptimizerStrategyMap; 84 }; 85 86 BackendRegistry& BackendRegistryInstance(); 87 88 } // namespace armnn 89