From b1aad4270fa8ad5c4aa62e27d564baf723b2cee5 Mon Sep 17 00:00:00 2001
From: Finn Williams <finn.williams@arm.com>
Date: Thu, 28 Oct 2021 19:07:32 +0100
Subject: IVGCVSW-6527 Support the new memory API in loaded network

 * enable external memory management for neon and ref backends
 * change m_TensorMemoryVector to hold shared pointers
 * change input layer backend Id to match backend id of connected layer

Signed-off-by: Finn Williams <finn.williams@arm.com>
Change-Id: I2216a724028312eb101b290df3f224177826b1a0
---
 src/armnn/WorkingMemHandle.hpp | 25 +++++++++++++++++--------
 1 file changed, 17 insertions(+), 8 deletions(-)

(limited to 'src/armnn/WorkingMemHandle.hpp')

diff --git a/src/armnn/WorkingMemHandle.hpp b/src/armnn/WorkingMemHandle.hpp
index 9078a8d54c..bca1d2d80c 100644
--- a/src/armnn/WorkingMemHandle.hpp
+++ b/src/armnn/WorkingMemHandle.hpp
@@ -14,6 +14,7 @@
 
 #include <unordered_map>
 #include <mutex>
+#include <backendsCommon/MemoryManager.hpp>
 
 namespace armnn
 {
@@ -45,11 +46,13 @@ public:
 
     WorkingMemHandle(NetworkId networkId,
                      std::vector<InputMemDescriptorCoords> inputLayerInfo,
-                     std::vector<OutputMemDescriptorCoords> ouputLayerInfo,
+                     std::vector<OutputMemDescriptorCoords> outputLayerInfo,
                      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);
+                     std::unique_ptr<MemoryManager> memoryManager,
+                     std::vector<std::pair<std::shared_ptr<TensorMemory>, MemorySource>> tensorMemory,
+                     std::vector<std::unique_ptr<ITensorHandle>> managedTensorHandles,
+                     std::vector<std::unique_ptr<ITensorHandle>> unmanagedTensorHandles);
 
     ~WorkingMemHandle()
     { Free(); }
@@ -128,11 +131,17 @@ private:
     std::vector<WorkingMemDescriptor> m_WorkingMemDescriptors;
     std::unordered_map<LayerGuid, WorkingMemDescriptor> m_WorkingMemDescriptorMap;
 
-    // Vector of IMemoryManagers that manage the WorkingMemHandle's memory
-    std::vector<std::shared_ptr<IMemoryManager>> m_MemoryManagers;
-    // TensorHandles owned by this WorkingMemHandle
-    // constant tensor's can be shared by multiple WorkingMemHandles and so will not be stored here
-    std::unordered_map<LayerGuid, std::vector<std::unique_ptr<ITensorHandle> > >  m_OwnedTensorHandles;
+    std::unique_ptr<MemoryManager> m_MemoryManager;
+
+    // Memory to be imported into the tensorHandles after allocation
+    std::vector<std::pair<std::shared_ptr<TensorMemory>, MemorySource>> m_TensorMemory;
+
+
+    // Tensors that will need to be allocated internally within armnn
+    std::vector<std::unique_ptr<ITensorHandle>> m_ManagedTensorHandles;
+
+    // Tensors that will be allocated externally by the user
+    std::vector<std::unique_ptr<ITensorHandle>> m_UnmanagedTensorHandles;
 
     std::unordered_map<LayerBindingId, bool> m_InputValidationMap;
     std::unordered_map<LayerBindingId, bool> m_OutputValidationMap;
-- 
cgit v1.2.1