1 files changed, 5 insertions, 0 deletions
diff --git a/include/armnn/BackendRegistry.hpp b/include/armnn/BackendRegistry.hpp
index fe6451cde0..c13aa9f8b6 100644
--- a/include/armnn/BackendRegistry.hpp
+++ b/include/armnn/BackendRegistry.hpp
@@ -7,6 +7,7 @@
 #include <armnn/Types.hpp>
 #include <armnn/BackendId.hpp>
 #include <armnn/Optional.hpp>
+#include <armnn/backends/ICustomAllocator.hpp>
 
 #include <memory>
 #include <unordered_map>
@@ -35,6 +36,8 @@ public:
     BackendIdSet GetBackendIds() const;
     std::string GetBackendIdsAsString() const;
     void SetProfilingService(armnn::Optional<profiling::ProfilingService&> profilingService);
+    void RegisterAllocator(const BackendId& id, std::shared_ptr<ICustomAllocator> alloc);
+    std::unordered_map<BackendId, std::shared_ptr<ICustomAllocator>> GetAllocators();
 
     BackendRegistry() {}
     virtual ~BackendRegistry() {}
@@ -50,6 +53,7 @@ public:
     };
 
     void Deregister(const BackendId& id);
+    void DeregisterAllocator(const BackendId &id);
 
 protected:
     using FactoryStorage = std::unordered_map<BackendId, FactoryFunction>;
@@ -63,6 +67,7 @@ private:
 
     FactoryStorage m_Factories;
     armnn::Optional<profiling::ProfilingService&> m_ProfilingService;
+    std::unordered_map<BackendId, std::shared_ptr<ICustomAllocator>> m_CustomMemoryAllocatorMap;
 };
 
 BackendRegistry& BackendRegistryInstance();