From fd313fef775ed210f8dab84452ea382a0b4164b0 Mon Sep 17 00:00:00 2001
From: David Monahan <David.Monahan@arm.com>
Date: Fri, 28 Jan 2022 13:06:35 +0000
Subject: IVGCVSW-6679 Add input and output workload slot pairs to
 LoadedNetwork

 * Added vectors to store the indexes of workload queues which contain inputs or outputs and their corresponding input / output slots

Signed-off-by: David Monahan <David.Monahan@arm.com>
Change-Id: Ide626726718216c7a778cce583da75af3ca2dc9a
---
 src/armnn/LoadedNetwork.cpp | 33 +++++++++++++++++++++++++++++++++
 src/armnn/LoadedNetwork.hpp |  5 +++++
 2 files changed, 38 insertions(+)

diff --git a/src/armnn/LoadedNetwork.cpp b/src/armnn/LoadedNetwork.cpp
index 965deeedc2..1d1aae53a5 100644
--- a/src/armnn/LoadedNetwork.cpp
+++ b/src/armnn/LoadedNetwork.cpp
@@ -312,6 +312,39 @@ LoadedNetwork::LoadedNetwork(std::unique_ptr<IOptimizedNetwork> net,
                     }
                     else
                     {
+                        if (layer->GetNumInputSlots() >= 1)
+                        {
+                            unsigned int slotIndex = 0;
+                            for (auto& inputSlot : layer->GetInputSlots())
+                            {
+                                if (inputSlot.GetOwningLayer().GetType() == LayerType::Input)
+                                {
+                                    m_InputWorkloadSlotPairs.push_back(
+                                        std::make_pair(m_WorkloadQueue.size(), slotIndex));
+                                }
+                                ++slotIndex;
+                            }
+                        }
+
+                        if (layer->GetNumOutputSlots() >= 1)
+                        {
+                            unsigned int slotIndex = 0;
+                            for (auto& outputSlot : layer->GetOutputSlots())
+                            {
+                                for (unsigned int i = 0; i < outputSlot.GetNumConnections(); i++)
+                                {
+                                    // If any of the connections on this outputSlot are connected to an Output then
+                                    // Add its index within layer->GetOutputSlots() to m_OutputWorkloadSlotPairs
+                                    if (outputSlot.GetConnection(i)->GetOwningLayer().GetType() == LayerType::Output)
+                                    {
+                                        m_OutputWorkloadSlotPairs.push_back(
+                                            std::make_pair(m_WorkloadQueue.size(), slotIndex));
+                                        continue;
+                                    }
+                                }
+                                ++slotIndex;
+                            }
+                        }
                         m_WorkloadQueue.push_back(std::move(workload));
                     }
 
diff --git a/src/armnn/LoadedNetwork.hpp b/src/armnn/LoadedNetwork.hpp
index 53c35d61f6..9de6307938 100644
--- a/src/armnn/LoadedNetwork.hpp
+++ b/src/armnn/LoadedNetwork.hpp
@@ -197,6 +197,11 @@ private:
     std::unique_ptr<MemoryManager> m_ExternalMemoryManager;
 
     std::unordered_map<BackendId, bool> m_SupportsExternallyManagedMemory;
+
+    // A set of vectors to record the workload queue indexes and their corresponding Input/Output Slot indexes
+    // which are connected to Inputs and Outputs for the network.
+    std::vector<std::pair<unsigned int, unsigned int>> m_InputWorkloadSlotPairs;
+    std::vector<std::pair<unsigned int, unsigned int>> m_OutputWorkloadSlotPairs;
 };
 
 }
-- 
cgit v1.2.1