1 files changed, 55 insertions, 11 deletions

diff --git a/src/backends/reference/RefTensorHandle.cpp b/src/backends/reference/RefTensorHandle.cpp
index b7670f6f6e..fe9310f423 100644
--- a/src/backends/reference/RefTensorHandle.cpp
+++ b/src/backends/reference/RefTensorHandle.cpp
@@ -7,39 +7,83 @@
 namespace armnn
 {
 
-RefTensorHandle::RefTensorHandle(const TensorInfo &tensorInfo):
+RefTensorHandle::RefTensorHandle(const TensorInfo &tensorInfo, std::shared_ptr<RefMemoryManager> &memoryManager):
     m_TensorInfo(tensorInfo),
-    m_Memory(nullptr)
+    m_MemoryManager(memoryManager),
+    m_Pool(nullptr),
+    m_UnmanagedMemory(nullptr)
 {
 
 }
 
 RefTensorHandle::~RefTensorHandle()
 {
-    ::operator delete(m_Memory);
+    if (!m_Pool)
+    {
+        // unmanaged
+        ::operator delete(m_UnmanagedMemory);
+    }
+}
+
+void RefTensorHandle::Manage()
+{
+    BOOST_ASSERT_MSG(!m_Pool, "RefTensorHandle::Manage() called twice");
+    BOOST_ASSERT_MSG(!m_UnmanagedMemory, "RefTensorHandle::Manage() called after Allocate()");
+
+    m_Pool = m_MemoryManager->Manage(m_TensorInfo.GetNumBytes());
 }
 
 void RefTensorHandle::Allocate()
 {
-    if (m_Memory == nullptr)
+    if (!m_UnmanagedMemory)
     {
-        m_Memory = ::operator new(m_TensorInfo.GetNumBytes());
+        if (!m_Pool)
+        {
+            // unmanaged
+            m_UnmanagedMemory = ::operator new(m_TensorInfo.GetNumBytes());
+        }
+        else
+        {
+            m_MemoryManager->Allocate(m_Pool);
+        }
     }
     else
     {
         throw InvalidArgumentException("RefTensorHandle::Allocate Trying to allocate a RefTensorHandle"
-                                           "that already has allocated memory.");
+                                       "that already has allocated memory.");
+    }
+}
+
+const void* RefTensorHandle::Map(bool /*unused*/) const
+{
+    return GetPointer();
+}
+
+void* RefTensorHandle::GetPointer() const
+{
+    if (m_UnmanagedMemory)
+    {
+        return m_UnmanagedMemory;
+    }
+    else
+    {
+        BOOST_ASSERT_MSG(m_Pool, "RefTensorHandle::GetPointer called on unmanaged, unallocated tensor handle");
+        return m_MemoryManager->GetPointer(m_Pool);
     }
 }
 
-void RefTensorHandle::CopyOutTo(void* memory) const
+void RefTensorHandle::CopyOutTo(void* dest) const
 {
-    memcpy(memory, m_Memory, m_TensorInfo.GetNumBytes());
+    const void *src = GetPointer();
+    BOOST_ASSERT(src);
+    memcpy(dest, src, m_TensorInfo.GetNumBytes());
 }
 
-void RefTensorHandle::CopyInFrom(const void* memory)
+void RefTensorHandle::CopyInFrom(const void* src)
 {
-    memcpy(m_Memory, memory, m_TensorInfo.GetNumBytes());
+    void *dest = GetPointer();
+    BOOST_ASSERT(dest);
+    memcpy(dest, src, m_TensorInfo.GetNumBytes());
 }
 
-}
\ No newline at end of file
+}