diff options
Diffstat (limited to 'src/armnn')
-rw-r--r-- | src/armnn/BackendRegistry.cpp | 22 | ||||
-rw-r--r-- | src/armnn/Runtime.cpp | 40 |
2 files changed, 62 insertions, 0 deletions
diff --git a/src/armnn/BackendRegistry.cpp b/src/armnn/BackendRegistry.cpp index 80daed9896..ade844fc39 100644 --- a/src/armnn/BackendRegistry.cpp +++ b/src/armnn/BackendRegistry.cpp @@ -128,4 +128,26 @@ std::unordered_map<BackendId, std::shared_ptr<ICustomAllocator>> BackendRegistry return m_CustomMemoryAllocatorMap; } +void BackendRegistry::RegisterMemoryOptimizerStrategy(const BackendId& id, + std::shared_ptr<IMemoryOptimizerStrategy> strategy) +{ + if (m_MemoryOptimizerStrategyMap.find(id) != m_MemoryOptimizerStrategyMap.end()) + { + throw InvalidArgumentException( + std::string(id) + " already has an memory optimizer strategy associated with it", + CHECK_LOCATION()); + } + m_MemoryOptimizerStrategyMap[id] = strategy; +} + +void BackendRegistry::DeregisterMemoryOptimizerStrategy(const BackendId &id) +{ + m_MemoryOptimizerStrategyMap.erase(id); +} + +MemoryOptimizerStrategiesMapRef BackendRegistry::GetMemoryOptimizerStrategies() +{ + return m_MemoryOptimizerStrategyMap; +} + } // namespace armnn diff --git a/src/armnn/Runtime.cpp b/src/armnn/Runtime.cpp index 085cf2cee8..50068ebe36 100644 --- a/src/armnn/Runtime.cpp +++ b/src/armnn/Runtime.cpp @@ -12,6 +12,7 @@ #include <armnn/backends/IBackendContext.hpp> #include <backendsCommon/DynamicBackendUtils.hpp> +#include <backendsCommon/MemoryOptimizerStrategyLibrary.hpp> #include <armnn/utility/PolymorphicDowncast.hpp> #include <common/include/LabelsAndEventClasses.hpp> @@ -373,6 +374,45 @@ RuntimeImpl::RuntimeImpl(const IRuntime::CreationOptions& options) BackendRegistryInstance().RegisterAllocator(id, customAllocatorMapIterator->second); } } + + // check if custom memory optimizer strategy map is set + if (!options.m_MemoryOptimizerStrategyMap.empty()) + { + auto customMemoryOptimizerStrategyMapIterator = options.m_MemoryOptimizerStrategyMap.find(id); + // if a memory optimizer strategy is provided make the backend use that instead of the default + if (customMemoryOptimizerStrategyMapIterator != options.m_MemoryOptimizerStrategyMap.end()) + { + // no errors.. register the memory optimizer strategy with the BackendRegistry + BackendRegistryInstance().RegisterMemoryOptimizerStrategy( + id, customMemoryOptimizerStrategyMapIterator->second); + + ARMNN_LOG(info) << "MemoryOptimizerStrategy " + << customMemoryOptimizerStrategyMapIterator->second->GetName() + << " set for the backend " << id << "."; + } + } + else + { + // check if to use one of the existing memory optimizer strategies is set + std::string memoryOptimizerStrategyName = ""; + ParseOptions(options.m_BackendOptions, id, [&](std::string name, const BackendOptions::Var& value) + { + if (name == "MemoryOptimizerStrategy") + { + memoryOptimizerStrategyName = ParseStringBackendOption(value, ""); + } + }); + if (memoryOptimizerStrategyName != "") + { + MemoryOptimizerStrategyLibrary memoryOptimizerStrategyLibrary; + if (memoryOptimizerStrategyLibrary.SetMemoryOptimizerStrategy(id, memoryOptimizerStrategyName)) + { + ARMNN_LOG(info) << "MemoryOptimizerStrategy " + << memoryOptimizerStrategyName << " set for the backend " << id << "."; + } + } + } + auto context = backend->CreateBackendContext(options); // backends are allowed to return nullptrs if they |