saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 17:15:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

allocator refactor (#467)

Browse source 
* update CPUAllocator.

* onnxruntime

* fix build break

* remove useless subclasses of CPUAllocator.

* refactor to get allocator from executionproviders instead of execution provider.
This commit is contained in:
Ke Zhang 2019-02-12 14:14:21 -08:00 committed by GitHub
parent 0c4fef9ac2
commit fc90a9b2fc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 35 additions and 57 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
include/onnxruntime/core/framework/allocator.h
View file

@ -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<OrtAllocatorInfo> allocator_info) {
ORT_ENFORCE(nullptr != allocator_info);
allocator_info_ = std::move(allocator_info);
}
CPUAllocator() {
allocator_info_ = std::make_unique<OrtAllocatorInfo>(CPU, OrtAllocatorType::OrtDeviceAllocator);
}
void* Alloc(size_t size) override;
void Free(void* p) override;
const OrtAllocatorInfo& Info() const override;
private:
std::unique_ptr<OrtAllocatorInfo> allocator_info_;
};
using AllocatorPtr = std::shared_ptr<IAllocator>;

3
onnxruntime/core/framework/allocator.cc
View file

@ -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

9
onnxruntime/core/framework/execution_providers.h
View file

@ -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(); }

9
onnxruntime/core/framework/utils.cc
View file

@ -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,

18
onnxruntime/core/providers/mkldnn/mkldnn_allocator.cc
View file

@ -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

20
onnxruntime/core/providers/mkldnn/mkldnn_allocator.h
View file

@ -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

11
onnxruntime/core/providers/mkldnn/mkldnn_execution_provider.cc
View file

@ -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<MKLDNNAllocator>(); }, std::numeric_limits<size_t>::max()});
[](int) { return std::make_unique<CPUAllocator>(std::make_unique<OrtAllocatorInfo>(MKLDNN, OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemTypeDefault)); }, std::numeric_limits<size_t>::max()});
InsertAllocator(CreateAllocator(default_allocator_info));
DeviceAllocatorRegistrationInfo cpu_allocator_info({OrtMemTypeCPUOutput,
[](int) { return std::make_unique<MKLDNNCPUAllocator>(); }, std::numeric_limits<size_t>::max()});
[](int) { return std::make_unique<CPUAllocator>(std::make_unique<OrtAllocatorInfo>(MKLDNN_CPU, OrtAllocatorType::OrtDeviceAllocator, 0, OrtMemTypeCPUOutput)); }, std::numeric_limits<size_t>::max()});
InsertAllocator(CreateAllocator(cpu_allocator_info));
}
@ -97,8 +100,6 @@ std::shared_ptr<KernelRegistry> GetMklDnnKernelRegistry() {
}
} // namespace mkl_dnn
std::shared_ptr<KernelRegistry> MKLDNNExecutionProvider::GetKernelRegistry() const {
static std::shared_ptr<KernelRegistry> kernel_registry = onnxruntime::mkl_dnn::GetMklDnnKernelRegistry();
return kernel_registry;