diff options
Diffstat (limited to 'src/armnn/NetworkUtils.cpp')
-rw-r--r-- | src/armnn/NetworkUtils.cpp | 55 |
1 files changed, 54 insertions, 1 deletions
diff --git a/src/armnn/NetworkUtils.cpp b/src/armnn/NetworkUtils.cpp index 9a4ce87b59..735a6244d5 100644 --- a/src/armnn/NetworkUtils.cpp +++ b/src/armnn/NetworkUtils.cpp @@ -5,6 +5,12 @@ #include "NetworkUtils.hpp" +#include "SubGraphSelector.hpp" + +#include <armnn/Exceptions.hpp> + +#include <backendsCommon/BackendRegistry.hpp> + namespace armnn { @@ -74,7 +80,6 @@ std::vector<ConvertFp32ToFp16Layer*> InsertConvertFp32ToFp16LayersAfter(Graph& g return convertLayers; } - std::vector<DebugLayer*> InsertDebugLayerAfter(Graph& graph, Layer& layer) { std::vector<DebugLayer*> debugLayers; @@ -97,10 +102,58 @@ std::vector<DebugLayer*> InsertDebugLayerAfter(Graph& graph, Layer& layer) debugLayer->GetOutputSlot().SetTensorInfo(debugInfo); + // NOTE: It is OK to do this because DebugLayer is only supported on CpuRef + debugLayer->SetBackendId(Compute::CpuRef); + debugLayers.emplace_back(debugLayer); } return debugLayers; } +PreCompiledLayer* CreatePreCompiledLayer(Graph& graph, + const SubGraph& subGraph, + unsigned int subGraphIndex, + const IBackendInternalUniquePtr& backendObjPtr) +{ + BOOST_ASSERT(backendObjPtr); + + IBackendInternal::ISubGraphConverterPtr converter = + backendObjPtr->CreateSubGraphConverter(std::make_shared<SubGraph>(subGraph)); + if (!converter) + { + return nullptr; + } + + try + { + // Attempt to convert and compile sub-graph + auto preCompiledObject = converter->GetOutput(); + } + catch (std::exception&) + { + return nullptr; + } + + // Create pre-compiled layer + std::string name = "pre-compiled" + std::to_string(subGraphIndex); + PreCompiledLayer* preCompiledLayer = graph.AddLayer<PreCompiledLayer>( + PreCompiledDescriptor(subGraph.GetNumInputSlots(), subGraph.GetNumOutputSlots()), name.c_str()); + + // Copy output tensor infos from sub-graph + for (unsigned int i = 0u; i < subGraph.GetNumOutputSlots(); i++) + { + preCompiledLayer->GetOutputSlot(i).SetTensorInfo(subGraph.GetOutputSlot(i)->GetTensorInfo()); + } + + // Assign pre-compiled object to layer + preCompiledLayer->SetPreCompiledObject(converter->GetOutput()); + + // Set the backend-id for the pre-compiled layer + BackendId backendId = backendObjPtr->GetId(); + preCompiledLayer->SetBackendId(backendId); + + return preCompiledLayer; +} + } // namespace armnn |