1 files changed, 34 insertions, 22 deletions
diff --git a/src/armnn/WorkingMemHandle.cpp b/src/armnn/WorkingMemHandle.cpp
index c1a48d482f..0cbef82e83 100644
--- a/src/armnn/WorkingMemHandle.cpp
+++ b/src/armnn/WorkingMemHandle.cpp
@@ -6,6 +6,7 @@
 #include "backendsCommon/CpuTensorHandle.hpp"
 #include "WorkingMemHandle.hpp"
 #include "Network.hpp"
+#include <armnn/backends/IMemoryManager.hpp>
 
 namespace armnn
 {
@@ -13,36 +14,47 @@ namespace armnn
 namespace experimental
 {
 
-WorkingMemHandle::WorkingMemHandle(NetworkId networkId,
-                                   std::vector<WorkingMemDescriptor> workingMemDescriptors,
-                                   std::unordered_map<LayerGuid, WorkingMemDescriptor> workingMemDescriptorMap) :
+WorkingMemHandle::WorkingMemHandle(
+        NetworkId networkId,
+        std::vector<WorkingMemDescriptor> workingMemDescriptors,
+        std::unordered_map<LayerGuid, WorkingMemDescriptor> workingMemDescriptorMap,
+        std::vector<std::shared_ptr<IMemoryManager>> memoryManagers,
+        std::unordered_map<LayerGuid, std::vector<std::unique_ptr<ITensorHandle> > > ownedTensorHandles) :
     m_NetworkId(networkId),
     m_WorkingMemDescriptors(workingMemDescriptors),
     m_WorkingMemDescriptorMap(workingMemDescriptorMap),
+    m_MemoryManagers(memoryManagers),
+    m_OwnedTensorHandles(std::move(ownedTensorHandles)),
     m_IsAllocated(false),
     m_Mutex()
-{}
+{
+}
 
-void WorkingMemHandle::FreeWorkingMemory()
+void WorkingMemHandle::Allocate()
 {
-    for (auto workingMemDescriptor : m_WorkingMemDescriptors)
+    if (m_IsAllocated)
+    {
+        return;
+    }
+    m_IsAllocated = true;
+
+    for (auto& mgr : m_MemoryManagers)
+    {
+        mgr->Acquire();
+    }
+}
+
+void WorkingMemHandle::Free()
+{
+    if (!m_IsAllocated)
+    {
+        return;
+    }
+    m_IsAllocated = false;
+
+    for (auto& mgr : m_MemoryManagers)
     {
-        for (auto input : workingMemDescriptor.m_Inputs)
-        {
-            if (input)
-            {
-                delete input;
-                input = nullptr;
-            }
-        }
-        for (auto output : workingMemDescriptor.m_Outputs)
-        {
-            if (output)
-            {
-                delete output;
-                output = nullptr;
-            }
-        }
+        mgr->Release();
     }
 }