5 files changed, 30 insertions, 9 deletions
diff --git a/src/armnn/layers/ConcatLayer.cpp b/src/armnn/layers/ConcatLayer.cpp
index 238fdb66d9..3a20e1b3f6 100644
--- a/src/armnn/layers/ConcatLayer.cpp
+++ b/src/armnn/layers/ConcatLayer.cpp
@@ -179,7 +179,8 @@ void ConcatLayer::CreateTensors(const TensorHandleFactoryRegistry& registry,
 
 void ConcatLayer::CreateTensorHandles(const TensorHandleFactoryRegistry& registry,
                                       const IWorkloadFactory& workloadFactory,
-                                      const bool isMemoryManaged)
+                                      const bool isMemoryManaged,
+                                      MemorySource memSource)
 {
     OutputSlot& slot = GetOutputSlot(0);
     ITensorHandleFactory::FactoryId factoryId = slot.GetTensorHandleFactoryId();
@@ -190,7 +191,15 @@ void ConcatLayer::CreateTensorHandles(const TensorHandleFactoryRegistry& registr
     }
     else
     {
-        ITensorHandleFactory* handleFactory = registry.GetFactory(factoryId);
+        ITensorHandleFactory* handleFactory;
+        if (memSource == MemorySource::Undefined)
+        {
+            handleFactory = registry.GetFactory(factoryId);
+        }
+        else
+        {
+            handleFactory = registry.GetFactory(factoryId, memSource);
+        }
         ARMNN_ASSERT(handleFactory);
         CreateTensors(registry, *handleFactory, isMemoryManaged);
     }
diff --git a/src/armnn/layers/ConcatLayer.hpp b/src/armnn/layers/ConcatLayer.hpp
index 3d9ba1815e..6a43318382 100644
--- a/src/armnn/layers/ConcatLayer.hpp
+++ b/src/armnn/layers/ConcatLayer.hpp
@@ -24,9 +24,11 @@ public:
     /// @param [in] registry Contains all the registered tensor handle factories available for use.
     /// @param [in] factory The workload factory which will create the workload.
     /// @param [in] IsMemoryManaged Determine whether or not to assign a memory manager during creation
+    /// @param [in] MemorySource Determine the source of memory e.g Malloc
     virtual void CreateTensorHandles(const TensorHandleFactoryRegistry& registry,
                                      const IWorkloadFactory& factory,
-                                     const bool IsMemoryManaged = true) override;
+                                     const bool IsMemoryManaged = true,
+                                     MemorySource memSource = MemorySource::Undefined) override;
 
     /// Creates a dynamically-allocated copy of this layer.
     /// @param [in] graph The graph into which this layer is being cloned.
diff --git a/src/armnn/layers/OutputLayer.hpp b/src/armnn/layers/OutputLayer.hpp
index 6315e25da1..fc6a8aa6b2 100644
--- a/src/armnn/layers/OutputLayer.hpp
+++ b/src/armnn/layers/OutputLayer.hpp
@@ -26,9 +26,10 @@ public:
     /// @param [in] IsMemoryManaged Determine whether or not to assign a memory manager during creation
     virtual void CreateTensorHandles(const TensorHandleFactoryRegistry& registry,
                                      const IWorkloadFactory& factory,
-                                     const bool IsMemoryManaged = true) override
+                                     const bool isMemoryManaged = true,
+                                     MemorySource memSource = MemorySource::Undefined) override
     {
-        IgnoreUnused(registry, factory, IsMemoryManaged);
+        IgnoreUnused(registry, factory, isMemoryManaged, memSource);
     }
 
     /// Creates a dynamically-allocated copy of this layer.
diff --git a/src/armnn/layers/SplitterLayer.cpp b/src/armnn/layers/SplitterLayer.cpp
index 5e6622e13a..adef9aa1a2 100644
--- a/src/armnn/layers/SplitterLayer.cpp
+++ b/src/armnn/layers/SplitterLayer.cpp
@@ -177,7 +177,8 @@ void SplitterLayer::CreateTensors(const TensorHandleFactoryRegistry& registry,
 
 void SplitterLayer::CreateTensorHandles(const TensorHandleFactoryRegistry& registry,
                                         const IWorkloadFactory& workloadFactory,
-                                        const bool isMemoryManaged)
+                                        const bool isMemoryManaged,
+                                        MemorySource memSource)
 {
     OutputSlot& slot = GetOutputSlot(0);
     ITensorHandleFactory::FactoryId factoryId = slot.GetTensorHandleFactoryId();
@@ -188,7 +189,15 @@ void SplitterLayer::CreateTensorHandles(const TensorHandleFactoryRegistry& regis
     }
     else
     {
-        ITensorHandleFactory* handleFactory = registry.GetFactory(factoryId);
+        ITensorHandleFactory* handleFactory;
+        if (memSource == MemorySource::Undefined)
+        {
+            handleFactory = registry.GetFactory(factoryId);
+        }
+        else
+        {
+            handleFactory = registry.GetFactory(factoryId, memSource);
+        }
         ARMNN_ASSERT(handleFactory);
         CreateTensors(registry, *handleFactory, isMemoryManaged);
     }
diff --git a/src/armnn/layers/SplitterLayer.hpp b/src/armnn/layers/SplitterLayer.hpp
index 9999009175..075b136da9 100644
--- a/src/armnn/layers/SplitterLayer.hpp
+++ b/src/armnn/layers/SplitterLayer.hpp
@@ -24,10 +24,10 @@ public:
     /// @param [in] registry Contains all the registered tensor handle factories available for use.
     /// @param [in] factory The workload factory which will create the workload.
     /// @param [in] IsMemoryManaged Determine whether or not to assign a memory manager during creation
-    //virtual void CreateTensorHandles(Graph& graph, const IWorkloadFactory& factory) override;
     virtual void CreateTensorHandles(const TensorHandleFactoryRegistry& registry,
                                      const IWorkloadFactory& factory,
-                                     const bool IsMemoryManaged = true) override;
+                                     const bool IsMemoryManaged = true,
+                                     MemorySource memSource = MemorySource::Undefined) override;
 
     /// Creates a dynamically-allocated copy of this layer.
     /// @param [in] graph The graph into which this layer is being cloned.