diff options
Diffstat (limited to 'include/armnn')
| -rw-r--r-- | include/armnn/BackendRegistry.hpp | 5 | ||||
| -rw-r--r-- | include/armnn/IRuntime.hpp | 20 | ||||
| -rw-r--r-- | include/armnn/backends/IBackendInternal.hpp | 5 | ||||
| -rw-r--r-- | include/armnn/backends/ICustomAllocator.hpp | 18 |
4 files changed, 35 insertions, 13 deletions
diff --git a/include/armnn/BackendRegistry.hpp b/include/armnn/BackendRegistry.hpp index fe6451cde0..c13aa9f8b6 100644 --- a/include/armnn/BackendRegistry.hpp +++ b/include/armnn/BackendRegistry.hpp @@ -7,6 +7,7 @@ #include <armnn/Types.hpp> #include <armnn/BackendId.hpp> #include <armnn/Optional.hpp> +#include <armnn/backends/ICustomAllocator.hpp> #include <memory> #include <unordered_map> @@ -35,6 +36,8 @@ public: BackendIdSet GetBackendIds() const; std::string GetBackendIdsAsString() const; void SetProfilingService(armnn::Optional<profiling::ProfilingService&> profilingService); + void RegisterAllocator(const BackendId& id, std::shared_ptr<ICustomAllocator> alloc); + std::unordered_map<BackendId, std::shared_ptr<ICustomAllocator>> GetAllocators(); BackendRegistry() {} virtual ~BackendRegistry() {} @@ -50,6 +53,7 @@ public: }; void Deregister(const BackendId& id); + void DeregisterAllocator(const BackendId &id); protected: using FactoryStorage = std::unordered_map<BackendId, FactoryFunction>; @@ -63,6 +67,7 @@ private: FactoryStorage m_Factories; armnn::Optional<profiling::ProfilingService&> m_ProfilingService; + std::unordered_map<BackendId, std::shared_ptr<ICustomAllocator>> m_CustomMemoryAllocatorMap; }; BackendRegistry& BackendRegistryInstance(); diff --git a/include/armnn/IRuntime.hpp b/include/armnn/IRuntime.hpp index 8c269dee49..97a9c2889e 100644 --- a/include/armnn/IRuntime.hpp +++ b/include/armnn/IRuntime.hpp @@ -16,6 +16,7 @@ #include <armnn/backends/ICustomAllocator.hpp> #include <memory> +#include <map> namespace armnn { @@ -103,8 +104,8 @@ public: : m_GpuAccTunedParameters(nullptr) , m_EnableGpuProfiling(false) , m_DynamicBackendsPath("") - , m_CustomAllocator(nullptr) , m_ProtectedMode(false) + , m_CustomAllocatorMap() {} /// If set, uses the GpuAcc tuned parameters from the given object when executing GPU workloads. @@ -118,17 +119,22 @@ public: /// Only a single path is allowed for the override std::string m_DynamicBackendsPath; - /// A Custom Allocator used for allocation of working memory in the backends. - /// Set this for when you need to allocate Protected Working Memory, required for ProtectedMode - /// Only supported for GpuAcc - ICustomAllocator* m_CustomAllocator; - /// Setting this flag will allow the user to create the Runtime in protected mode. /// It will run all the inferences on protected memory and will make sure that /// INetworkProperties::m_ImportEnabled set to true with MemorySource::DmaBufProtected option - /// This will use Protected Memory Allocator associated with the backend + /// This requires that the backend supports Protected Memory and has an allocator capable of + /// allocating Protected Memory associated with it. bool m_ProtectedMode; + /// @brief A map to define a custom memory allocator for specific backend Ids. + /// + /// @details A Custom Allocator is used for allocation of working memory in the backends. + /// Set this if you need to take control of how memory is allocated on a backend. Required for + /// Protected Mode in order to correctly allocate Protected Memory + /// + /// @note Only supported for GpuAcc + std::map<BackendId, std::shared_ptr<ICustomAllocator>> m_CustomAllocatorMap; + struct ExternalProfilingOptions { ExternalProfilingOptions() diff --git a/include/armnn/backends/IBackendInternal.hpp b/include/armnn/backends/IBackendInternal.hpp index 3b4ef95703..626746465f 100644 --- a/include/armnn/backends/IBackendInternal.hpp +++ b/include/armnn/backends/IBackendInternal.hpp @@ -199,10 +199,13 @@ public: /// Signals the backend to use a custom memory allocator provided by the user /// + /// \param allocator - a pointer to the provided ICustomAllocator to use with this backend /// \param errMsg - Optional string variable to return error messages /// \return - Returns true if switching to custom allocator was successful - virtual bool UseCustomMemoryAllocator(armnn::Optional<std::string&> errMsg) + virtual bool UseCustomMemoryAllocator(std::shared_ptr<ICustomAllocator> allocator, + armnn::Optional<std::string&> errMsg) { + IgnoreUnused(allocator); if (errMsg) { std::stringstream message; diff --git a/include/armnn/backends/ICustomAllocator.hpp b/include/armnn/backends/ICustomAllocator.hpp index 1d4df0cb86..92cbcc2641 100644 --- a/include/armnn/backends/ICustomAllocator.hpp +++ b/include/armnn/backends/ICustomAllocator.hpp @@ -7,6 +7,7 @@ #include <cstddef> #include <memory> +#include <armnn/MemorySources.hpp> namespace armnn { @@ -23,13 +24,20 @@ public: * @param[in] alignment Alignment that the returned pointer should comply with * * @return A pointer to the allocated memory + * The returned pointer must be host write accessible */ - virtual void *allocate(size_t size, size_t alignment) = 0; - /** Interface to be implemented by the child class to free the allocated tensor */ - virtual void free(void *ptr) = 0; + virtual void* allocate(size_t size, size_t alignment) = 0; - // Utility Function to define the Custom Memory Allocators capabilities - virtual bool SupportsProtectedMemory() = 0; + /** Interface to be implemented by the child class to free the allocated bytes */ + virtual void free(void* ptr) = 0; + + // Used to specify what type of memory is being allocated by this allocator. + // Supported types are: + // MemorySource::Malloc + // Unsupported types are: + // MemorySource::DmaBuf + // MemorySource::DmaBufProtected + virtual armnn::MemorySource GetMemorySourceType() = 0; }; } // namespace armnn
\ No newline at end of file |