From e90b6f06d13aa64812e3c55bc43240f9111de29c Mon Sep 17 00:00:00 2001 From: Ryan Hill <38674843+RyanUnderhill@users.noreply.github.com> Date: Tue, 27 Oct 2020 17:28:17 -0700 Subject: [PATCH] Factor out IAllocator so that it can be shared with shared providers (#5567) * Factor out IAllocator so shared providers can use it directly. --- .../onnxruntime/core/framework/allocator.h | 9 -- include/onnxruntime/core/graph/node_arg.h | 1 + onnxruntime/core/framework/allocatormgr.h | 3 +- onnxruntime/core/framework/node_index_info.h | 1 + .../core/framework/provider_bridge_ort.cc | 88 +++++-------------- .../providers/dnnl/dnnl_execution_provider.cc | 10 +-- .../core/providers/dnnl/subgraph/dnnl_conv.h | 5 +- .../dnnl/subgraph/dnnl_conv_batchnorm.h | 8 +- .../providers/dnnl/subgraph/dnnl_kernel.h | 2 +- .../core/providers/dnnl/subgraph/dnnl_pool.h | 2 +- .../providers/shared_library/provider_api.h | 23 ++--- .../provider_bridge_provider.cc | 62 +++---------- .../shared_library/provider_interfaces.h | 84 ++++-------------- .../tensorrt/tensorrt_execution_provider.cc | 11 ++- .../tensorrt/tensorrt_execution_provider.h | 15 ++-- orttraining/orttraining/models/bert/main.cc | 1 + orttraining/orttraining/models/gpt2/main.cc | 1 + orttraining/orttraining/models/mnist/main.cc | 3 +- 18 files changed, 93 insertions(+), 236 deletions(-) diff --git a/include/onnxruntime/core/framework/allocator.h b/include/onnxruntime/core/framework/allocator.h index 6ef1f3ec37..7559d44d43 100644 --- a/include/onnxruntime/core/framework/allocator.h +++ b/include/onnxruntime/core/framework/allocator.h @@ -3,17 +3,8 @@ #pragma once -#include -#include -#include -#include -#include - #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" diff --git a/include/onnxruntime/core/graph/node_arg.h b/include/onnxruntime/core/graph/node_arg.h index 31ba1a5fce..eeb53d2903 100644 --- a/include/onnxruntime/core/graph/node_arg.h +++ b/include/onnxruntime/core/graph/node_arg.h @@ -3,6 +3,7 @@ #pragma once +#include "core/graph/basic_types.h" #include "core/graph/onnx_protobuf.h" namespace onnxruntime { diff --git a/onnxruntime/core/framework/allocatormgr.h b/onnxruntime/core/framework/allocatormgr.h index cb7e141327..851fbba6f9 100644 --- a/onnxruntime/core/framework/allocatormgr.h +++ b/onnxruntime/core/framework/allocatormgr.h @@ -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 { diff --git a/onnxruntime/core/framework/node_index_info.h b/onnxruntime/core/framework/node_index_info.h index 19b4a202f5..b38a5fad7d 100644 --- a/onnxruntime/core/framework/node_index_info.h +++ b/onnxruntime/core/framework/node_index_info.h @@ -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; diff --git a/onnxruntime/core/framework/provider_bridge_ort.cc b/onnxruntime/core/framework/provider_bridge_ort.cc index cbc9aaa850..b5420c756d 100644 --- a/onnxruntime/core/framework/provider_bridge_ort.cc +++ b/onnxruntime/core/framework/provider_bridge_ort.cc @@ -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 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 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 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 p_; -}; - struct Provider_TensorShapeProto_Dimension_Iterator_Impl : Provider_TensorShapeProto_Dimension_Iterator { Provider_TensorShapeProto_Dimension_Iterator_Impl(google::protobuf::internal::RepeatedPtrIterator&& 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(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(allocator.get())->p_; + return outer_->Provider_GetAllocator(id, mem_type); } std::unique_ptr Provider_GetDataTransfer() const override { return IExecutionProvider::GetDataTransfer(); } std::unique_ptr GetDataTransfer() const override { return outer_->Provider_GetDataTransfer(); } - void Provider_InsertAllocator(Provider_AllocatorPtr allocator) override { - IExecutionProvider::InsertAllocator(static_cast(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; } - Provider_AllocatorPtr CreateAllocator(const Provider_AllocatorCreationInfo& info) override { - AllocatorCreationInfo info_real{ - [&info](int value) -> std::unique_ptr { - 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(std::move(allocator)); - else - return std::move(static_cast(&*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(onnxruntime::CreateAllocator(info_real)); + AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info) override { + return onnxruntime::CreateAllocator(info); } - std::unique_ptr CreateCPUAllocator( - const OrtMemoryInfo& memory_info) override { - return onnxruntime::make_unique( - onnxruntime::make_unique(memory_info)); + std::unique_ptr CreateCPUAllocator(const OrtMemoryInfo& memory_info) override { + return onnxruntime::make_unique(memory_info); }; std::unique_ptr Create_IExecutionProvider_Router( @@ -268,12 +216,12 @@ struct ProviderHostImpl : ProviderHost { }; #ifdef USE_TENSORRT - std::unique_ptr CreateCUDAAllocator(int16_t device_id, const char* name) override { - return onnxruntime::make_unique(onnxruntime::make_unique(device_id, name)); + std::unique_ptr CreateCUDAAllocator(int16_t device_id, const char* name) override { + return onnxruntime::make_unique(device_id, name); } - std::unique_ptr CreateCUDAPinnedAllocator(int16_t device_id, const char* name) override { - return onnxruntime::make_unique(onnxruntime::make_unique(device_id, name)); + std::unique_ptr CreateCUDAPinnedAllocator(int16_t device_id, const char* name) override { + return onnxruntime::make_unique(device_id, name); } std::unique_ptr CreateGPUDataTransfer() override { return onnxruntime::make_unique(); } @@ -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(); } diff --git a/onnxruntime/core/providers/dnnl/dnnl_execution_provider.cc b/onnxruntime/core/providers/dnnl/dnnl_execution_provider.cc index 67b24f5f20..da52d01447 100644 --- a/onnxruntime/core/providers/dnnl/dnnl_execution_provider.cc +++ b/onnxruntime/core/providers/dnnl/dnnl_execution_provider.cc @@ -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); diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv.h index 12c6a4f94d..1baf516a9c 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv.h +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv.h @@ -386,8 +386,7 @@ class DnnlConv : public DnnlKernel { if (filter_dst_mem == nullptr) { dnnl::memory src = dnnl::memory({{filter_dims_mkl}, DnnnType(), filter_format_}, cpu_engine, (void*)filter_data); - IAllocatorUniquePtr filter_reorder_buffer = - Provider_IAllocator::MakeUniquePtr(alloc_, filter_size_); + IAllocatorUniquePtr filter_reorder_buffer = IAllocator::MakeUniquePtr(alloc_, filter_size_); filter_dst_mem = onnxruntime::make_unique( 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(alloc_, src_size); + src_reorder_buffer_ = IAllocator::MakeUniquePtr(alloc_, src_size); src_mem_->set_data_handle(src_reorder_buffer_.get()); } else { if (mklnode_ptr_->parent_nodes.empty()) { diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv_batchnorm.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv_batchnorm.h index 8173855da5..9b1a8b2104 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv_batchnorm.h +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_conv_batchnorm.h @@ -395,8 +395,7 @@ class DnnlConvBatchNorm : public DnnlKernel { if (filter_dst_mem == nullptr) { dnnl::memory src = dnnl::memory({{filter_dims_mkl}, DnnnType(), filter_format_}, cpu_engine, (void*)weights_scaled_by_axis.data()); - IAllocatorUniquePtr filter_reorder_buffer = - Provider_IAllocator::MakeUniquePtr(alloc_, filter_size_); + IAllocatorUniquePtr filter_reorder_buffer = IAllocator::MakeUniquePtr(alloc_, filter_size_); filter_dst_mem = onnxruntime::make_unique( dnnl::memory(conv_fwd_pd_->weights_desc(), cpu_engine, filter_reorder_buffer.get())); @@ -410,8 +409,7 @@ class DnnlConvBatchNorm : public DnnlKernel { std::shared_ptr bias_mem = provider_->GetBiasMemoryBuffer(mklnode_ptr_->weight_name); if (bias_mem == nullptr) { auto bias_size = conv_fwd_pd_.get()->bias_desc().get_size(); - IAllocatorUniquePtr bias_buffer = - Provider_IAllocator::MakeUniquePtr(alloc_, bias_size); + IAllocatorUniquePtr bias_buffer = IAllocator::MakeUniquePtr(alloc_, bias_size); bias_mem = onnxruntime::make_unique( dnnl::memory(conv_fwd_pd_->bias_desc(), cpu_engine, bias_buffer.get())); float* bias_buffer_data = static_cast(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(alloc_, src_size); + src_reorder_buffer_ = IAllocator::MakeUniquePtr(alloc_, src_size); src_mem_->set_data_handle(src_reorder_buffer_.get()); } else { if (mklnode_ptr_->parent_nodes.empty()) { diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.h index 454c66b012..f215509b21 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.h +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_kernel.h @@ -111,7 +111,7 @@ class DnnlKernel { // memory used for reorders std::unique_ptr reorder_dst_mem_to_; - Provider_AllocatorPtr alloc_; + AllocatorPtr alloc_; DNNLExecutionProvider* provider_; }; diff --git a/onnxruntime/core/providers/dnnl/subgraph/dnnl_pool.h b/onnxruntime/core/providers/dnnl/subgraph/dnnl_pool.h index d7fc1a9c81..1e396eff93 100644 --- a/onnxruntime/core/providers/dnnl/subgraph/dnnl_pool.h +++ b/onnxruntime/core/providers/dnnl/subgraph/dnnl_pool.h @@ -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(alloc_, src_size); + src_reorder_buffer_ = IAllocator::MakeUniquePtr(alloc_, src_size); src_mem_->set_data_handle(src_reorder_buffer_.get()); } else { if (mklnode_ptr_->parent_nodes.empty()) { diff --git a/onnxruntime/core/providers/shared_library/provider_api.h b/onnxruntime/core/providers/shared_library/provider_api.h index fb4d0a8d91..bc14550b38 100644 --- a/onnxruntime/core/providers/shared_library/provider_api.h +++ b/onnxruntime/core/providers/shared_library/provider_api.h @@ -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 using IAllocatorUniquePtr = std::unique_ptr>; -std::unique_ptr Provider_CreateCPUAllocator(const OrtMemoryInfo& memory_info); -std::unique_ptr Provider_CreateCUDAAllocator(int16_t device_id, const char* name); -std::unique_ptr Provider_CreateCUDAPinnedAllocator(int16_t device_id, const char* name); -Provider_AllocatorPtr CreateAllocator(const Provider_AllocatorCreationInfo& info); +std::unique_ptr CreateCPUAllocator(const OrtMemoryInfo& memory_info); +std::unique_ptr CreateCUDAAllocator(int16_t device_id, const char* name); +std::unique_ptr CreateCUDAPinnedAllocator(int16_t device_id, const char* name); std::unique_ptr Provider_CreateGPUDataTransfer(); @@ -192,8 +183,6 @@ struct Category { static const char* onnxruntime; ///< General output }; -constexpr const char* SEVERITY_PREFIX = "VIWEF"; - } // namespace logging namespace math { diff --git a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc index c5ed986da0..fd296d3e1b 100644 --- a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc +++ b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc @@ -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_CreateCPUAllocator(const OrtMemoryInfo& info) { +std::unique_ptr 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_CreateCUDAAllocator(int16_t device_id, const char* name) { +std::unique_ptr CreateCUDAAllocator(int16_t device_id, const char* name) { return g_host->CreateCUDAAllocator(device_id, name); } -std::unique_ptr Provider_CreateCUDAPinnedAllocator(int16_t device_id, const char* name) { +std::unique_ptr 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(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; diff --git a/onnxruntime/core/providers/shared_library/provider_interfaces.h b/onnxruntime/core/providers/shared_library/provider_interfaces.h index 8ff9cebe73..09ac4d6ce1 100644 --- a/onnxruntime/core/providers/shared_library/provider_interfaces.h +++ b/onnxruntime/core/providers/shared_library/provider_interfaces.h @@ -74,62 +74,6 @@ struct Provider_IExecutionProviderFactory { class DataTypeImpl; using MLDataType = const DataTypeImpl*; -template -using Provider_IAllocatorUniquePtr = std::unique_ptr>; - -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 - static Provider_IAllocatorUniquePtr MakeUniquePtr(std::shared_ptr 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::value) { - // TODO: Use internal implementation to get correct sizes - return nullptr; - } - return Provider_IAllocatorUniquePtr{ - static_cast(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; -using Provider_AllocatorFactory = std::function(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; +// using Provider_NodeAttributes = std::unordered_map; using Provider_InitializedTensorSet = std::unordered_map; @@ -188,9 +132,9 @@ struct Provider_IExecutionProvider_Router { virtual std::vector> Provider_GetCapability(const onnxruntime::Provider_GraphViewer& graph, const std::vector& 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_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& fused_nodes, std::vector& 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 CreateCPUAllocator(const OrtMemoryInfo& memory_info) = 0; + virtual std::unique_ptr CreateCPUAllocator(const OrtMemoryInfo& memory_info) = 0; #ifdef USE_TENSORRT - virtual std::unique_ptr CreateCUDAAllocator(int16_t device_id, const char* name) = 0; - virtual std::unique_ptr CreateCUDAPinnedAllocator(int16_t device_id, const char* name) = 0; + virtual std::unique_ptr CreateCUDAAllocator(int16_t device_id, const char* name) = 0; + virtual std::unique_ptr CreateCUDAPinnedAllocator(int16_t device_id, const char* name) = 0; virtual std::unique_ptr 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 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; diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 1dc0d5f863..7763974622 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -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 { diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h index 481a543c47..70dd5b5d8f 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h @@ -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& fused_nodes, std::vector& 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 diff --git a/orttraining/orttraining/models/bert/main.cc b/orttraining/orttraining/models/bert/main.cc index a31e613178..bfa4bd12c4 100644 --- a/orttraining/orttraining/models/bert/main.cc +++ b/orttraining/orttraining/models/bert/main.cc @@ -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" diff --git a/orttraining/orttraining/models/gpt2/main.cc b/orttraining/orttraining/models/gpt2/main.cc index 4c418ce6fa..6e2277654a 100644 --- a/orttraining/orttraining/models/gpt2/main.cc +++ b/orttraining/orttraining/models/gpt2/main.cc @@ -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" diff --git a/orttraining/orttraining/models/mnist/main.cc b/orttraining/orttraining/models/mnist/main.cc index 70287acc77..9cf6a65e60 100644 --- a/orttraining/orttraining/models/mnist/main.cc +++ b/orttraining/orttraining/models/mnist/main.cc @@ -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(MPIContext::GetInstance().GetLocalRank()); params.providers.emplace(kCudaExecutionProvider, CreateExecutionProviderFactory_CUDA(device_id)); }