From 549cb7ad5fe0686a5216557ef7ee21b520bf7936 Mon Sep 17 00:00:00 2001
From: Narumol Prangnawarat <narumol.prangnawarat@arm.com>
Date: Fri, 10 Jul 2020 17:50:53 +0100
Subject: IVGCVSW-5078 Fix segmentation fault on Yolo V3 Big App

 * Register TensorHandleFactories when create Workload Factory

Signed-off-by: Narumol Prangnawarat <narumol.prangnawarat@arm.com>
Change-Id: I6217c6f00d27b67e42b8043cbaf72d9db5789bfb
---
 src/armnn/LoadedNetwork.cpp           | 3 +--
 src/backends/cl/ClBackend.cpp         | 1 +
 src/backends/neon/NeonBackend.cpp     | 2 ++
 src/backends/reference/RefBackend.cpp | 1 +
 4 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/armnn/LoadedNetwork.cpp b/src/armnn/LoadedNetwork.cpp
index b35dfd1107..0dd75aa649 100644
--- a/src/armnn/LoadedNetwork.cpp
+++ b/src/armnn/LoadedNetwork.cpp
@@ -121,6 +121,7 @@ LoadedNetwork::LoadedNetwork(std::unique_ptr<OptimizedNetwork> net,
                              m_OptimizedNetwork(std::move(net)),
                              m_IsImportEnabled(networkProperties.m_ImportEnabled),
                              m_IsExportEnabled(networkProperties.m_ExportEnabled),
+                             m_TensorHandleFactoryRegistry(),
                              m_ProfilingService(profilingService)
 {
     // Create a profiler and register it for the current thread.
@@ -144,8 +145,6 @@ LoadedNetwork::LoadedNetwork(std::unique_ptr<OptimizedNetwork> net,
 
             if (backend->SupportsTensorAllocatorAPI())
             {
-                backend->RegisterTensorHandleFactories(m_TensorHandleFactoryRegistry);
-
                 auto workloadFactory = backend->CreateWorkloadFactory(m_TensorHandleFactoryRegistry);
                 m_WorkloadFactories.emplace(
                     std::make_pair(backendId, std::make_pair(std::move(workloadFactory), nullptr)));
diff --git a/src/backends/cl/ClBackend.cpp b/src/backends/cl/ClBackend.cpp
index 0a898ec2c0..f9a8993baf 100644
--- a/src/backends/cl/ClBackend.cpp
+++ b/src/backends/cl/ClBackend.cpp
@@ -50,6 +50,7 @@ IBackendInternal::IWorkloadFactoryPtr ClBackend::CreateWorkloadFactory(
     auto memoryManager = std::make_shared<ClMemoryManager>(std::make_unique<arm_compute::CLBufferAllocator>());
 
     registry.RegisterMemoryManager(memoryManager);
+    registry.RegisterFactory(std::make_unique<ClTensorHandleFactory>(memoryManager));
 
     return std::make_unique<ClWorkloadFactory>(
             PolymorphicPointerDowncast<ClMemoryManager>(memoryManager));
diff --git a/src/backends/neon/NeonBackend.cpp b/src/backends/neon/NeonBackend.cpp
index 841ed27006..01cc6d8119 100644
--- a/src/backends/neon/NeonBackend.cpp
+++ b/src/backends/neon/NeonBackend.cpp
@@ -53,6 +53,8 @@ IBackendInternal::IWorkloadFactoryPtr NeonBackend::CreateWorkloadFactory(
                                                              BaseMemoryManager::MemoryAffinity::Offset);
 
     tensorHandleFactoryRegistry.RegisterMemoryManager(memoryManager);
+    tensorHandleFactoryRegistry.RegisterFactory(std::make_unique<NeonTensorHandleFactory>(memoryManager));
+
     return std::make_unique<NeonWorkloadFactory>(
         PolymorphicPointerDowncast<NeonMemoryManager>(memoryManager));
 }
diff --git a/src/backends/reference/RefBackend.cpp b/src/backends/reference/RefBackend.cpp
index 590fde3e16..e93b317dce 100644
--- a/src/backends/reference/RefBackend.cpp
+++ b/src/backends/reference/RefBackend.cpp
@@ -37,6 +37,7 @@ IBackendInternal::IWorkloadFactoryPtr RefBackend::CreateWorkloadFactory(
     auto memoryManager = std::make_shared<RefMemoryManager>();
 
     tensorHandleFactoryRegistry.RegisterMemoryManager(memoryManager);
+    tensorHandleFactoryRegistry.RegisterFactory(std::make_unique<RefTensorHandleFactory>(memoryManager));
 
     return std::make_unique<RefWorkloadFactory>(PolymorphicPointerDowncast<RefMemoryManager>(memoryManager));
 }
-- 
cgit v1.2.1