From ca3a3e0fd86a07de9e073ba31dc2b42d6ca84536 Mon Sep 17 00:00:00 2001
From: Finn Williams <Finn.Williams@arm.com>
Date: Fri, 11 Jun 2021 15:04:02 +0100
Subject: IVGCVSW-6062 Rework the async threadpool

!armnn:5801

Signed-off-by: Finn Williams <Finn.Williams@arm.com>
Change-Id: I9964d0899ce752441f380edddbd974010257b2dd
---
 ArmnnPreparedModel_1_3.cpp | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

(limited to 'ArmnnPreparedModel_1_3.cpp')

diff --git a/ArmnnPreparedModel_1_3.cpp b/ArmnnPreparedModel_1_3.cpp
index 63918725..cd3df516 100644
--- a/ArmnnPreparedModel_1_3.cpp
+++ b/ArmnnPreparedModel_1_3.cpp
@@ -169,7 +169,8 @@ ArmnnPreparedModel_1_3<HalVersion>::ArmnnPreparedModel_1_3(armnn::NetworkId netw
                                                            const std::string& requestInputsAndOutputsDumpDir,
                                                            const bool gpuProfilingEnabled,
                                                            V1_3::Priority priority,
-                                                           const bool asyncModelExecutionEnabled)
+                                                           const bool asyncModelExecutionEnabled,
+                                                           const unsigned int numberOfThreads)
     : m_NetworkId(networkId)
     , m_Runtime(runtime)
     , m_Model(model)
@@ -184,7 +185,14 @@ ArmnnPreparedModel_1_3<HalVersion>::ArmnnPreparedModel_1_3(armnn::NetworkId netw
 
     if (asyncModelExecutionEnabled)
     {
-        m_WorkingMemHandle = m_Runtime->CreateWorkingMemHandle(networkId);
+        std::vector<std::shared_ptr<armnn::IWorkingMemHandle>> memHandles;
+        for (int i=0; i < numberOfThreads; ++i)
+        {
+            memHandles.emplace_back(m_Runtime->CreateWorkingMemHandle(networkId));
+        }
+
+        m_WorkingMemHandle = memHandles.back();
+        m_Threadpool = std::make_unique<armnn::Threadpool>(numberOfThreads, runtime, memHandles);
     }
 }
 
@@ -821,11 +829,11 @@ void ArmnnPreparedModel_1_3<HalVersion>::ScheduleGraphForExecution(
                                                           outputTensors,
                                                           callbackContext);
 
-    m_Runtime->Schedule(m_NetworkId,
-                        *tpCb->m_InputTensors,
-                        *tpCb->m_OutputTensors,
-                        priority,
-                        tpCb);
+    m_Threadpool->Schedule(m_NetworkId,
+                           *tpCb->m_InputTensors,
+                           *tpCb->m_OutputTensors,
+                           priority,
+                           tpCb);
     ALOGV("ArmnnPreparedModel_1_3::ScheduleGraphForExecution end");
 }
 
-- 
cgit v1.2.1