diff --git a/include/onnxruntime/core/framework/allocator.h b/include/onnxruntime/core/framework/allocator.h index 085284d1bb..082ecf4fa3 100644 --- a/include/onnxruntime/core/framework/allocator.h +++ b/include/onnxruntime/core/framework/allocator.h @@ -22,14 +22,14 @@ struct OrtAllocatorInfo { OrtMemType mem_type; OrtAllocatorType type; - constexpr OrtAllocatorInfo(const char* name1, OrtAllocatorType type, int id1 = 0, OrtMemType mem_type1 = OrtMemTypeDefault) + constexpr OrtAllocatorInfo(const char* name_, OrtAllocatorType type_, int id_ = 0, OrtMemType mem_type_ = OrtMemTypeDefault) #if (defined(__GNUC__) || defined(__clang__)) __attribute__((nonnull)) #endif - : name(name1), - id(id1), - mem_type(mem_type1), - type(type) { + : name(name_), + id(id_), + mem_type(mem_type_), + type(type_) { } inline bool operator==(const OrtAllocatorInfo& other) const { @@ -182,9 +182,21 @@ class IDeviceAllocator : public IAllocator { class CPUAllocator : public IDeviceAllocator { public: + explicit CPUAllocator(std::unique_ptr allocator_info) { + ORT_ENFORCE(nullptr != allocator_info); + allocator_info_ = std::move(allocator_info); + } + + CPUAllocator() { + allocator_info_ = std::make_unique(CPU, OrtAllocatorType::OrtDeviceAllocator); + } + void* Alloc(size_t size) override; void Free(void* p) override; const OrtAllocatorInfo& Info() const override; + + private: + std::unique_ptr allocator_info_; }; using AllocatorPtr = std::shared_ptr; diff --git a/onnxruntime/core/framework/allocator.cc b/onnxruntime/core/framework/allocator.cc index 25deadb2e0..e160ce4033 100644 --- a/onnxruntime/core/framework/allocator.cc +++ b/onnxruntime/core/framework/allocator.cc @@ -37,8 +37,7 @@ void CPUAllocator::Free(void* p) { } const OrtAllocatorInfo& CPUAllocator::Info() const { - static constexpr OrtAllocatorInfo cpuAllocatorInfo(CPU, OrtAllocatorType::OrtDeviceAllocator); - return cpuAllocatorInfo; + return *allocator_info_; } } // namespace onnxruntime diff --git a/onnxruntime/core/framework/execution_providers.h b/onnxruntime/core/framework/execution_providers.h index fdb51f0e18..dc471b880f 100644 --- a/onnxruntime/core/framework/execution_providers.h +++ b/onnxruntime/core/framework/execution_providers.h @@ -74,6 +74,15 @@ class ExecutionProviders { return exec_providers_[it->second].get(); } + AllocatorPtr GetAllocator(const OrtAllocatorInfo& allocator_info) const { + auto exec_provider = Get(allocator_info); + if (exec_provider == nullptr) { + return nullptr; + } + + return exec_provider->GetAllocator(allocator_info.id, allocator_info.mem_type); + } + bool Empty() const { return exec_providers_.empty(); } size_t NumProviders() const { return exec_providers_.size(); } diff --git a/onnxruntime/core/framework/utils.cc b/onnxruntime/core/framework/utils.cc index be1b289725..fe996fdc1a 100644 --- a/onnxruntime/core/framework/utils.cc +++ b/onnxruntime/core/framework/utils.cc @@ -30,16 +30,11 @@ const KernelDef* GetKernelDef(const KernelRegistryManager& kernel_registry, } AllocatorPtr GetAllocator(const ExecutionProviders& exec_providers, const OrtAllocatorInfo& allocator_info) { - auto exec_provider = exec_providers.Get(allocator_info); - if (exec_provider == nullptr) { - return nullptr; - } - - return exec_provider->GetAllocator(allocator_info.id, allocator_info.mem_type); + return exec_providers.GetAllocator(allocator_info); } AllocatorPtr GetAllocator(const SessionState& session_state, const OrtAllocatorInfo& allocator_info) { - return GetAllocator(session_state.GetExecutionProviders(), allocator_info); + return session_state.GetExecutionProviders().GetAllocator(allocator_info); } common::Status AllocateHelper(const IExecutionProvider& execution_provider, diff --git a/onnxruntime/core/providers/mkldnn/mkldnn_allocator.cc b/onnxruntime/core/providers/mkldnn/mkldnn_allocator.cc deleted file mode 100644 index a58af56b44..0000000000 --- a/onnxruntime/core/providers/mkldnn/mkldnn_allocator.cc +++ /dev/null @@ -1,18 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "mkldnn_allocator.h" -#include "core/framework/allocatormgr.h" - -namespace onnxruntime { - -const OrtAllocatorInfo& MKLDNNAllocator::Info() const { - static constexpr OrtAllocatorInfo mkl_allocator_info(MKLDNN, OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemTypeDefault); - return mkl_allocator_info; -} - -const OrtAllocatorInfo& MKLDNNCPUAllocator::Info() const { - static constexpr OrtAllocatorInfo mkl_cpu_allocator_info(MKLDNN_CPU, OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemTypeCPUOutput); - return mkl_cpu_allocator_info; -} -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/mkldnn/mkldnn_allocator.h b/onnxruntime/core/providers/mkldnn/mkldnn_allocator.h deleted file mode 100644 index ac888e8321..0000000000 --- a/onnxruntime/core/providers/mkldnn/mkldnn_allocator.h +++ /dev/null @@ -1,20 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once -#include "core/framework/allocator.h" - -// Placeholder for an MKL allocators -namespace onnxruntime { -constexpr const char* MKLDNN = "MklDnn"; -constexpr const char* MKLDNN_CPU = "MklDnnCpu"; - -class MKLDNNAllocator : public CPUAllocator { - public: - const OrtAllocatorInfo& Info() const override; -}; -class MKLDNNCPUAllocator : public CPUAllocator { - public: - const OrtAllocatorInfo& Info() const override; -}; -} // namespace onnxruntime diff --git a/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.cc b/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.cc index 5862e21e5f..1a6f1b8844 100644 --- a/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.cc +++ b/onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.cc @@ -1,7 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "mkldnn_allocator.h" #include "mkldnn_execution_provider.h" #include "core/framework/allocator.h" #include "core/framework/memcpy.h" @@ -9,6 +8,10 @@ #include "mkldnn_fwd.h" namespace onnxruntime { + +constexpr const char* MKLDNN = "MklDnn"; +constexpr const char* MKLDNN_CPU = "MklDnnCpu"; + namespace mkl_dnn { ONNX_OPERATOR_KERNEL_EX( @@ -31,11 +34,11 @@ ONNX_OPERATOR_KERNEL_EX( MKLDNNExecutionProvider::MKLDNNExecutionProvider(const MKLDNNExecutionProviderInfo& /*info*/) { DeviceAllocatorRegistrationInfo default_allocator_info({OrtMemTypeDefault, - [](int) { return std::make_unique(); }, std::numeric_limits::max()}); + [](int) { return std::make_unique(std::make_unique(MKLDNN, OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemTypeDefault)); }, std::numeric_limits::max()}); InsertAllocator(CreateAllocator(default_allocator_info)); DeviceAllocatorRegistrationInfo cpu_allocator_info({OrtMemTypeCPUOutput, - [](int) { return std::make_unique(); }, std::numeric_limits::max()}); + [](int) { return std::make_unique(std::make_unique(MKLDNN_CPU, OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemTypeCPUOutput)); }, std::numeric_limits::max()}); InsertAllocator(CreateAllocator(cpu_allocator_info)); } @@ -97,8 +100,6 @@ std::shared_ptr GetMklDnnKernelRegistry() { } } // namespace mkl_dnn - - std::shared_ptr MKLDNNExecutionProvider::GetKernelRegistry() const { static std::shared_ptr kernel_registry = onnxruntime::mkl_dnn::GetMklDnnKernelRegistry(); return kernel_registry;