diff options
Diffstat (limited to 'src')
-rw-r--r-- | src/armnn/LoadedNetwork.cpp | 33 | ||||
-rw-r--r-- | src/armnn/LoadedNetwork.hpp | 5 |
2 files changed, 38 insertions, 0 deletions
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; }; } |