From c1c872f12797ef6fe52c4589113e7efc353e56eb Mon Sep 17 00:00:00 2001
From: Jan Eilers <jan.eilers@arm.com>
Date: Thu, 22 Jul 2021 13:17:04 +0100
Subject: Adds CustomAllocator interface and Sample App

 * Updates the runtime options with a CustomAllocatorMap which allows to define a CustomAllocator for specific backends
 * Change IBackendInternal interface to use a shared pointer to a custom allocator
 * Update ClBackend.hpp/cpp to use the CustomAllocator
 * Adds an example application and unit test which uses a CustomAllocator for GpuAcc
 * Refactor of the interface to use MemorySource instead of the user Mapping cl_mem directly
 * Modify the BackendRegistry to also hold a registry of CustomAllocators
 * BackendRegistry Deregister will also deregister any allocators associated with that backend id
 * set_global_allocator within the BaseMemoryManager so that it always matches the currently used allocator

Signed-off-by: Jan Eilers <jan.eilers@arm.com>
Change-Id: I156d819686021865f4375e6cb7a5c5dec8fee9e8
Signed-off-by: David Monahan <david.monahan@arm.com>
---
 src/armnn/BackendRegistry.cpp | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

(limited to 'src/armnn/BackendRegistry.cpp')

diff --git a/src/armnn/BackendRegistry.cpp b/src/armnn/BackendRegistry.cpp
index ff63c8236a..80daed9896 100644
--- a/src/armnn/BackendRegistry.cpp
+++ b/src/armnn/BackendRegistry.cpp
@@ -39,6 +39,7 @@ void BackendRegistry::Register(const BackendId& id, BackendRegistry::FactoryFunc
 void BackendRegistry::Deregister(const BackendId& id)
 {
     m_Factories.erase(id);
+    DeregisterAllocator(id);
 
     if (m_ProfilingService.has_value() && m_ProfilingService.value().IsProfilingEnabled())
     {
@@ -106,5 +107,25 @@ void BackendRegistry::SetProfilingService(armnn::Optional<profiling::ProfilingSe
     m_ProfilingService = profilingService;
 }
 
+void BackendRegistry::RegisterAllocator(const BackendId& id, std::shared_ptr<ICustomAllocator> alloc)
+{
+    if (m_CustomMemoryAllocatorMap.find(id) != m_CustomMemoryAllocatorMap.end())
+    {
+        throw InvalidArgumentException(
+            std::string(id) + " already has an allocator associated with it",
+            CHECK_LOCATION());
+    }
+    m_CustomMemoryAllocatorMap[id] = alloc;
+}
+
+void BackendRegistry::DeregisterAllocator(const BackendId& id)
+{
+    m_CustomMemoryAllocatorMap.erase(id);
+}
+
+std::unordered_map<BackendId, std::shared_ptr<ICustomAllocator>> BackendRegistry::GetAllocators()
+{
+    return m_CustomMemoryAllocatorMap;
+}
 
 } // namespace armnn
-- 
cgit v1.2.1