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

Factor out IAllocator so that it can be shared with shared providers (#5567)

Browse source 
* Factor out IAllocator so shared providers can use it directly.
This commit is contained in:
Ryan Hill 2020-10-27 17:28:17 -07:00 committed by GitHub
parent e5b0d192f4
commit e90b6f06d1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
18 changed files with 93 additions and 236 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -3,17 +3,8 @@
#pragma once
#include <functional>
#include <map>
#include <string>
#include <cstring>
#include <type_traits>
#include "core/common/common.h"
#include "core/common/exceptions.h"
#include "core/common/status.h"
#include "core/framework/fence.h"
#include "core/graph/basic_types.h"
#include "core/session/onnxruntime_c_api.h"
#include "ortdevice.h"
#include "ortmemoryinfo.h"

1
include/onnxruntime/core/graph/node_arg.h
View file

@ -3,6 +3,7 @@
#pragma once
#include "core/graph/basic_types.h"
#include "core/graph/onnx_protobuf.h"
namespace onnxruntime {

3
onnxruntime/core/framework/allocatormgr.h
View file

@ -4,8 +4,7 @@
#pragma once
#include "core/common/common.h"
#include "core/framework/arena.h"
#include "core/framework/bfc_arena.h"
#include "core/framework/allocator.h"
#include "core/session/onnxruntime_c_api.h"
namespace onnxruntime {

1
onnxruntime/core/framework/node_index_info.h
View file

@ -7,6 +7,7 @@
#include "core/common/common.h"
#include "core/framework/ml_value.h"
#include "core/graph/basic_types.h"
namespace onnxruntime {
class GraphNodes;

88
onnxruntime/core/framework/provider_bridge_ort.cc
View file

@ -92,37 +92,6 @@ namespace onnxruntime {
ProviderHost* g_host{};
struct Provider_AllocatorPtr_Impl : Provider_IAllocator {
Provider_AllocatorPtr_Impl(AllocatorPtr p) : Provider_IAllocator{p->Info()}, p_{p} {}
void* Alloc(size_t size) override { return p_->Alloc(size); }
void Free(void* p) override { return p_->Free(p); }
AllocatorPtr p_;
};
// This is really a IAllocator, but we wrap it with this class to make it into a Provider_IAllocator
struct Provider_IAllocator_Impl : Provider_IAllocator {
Provider_IAllocator_Impl(std::unique_ptr<IAllocator> p) : Provider_IAllocator{p->Info()}, p_{std::move(p)} {}
void* Alloc(size_t size) override { return p_->Alloc(size); }
void Free(void* p) override { return p_->Free(p); }
bool IsProviderInterface() const override { return false; }
std::unique_ptr<IAllocator> p_;
};
// This is really a Provider_IAllocator, but we wrap it with this class to make it into a IAllocator
struct ProviderAllocator : IAllocator {
ProviderAllocator(std::shared_ptr<Provider_IAllocator> p) : IAllocator{p->memory_info_}, p_{std::move(p)} {}
void* Alloc(size_t size) override { return p_->Alloc(size); }
void Free(void* p) override { return p_->Free(p); }
std::shared_ptr<Provider_IAllocator> p_;
};
struct Provider_TensorShapeProto_Dimension_Iterator_Impl : Provider_TensorShapeProto_Dimension_Iterator {
Provider_TensorShapeProto_Dimension_Iterator_Impl(google::protobuf::internal::RepeatedPtrIterator<const onnx::TensorShapeProto_Dimension>&& v) : v_{std::move(v)} {}
@ -206,22 +175,19 @@ struct Provider_IExecutionProvider_Router_Impl : Provider_IExecutionProvider_Rou
return outer_->Provider_Compile(fused_nodes, node_compute_funcs);
}
Provider_AllocatorPtr Provider_GetAllocator(int id, OrtMemType mem_type) const override {
return std::make_shared<Provider_AllocatorPtr_Impl>(IExecutionProvider::GetAllocator(id, mem_type));
AllocatorPtr Provider_GetAllocator(int id, OrtMemType mem_type) const override {
return IExecutionProvider::GetAllocator(id, mem_type);
}
AllocatorPtr GetAllocator(int id, OrtMemType mem_type) const override {
auto allocator = outer_->Provider_GetAllocator(id, mem_type);
if (!allocator)
return nullptr;
return static_cast<Provider_AllocatorPtr_Impl*>(allocator.get())->p_;
return outer_->Provider_GetAllocator(id, mem_type);
}
std::unique_ptr<Provider_IDataTransfer> Provider_GetDataTransfer() const override { return IExecutionProvider::GetDataTransfer(); }
std::unique_ptr<IDataTransfer> GetDataTransfer() const override { return outer_->Provider_GetDataTransfer(); }
void Provider_InsertAllocator(Provider_AllocatorPtr allocator) override {
IExecutionProvider::InsertAllocator(static_cast<Provider_AllocatorPtr_Impl*>(allocator.get())->p_);
void Provider_InsertAllocator(AllocatorPtr allocator) override {
IExecutionProvider::InsertAllocator(allocator);
}
const logging::Logger* GetLogger() const override { return IExecutionProvider::GetLogger(); }
@ -236,30 +202,12 @@ struct ProviderHostImpl : ProviderHost {
DataTypeImpl_GetTensorType_float = &DataTypeImpl::GetTensorType<float>;
}
Provider_AllocatorPtr CreateAllocator(const Provider_AllocatorCreationInfo& info) override {
AllocatorCreationInfo info_real{
[&info](int value) -> std::unique_ptr<IAllocator> {
auto allocator = info.factory(value);
// If the allocator is a provider interface, we need to wrap it with ProviderAllocator to turn it into an IAllocator
// Otherwise it's really a Provider_IAllocator_Impl, so we can just unwrap it to get back to the IAllocator inside
if (allocator->IsProviderInterface())
return onnxruntime::make_unique<ProviderAllocator>(std::move(allocator));
else
return std::move(static_cast<Provider_IAllocator_Impl*>(&*allocator)->p_);
},
info.device_id,
info.use_arena,
info.arena_cfg};
// info_real will always return a unique_ptr to an IAllocator, which might be a native IAllocator or a provider interface wrapped by ProviderAllocator.
// Either way we wrap it in a Provider_IAllocator_Impl to be unwrapped by Provider_InsertAllocator
return std::make_shared<Provider_AllocatorPtr_Impl>(onnxruntime::CreateAllocator(info_real));
AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info) override {
return onnxruntime::CreateAllocator(info);
}
std::unique_ptr<Provider_IAllocator> CreateCPUAllocator(
const OrtMemoryInfo& memory_info) override {
return onnxruntime::make_unique<Provider_IAllocator_Impl>(
onnxruntime::make_unique<CPUAllocator>(memory_info));
std::unique_ptr<IAllocator> CreateCPUAllocator(const OrtMemoryInfo& memory_info) override {
return onnxruntime::make_unique<CPUAllocator>(memory_info);
};
std::unique_ptr<Provider_IExecutionProvider_Router> Create_IExecutionProvider_Router(
@ -268,12 +216,12 @@ struct ProviderHostImpl : ProviderHost {
};
#ifdef USE_TENSORRT
std::unique_ptr<Provider_IAllocator> CreateCUDAAllocator(int16_t device_id, const char* name) override {
return onnxruntime::make_unique<Provider_IAllocator_Impl>(onnxruntime::make_unique<CUDAAllocator>(device_id, name));
std::unique_ptr<IAllocator> CreateCUDAAllocator(int16_t device_id, const char* name) override {
return onnxruntime::make_unique<CUDAAllocator>(device_id, name);
}
std::unique_ptr<Provider_IAllocator> CreateCUDAPinnedAllocator(int16_t device_id, const char* name) override {
return onnxruntime::make_unique<Provider_IAllocator_Impl>(onnxruntime::make_unique<CUDAPinnedAllocator>(device_id, name));
std::unique_ptr<IAllocator> CreateCUDAPinnedAllocator(int16_t device_id, const char* name) override {
return onnxruntime::make_unique<CUDAPinnedAllocator>(device_id, name);
}
std::unique_ptr<Provider_IDataTransfer> CreateGPUDataTransfer() override { return onnxruntime::make_unique<GPUDataTransfer>(); }
@ -313,6 +261,16 @@ struct ProviderHostImpl : ProviderHost {
return ::onnxruntime::LogRuntimeError(session_id, status, file, function, line);
}
// IAllocator
bool IAllocator__CalcMemSizeForArrayWithAlignment(size_t nmemb, size_t size, size_t alignment, size_t* out) override { return IAllocator::CalcMemSizeForArrayWithAlignment(nmemb, size, alignment, out); }
// Status
std::string Status__ToString(const Status* p) override { return p->ToString(); }
// TensorShape
int64_t TensorShape__SizeHelper(const TensorShape* p, size_t start, size_t end) override { return p->SizeHelper(start, end); }
std::string TensorShape__ToString(const TensorShape* p) override { return p->ToString(); }
// CPUIDInfo
const CPUIDInfo& CPUIDInfo__GetCPUIDInfo() override { return CPUIDInfo::GetCPUIDInfo(); }
bool CPUIDInfo__HasAVX2(const CPUIDInfo* p) override { return p->HasAVX2(); }

10
onnxruntime/core/providers/dnnl/dnnl_execution_provider.cc
View file

@ -18,16 +18,16 @@ constexpr const char* DNNL_CPU = "DnnlCpu";
DNNLExecutionProvider::DNNLExecutionProvider(const DNNLExecutionProviderInfo& info)
: Provider_IExecutionProvider{onnxruntime::kDnnlExecutionProvider} {
Provider_AllocatorCreationInfo default_memory_info(
AllocatorCreationInfo default_memory_info(
{[](int) {
return onnxruntime::Provider_CreateCPUAllocator(OrtMemoryInfo(DNNL, OrtAllocatorType::OrtDeviceAllocator));
return onnxruntime::CreateCPUAllocator(OrtMemoryInfo(DNNL, OrtAllocatorType::OrtDeviceAllocator));
}},
0, info.create_arena);
Provider_AllocatorCreationInfo cpu_memory_info(
AllocatorCreationInfo cpu_memory_info(
{[](int) {
return onnxruntime::Provider_CreateCPUAllocator(OrtMemoryInfo(DNNL_CPU, OrtAllocatorType::OrtDeviceAllocator, OrtDevice(), 0,
OrtMemTypeCPUOutput));
return onnxruntime::CreateCPUAllocator(OrtMemoryInfo(DNNL_CPU, OrtAllocatorType::OrtDeviceAllocator, OrtDevice(), 0,
OrtMemTypeCPUOutput));
}},
0, info.create_arena);

5
onnxruntime/core/providers/dnnl/subgraph/dnnl_conv.h
View file

@ -386,8 +386,7 @@ class DnnlConv : public DnnlKernel {
if (filter_dst_mem == nullptr) {
dnnl::memory src = dnnl::memory({{filter_dims_mkl}, DnnnType<T>(), filter_format_}, cpu_engine, (void*)filter_data);
IAllocatorUniquePtr<void> filter_reorder_buffer =
Provider_IAllocator::MakeUniquePtr<void>(alloc_, filter_size_);
IAllocatorUniquePtr<void> filter_reorder_buffer = IAllocator::MakeUniquePtr<void>(alloc_, filter_size_);
filter_dst_mem = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->weights_desc(), cpu_engine, filter_reorder_buffer.get()));
@ -437,7 +436,7 @@ class DnnlConv : public DnnlKernel {
}
auto src_size = conv_fwd_pd_.get()->src_desc().get_size();
src_reorder_buffer_ = Provider_IAllocator::MakeUniquePtr<void>(alloc_, src_size);
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, src_size);
src_mem_->set_data_handle(src_reorder_buffer_.get());
} else {
if (mklnode_ptr_->parent_nodes.empty()) {

8
onnxruntime/core/providers/dnnl/subgraph/dnnl_conv_batchnorm.h
View file

@ -395,8 +395,7 @@ class DnnlConvBatchNorm : public DnnlKernel {
if (filter_dst_mem == nullptr) {
dnnl::memory src = dnnl::memory({{filter_dims_mkl}, DnnnType<T>(), filter_format_}, cpu_engine, (void*)weights_scaled_by_axis.data());
IAllocatorUniquePtr<void> filter_reorder_buffer =
Provider_IAllocator::MakeUniquePtr<void>(alloc_, filter_size_);
IAllocatorUniquePtr<void> filter_reorder_buffer = IAllocator::MakeUniquePtr<void>(alloc_, filter_size_);
filter_dst_mem = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->weights_desc(), cpu_engine, filter_reorder_buffer.get()));
@ -410,8 +409,7 @@ class DnnlConvBatchNorm : public DnnlKernel {
std::shared_ptr<dnnl::memory> bias_mem = provider_->GetBiasMemoryBuffer(mklnode_ptr_->weight_name);
if (bias_mem == nullptr) {
auto bias_size = conv_fwd_pd_.get()->bias_desc().get_size();
IAllocatorUniquePtr<void> bias_buffer =
Provider_IAllocator::MakeUniquePtr<void>(alloc_, bias_size);
IAllocatorUniquePtr<void> bias_buffer = IAllocator::MakeUniquePtr<void>(alloc_, bias_size);
bias_mem = onnxruntime::make_unique<dnnl::memory>(
dnnl::memory(conv_fwd_pd_->bias_desc(), cpu_engine, bias_buffer.get()));
float* bias_buffer_data = static_cast<float*>(bias_buffer.get());
@ -480,7 +478,7 @@ class DnnlConvBatchNorm : public DnnlKernel {
}
auto src_size = conv_fwd_pd_.get()->src_desc().get_size();
src_reorder_buffer_ = Provider_IAllocator::MakeUniquePtr<void>(alloc_, src_size);
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, src_size);
src_mem_->set_data_handle(src_reorder_buffer_.get());
} else {
if (mklnode_ptr_->parent_nodes.empty()) {

2
onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.h
View file

@ -111,7 +111,7 @@ class DnnlKernel {
// memory used for reorders
std::unique_ptr<dnnl::memory> reorder_dst_mem_to_;
Provider_AllocatorPtr alloc_;
AllocatorPtr alloc_;
DNNLExecutionProvider* provider_;
};

2
onnxruntime/core/providers/dnnl/subgraph/dnnl_pool.h
View file

@ -198,7 +198,7 @@ class DnnlPool : public DnnlKernel {
}
auto src_size = fwd_primitive_desc_.get()->src_desc().get_size();
src_reorder_buffer_ = Provider_IAllocator::MakeUniquePtr<void>(alloc_, src_size);
src_reorder_buffer_ = IAllocator::MakeUniquePtr<void>(alloc_, src_size);
src_mem_->set_data_handle(src_reorder_buffer_.get());
} else {
if (mklnode_ptr_->parent_nodes.empty()) {

23
onnxruntime/core/providers/shared_library/provider_api.h
View file

@ -15,22 +15,14 @@
#include "onnx/common/stl_backports.h"
#include "core/common/common.h"
#include "core/common/const_pointer_container.h"
#include "core/session/onnxruntime_c_api.h"
#include "core/framework/ortdevice.h"
#include "core/framework/ortmemoryinfo.h"
#include "core/common/logging/severity.h"
#include "core/framework/allocator.h"
#include "core/framework/allocatormgr.h"
#include "core/framework/tensor_shape.h"
namespace onnxruntime {
namespace logging {
enum class Severity {
kVERBOSE = 0,
kINFO = 1,
kWARNING = 2,
kERROR = 3,
kFATAL = 4
};
enum class DataType {
SYSTEM = 0, ///< System data.
USER = 1 ///< Contains potentially sensitive user data.
@ -175,10 +167,9 @@ class DataTypeImpl {
template <typename T>
using IAllocatorUniquePtr = std::unique_ptr<T, std::function<void(T*)>>;
std::unique_ptr<Provider_IAllocator> Provider_CreateCPUAllocator(const OrtMemoryInfo& memory_info);
std::unique_ptr<Provider_IAllocator> Provider_CreateCUDAAllocator(int16_t device_id, const char* name);
std::unique_ptr<Provider_IAllocator> Provider_CreateCUDAPinnedAllocator(int16_t device_id, const char* name);
Provider_AllocatorPtr CreateAllocator(const Provider_AllocatorCreationInfo& info);
std::unique_ptr<IAllocator> CreateCPUAllocator(const OrtMemoryInfo& memory_info);
std::unique_ptr<IAllocator> CreateCUDAAllocator(int16_t device_id, const char* name);
std::unique_ptr<IAllocator> CreateCUDAPinnedAllocator(int16_t device_id, const char* name);
std::unique_ptr<Provider_IDataTransfer> Provider_CreateGPUDataTransfer();
@ -192,8 +183,6 @@ struct Category {
static const char* onnxruntime; ///< General output
};
constexpr const char* SEVERITY_PREFIX = "VIWEF";
} // namespace logging
namespace math {

62
onnxruntime/core/providers/shared_library/provider_bridge_provider.cc
View file

@ -48,7 +48,7 @@ void operator delete(void* p, size_t /*size*/) { return onnxruntime::g_host->Hea
namespace onnxruntime {
Provider_AllocatorPtr CreateAllocator(const Provider_AllocatorCreationInfo& info) {
AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info) {
return g_host->CreateAllocator(info);
}
@ -85,13 +85,7 @@ int64_t TensorShape::Size() const {
}
int64_t TensorShape::SizeHelper(size_t start, size_t end) const {
// Must return 1 for an empty sequence
int64_t size = 1;
for (size_t i = start; i < end; i++) {
if ((*this)[i] < 0) return -1;
size *= (*this)[i];
}
return size;
return g_host->TensorShape__SizeHelper(this, start, end);
}
TensorShape TensorShape::Slice(size_t dimstart, size_t dimend) const {
@ -104,37 +98,27 @@ TensorShape TensorShape::Slice(size_t dimstart) const {
}
std::string TensorShape::ToString() const {
std::string result;
result.append("{");
bool first = true;
for (auto dim : (*this)) {
if (!first) {
result.append(",");
}
result.append(std::to_string(dim));
first = false;
}
result.append("}");
return result;
return g_host->TensorShape__ToString(this);
}
Provider_AllocatorPtr CreateAllocator(Provider_AllocatorCreationInfo info) {
AllocatorPtr CreateAllocator(AllocatorCreationInfo info) {
return g_host->CreateAllocator(info);
}
std::unique_ptr<Provider_IAllocator> Provider_CreateCPUAllocator(const OrtMemoryInfo& info) {
std::unique_ptr<IAllocator> CreateCPUAllocator(const OrtMemoryInfo& info) {
return g_host->CreateCPUAllocator(info);
}
bool IAllocator::CalcMemSizeForArrayWithAlignment(size_t nmemb, size_t size, size_t alignment, size_t* out) noexcept {
return g_host->IAllocator__CalcMemSizeForArrayWithAlignment(nmemb, size, alignment, out);
}
#ifdef USE_TENSORRT
std::unique_ptr<Provider_IAllocator> Provider_CreateCUDAAllocator(int16_t device_id, const char* name) {
std::unique_ptr<IAllocator> CreateCUDAAllocator(int16_t device_id, const char* name) {
return g_host->CreateCUDAAllocator(device_id, name);
}
std::unique_ptr<Provider_IAllocator> Provider_CreateCUDAPinnedAllocator(int16_t device_id, const char* name) {
std::unique_ptr<IAllocator> CreateCUDAPinnedAllocator(int16_t device_id, const char* name) {
return g_host->CreateCUDAPinnedAllocator(device_id, name);
}
@ -181,29 +165,7 @@ const std::string& Status::ErrorMessage() const noexcept {
return IsOK() ? EmptyString() : state_->msg;
}
std::string Status::ToString() const {
if (state_ == nullptr) {
return std::string("OK");
}
std::string result;
if (common::SYSTEM == state_->category) {
result += "SystemError";
result += " : ";
result += std::to_string(errno);
} else if (common::ONNXRUNTIME == state_->category) {
result += "[ONNXRuntimeError]";
result += " : ";
result += std::to_string(Code());
result += " : ";
result += StatusCodeToString(static_cast<StatusCode>(Code()));
result += " : ";
result += state_->msg;
}
return result;
}
std::string Status::ToString() const { return g_host->Status__ToString(this); }
const std::string& Status::EmptyString() noexcept {
static std::string s_empty;

84
onnxruntime/core/providers/shared_library/provider_interfaces.h
View file

@ -74,62 +74,6 @@ struct Provider_IExecutionProviderFactory {
class DataTypeImpl;
using MLDataType = const DataTypeImpl*;
template <typename T>
using Provider_IAllocatorUniquePtr = std::unique_ptr<T, std::function<void(T*)>>;
struct Provider_IAllocator {
Provider_IAllocator(const OrtMemoryInfo& info) : memory_info_{info} {}
virtual ~Provider_IAllocator() {}
virtual void* Alloc(size_t size) = 0;
virtual void Free(void* p) = 0;
const OrtMemoryInfo& Info() const { return memory_info_; };
virtual bool IsProviderInterface() const { return true; }
template <typename T>
static Provider_IAllocatorUniquePtr<T> MakeUniquePtr(std::shared_ptr<Provider_IAllocator> allocator, size_t count_or_bytes) {
if (allocator == nullptr) return nullptr;
size_t alloc_size = count_or_bytes;
// if T is not void, 'count_or_bytes' == number of items so allow for that
if (!std::is_void<T>::value) {
// TODO: Use internal implementation to get correct sizes
return nullptr;
}
return Provider_IAllocatorUniquePtr<T>{
static_cast<T*>(allocator->Alloc(alloc_size)), // allocate
[=](T* ptr) { allocator->Free(ptr); }}; // capture IAllocator so it's always valid, and use as deleter
}
const OrtMemoryInfo memory_info_;
Provider_IAllocator(const Provider_IAllocator&) = delete;
void operator=(const Provider_IAllocator&) = delete;
};
using Provider_AllocatorPtr = std::shared_ptr<Provider_IAllocator>;
using Provider_AllocatorFactory = std::function<std::unique_ptr<Provider_IAllocator>(int)>;
using DeviceId = int16_t;
struct Provider_AllocatorCreationInfo {
Provider_AllocatorCreationInfo(Provider_AllocatorFactory device_alloc_factory0,
DeviceId device_id0 = 0,
bool use_arena0 = true,
OrtArenaCfg arena_cfg0 = {0, -1, -1, -1})
: factory(device_alloc_factory0),
device_id(device_id0),
use_arena(use_arena0),
arena_cfg(arena_cfg0) {
}
Provider_AllocatorFactory factory;
DeviceId device_id;
bool use_arena;
OrtArenaCfg arena_cfg;
};
struct Provider_OpKernel {
Provider_OpKernel() {}
virtual ~Provider_OpKernel() = default;
@ -143,7 +87,7 @@ struct Provider_OpKernel {
using NodeIndex = size_t;
using Provider_NodeArgInfo = Provider_ValueInfoProto;
// We can't just reinterpret_cast this one, since it's an unordered_map of object BY VALUE (can't do anything by value on the real types)
//using Provider_NodeAttributes = std::unordered_map<std::string, ONNX_NAMESPACE::Provider_AttributeProto_Copyable>;
// using Provider_NodeAttributes = std::unordered_map<std::string, ONNX_NAMESPACE::Provider_AttributeProto_Copyable>;
using Provider_InitializedTensorSet = std::unordered_map<std::string, const Provider_TensorProto*>;
@ -188,9 +132,9 @@ struct Provider_IExecutionProvider_Router {
virtual std::vector<std::unique_ptr<Provider_ComputeCapability>> Provider_GetCapability(const onnxruntime::Provider_GraphViewer& graph,
const std::vector<const Provider_KernelRegistry*>& kernel_registries) const = 0;
virtual Provider_AllocatorPtr Provider_GetAllocator(int id, OrtMemType mem_type) const = 0;
virtual AllocatorPtr Provider_GetAllocator(int id, OrtMemType mem_type) const = 0;
virtual std::unique_ptr<Provider_IDataTransfer> Provider_GetDataTransfer() const = 0;
virtual void Provider_InsertAllocator(Provider_AllocatorPtr allocator) = 0;
virtual void Provider_InsertAllocator(AllocatorPtr allocator) = 0;
virtual const logging::Logger* GetLogger() const = 0;
void operator=(const Provider_IExecutionProvider_Router&) = delete;
@ -209,8 +153,8 @@ struct Provider_IExecutionProvider {
virtual common::Status Provider_Compile(const std::vector<Provider_Node*>& fused_nodes, std::vector<NodeComputeInfo>& node_compute_funcs) = 0;
virtual Provider_AllocatorPtr Provider_GetAllocator(int id, OrtMemType mem_type) const { return p_->Provider_GetAllocator(id, mem_type); }
virtual void Provider_InsertAllocator(Provider_AllocatorPtr allocator) { return p_->Provider_InsertAllocator(allocator); }
virtual AllocatorPtr Provider_GetAllocator(int id, OrtMemType mem_type) const { return p_->Provider_GetAllocator(id, mem_type); }
virtual void Provider_InsertAllocator(AllocatorPtr allocator) { return p_->Provider_InsertAllocator(allocator); }
virtual const logging::Logger* GetLogger() const { return p_->GetLogger(); }
@ -230,15 +174,15 @@ struct Provider {
// calls the virtual function (which will lead to infinite recursion in the bridge). There is no known way to get the non virtual member
// function pointer implementation in this case.
struct ProviderHost {
virtual Provider_AllocatorPtr CreateAllocator(const Provider_AllocatorCreationInfo& info) = 0;
virtual AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info) = 0;
virtual logging::Logger* LoggingManager_GetDefaultLogger() = 0;
virtual std::unique_ptr<Provider_IAllocator> CreateCPUAllocator(const OrtMemoryInfo& memory_info) = 0;
virtual std::unique_ptr<IAllocator> CreateCPUAllocator(const OrtMemoryInfo& memory_info) = 0;
#ifdef USE_TENSORRT
virtual std::unique_ptr<Provider_IAllocator> CreateCUDAAllocator(int16_t device_id, const char* name) = 0;
virtual std::unique_ptr<Provider_IAllocator> CreateCUDAPinnedAllocator(int16_t device_id, const char* name) = 0;
virtual std::unique_ptr<IAllocator> CreateCUDAAllocator(int16_t device_id, const char* name) = 0;
virtual std::unique_ptr<IAllocator> CreateCUDAPinnedAllocator(int16_t device_id, const char* name) = 0;
virtual std::unique_ptr<Provider_IDataTransfer> CreateGPUDataTransfer() = 0;
virtual void cuda__Impl_Cast(const int64_t* input_data, int32_t* output_data, size_t count) = 0;
@ -267,6 +211,16 @@ struct ProviderHost {
virtual std::vector<std::string> GetStackTrace() = 0;
// IAllocator
virtual bool IAllocator__CalcMemSizeForArrayWithAlignment(size_t nmemb, size_t size, size_t alignment, size_t* out) = 0;
// Status
virtual std::string Status__ToString(const Status* p) = 0;
// TensorShape
virtual int64_t TensorShape__SizeHelper(const TensorShape* p, size_t start, size_t end) = 0;
virtual std::string TensorShape__ToString(const TensorShape* p) = 0;
// CPUIDInfo
virtual const CPUIDInfo& CPUIDInfo__GetCPUIDInfo() = 0;
virtual bool CPUIDInfo__HasAVX2(const CPUIDInfo* p) = 0;

11
onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
View file

@ -10,7 +10,6 @@
#include "core/common/safeint.h"
#include "tensorrt_execution_provider.h"
#include "core/providers/cuda/cuda_allocator.h"
#include "core/providers/cuda/shared_inc/cuda_call.h"
#include "core/providers/cuda/math/unary_elementwise_ops_impl.h"
#include "cuda_runtime_api.h"
@ -174,13 +173,13 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv
: Provider_IExecutionProvider{onnxruntime::kTensorrtExecutionProvider}, device_id_(info.device_id) {
CUDA_CALL_THROW(cudaSetDevice(device_id_));
Provider_AllocatorCreationInfo default_memory_info(
[](int id) { return Provider_CreateCUDAAllocator(id, TRT); }, device_id_);
AllocatorCreationInfo default_memory_info(
[](int id) { return CreateCUDAAllocator(id, TRT); }, device_id_);
allocator_ = CreateAllocator(default_memory_info);
Provider_InsertAllocator(allocator_);
Provider_AllocatorCreationInfo pinned_allocator_info(
[](int) { return Provider_CreateCUDAPinnedAllocator(0, TRT_PINNED); }, device_id_);
AllocatorCreationInfo pinned_allocator_info(
[](int) { return CreateCUDAPinnedAllocator(0, TRT_PINNED); }, device_id_);
Provider_InsertAllocator(CreateAllocator(pinned_allocator_info));
// Get environment variables
@ -227,7 +226,7 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv
TensorrtExecutionProvider::~TensorrtExecutionProvider() {}
Provider_AllocatorPtr TensorrtExecutionProvider::Provider_GetAllocator(int id, OrtMemType mem_type) const {
AllocatorPtr TensorrtExecutionProvider::Provider_GetAllocator(int id, OrtMemType mem_type) const {
if (mem_type == OrtMemTypeDefault) {
return allocator_;
} else {

15
onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h
View file

@ -32,11 +32,14 @@ class TensorrtLogger : public nvinfer1::ILogger {
strftime(&buf[0], 256,
"%Y-%m-%d %H:%M:%S",
std::gmtime(&rawtime));
const char* sevstr = (severity == Severity::kINTERNAL_ERROR ? " BUG" : severity == Severity::kERROR ? " ERROR" : severity == Severity::kWARNING ? "WARNING" : severity == Severity::kINFO ? " INFO" : "UNKNOWN");
if(severity <= Severity::kERROR)
LOGS_DEFAULT(ERROR) << "[" << buf << " " << sevstr << "] " << msg;
const char* sevstr = (severity == Severity::kINTERNAL_ERROR ? " BUG" : severity == Severity::kERROR ? " ERROR"
: severity == Severity::kWARNING ? "WARNING"
: severity == Severity::kINFO ? " INFO"
: "UNKNOWN");
if (severity <= Severity::kERROR)
LOGS_DEFAULT(ERROR) << "[" << buf << " " << sevstr << "] " << msg;
else
LOGS_DEFAULT(WARNING) << "[" << buf << " " << sevstr << "] " << msg;
LOGS_DEFAULT(WARNING) << "[" << buf << " " << sevstr << "] " << msg;
}
}
};
@ -101,7 +104,7 @@ class TensorrtExecutionProvider : public Provider_IExecutionProvider {
common::Status Provider_Compile(const std::vector<Provider_Node*>& fused_nodes,
std::vector<NodeComputeInfo>& node_compute_funcs) override;
Provider_AllocatorPtr Provider_GetAllocator(int id, OrtMemType mem_type) const override;
AllocatorPtr Provider_GetAllocator(int id, OrtMemType mem_type) const override;
private:
size_t max_workspace_size_ = 1 << 30; // 1GB
@ -140,7 +143,7 @@ class TensorrtExecutionProvider : public Provider_IExecutionProvider {
void RemoveTensorRTGraphCycles(SubGraphCollection_t& supported_nodes_vector, const onnxruntime::Provider_GraphViewer& graph) const;
Provider_AllocatorPtr allocator_;
AllocatorPtr allocator_;
};
} // namespace onnxruntime

1
orttraining/orttraining/models/bert/main.cc
View file

@ -8,6 +8,7 @@
#include "core/common/logging/sinks/clog_sink.h"
#include "core/common/profiler.h"
#include "core/session/environment.h"
#include "core/framework/bfc_arena.h"
#include "core/framework/random_seed.h"
#include "core/providers/cuda/cuda_allocator.h"
#include "orttraining/core/session/training_session.h"

1
orttraining/orttraining/models/gpt2/main.cc
View file

@ -8,6 +8,7 @@
#include "core/common/logging/sinks/clog_sink.h"
#include "core/session/environment.h"
#include "core/framework/random_seed.h"
#include "core/framework/bfc_arena.h"
#include "core/providers/cuda/cuda_allocator.h"
#include "orttraining/core/framework/mpi_context.h"
#include "orttraining/core/framework/tensorboard/event_writer.h"

3
orttraining/orttraining/models/mnist/main.cc
View file

@ -4,6 +4,7 @@
#include "cxxopts.hpp"
#include "core/common/logging/logging.h"
#include "core/common/logging/sinks/clog_sink.h"
#include "core/framework/bfc_arena.h"
#include "core/platform/env.h"
#include "core/session/environment.h"
#include "orttraining/core/session/training_session.h"
@ -83,7 +84,7 @@ Status ParseArguments(int argc, char* argv[], TrainingRunner::Parameters& params
#ifdef USE_CUDA
bool use_cuda = flags.count("use_cuda") > 0;
if (use_cuda) {
// Use local rank as device ID of the associated CUDA EP.
// Use local rank as device ID of the associated CUDA EP.
OrtDevice::DeviceId device_id = static_cast<OrtDevice::DeviceId>(MPIContext::GetInstance().GetLocalRank());
params.providers.emplace(kCudaExecutionProvider, CreateExecutionProviderFactory_CUDA(device_id));
}