• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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