From b5d1a49b300e2e2b2182f90b21b0dd749663b337 Mon Sep 17 00:00:00 2001 From: Hector Li Date: Wed, 27 Jan 2021 00:14:43 -0800 Subject: [PATCH] Share allocator between CUDA EP & TRT EP. (#6332) * Share allocator between CUDA EP & TRT EP. limitation: 1. Does not cover the per-thread allocator created by CUDA EP, still need to figure out the way to remove it 2. Need to have more identifiers to make it able to share CPU allocator across all EPs --- .../onnxruntime/core/framework/allocator.h | 3 - .../core/framework/execution_provider.h | 15 ++- onnxruntime/core/framework/allocator.cc | 8 -- onnxruntime/core/framework/allocatormgr.cc | 37 +++++++ onnxruntime/core/framework/allocatormgr.h | 27 +++++ .../core/framework/execution_provider.cc | 14 +++ .../core/framework/provider_bridge_ort.cc | 11 ++- .../providers/cuda/cuda_execution_provider.cc | 90 ++++++++++------- .../providers/cuda/cuda_execution_provider.h | 4 +- onnxruntime/core/providers/cuda/cuda_kernel.h | 2 +- .../providers/rocm/rocm_execution_provider.cc | 6 +- .../providers/rocm/rocm_execution_provider.h | 2 - onnxruntime/core/providers/rocm/rocm_kernel.h | 2 +- .../providers/shared_library/provider_api.h | 3 + .../provider_bridge_provider.cc | 16 +++ .../shared_library/provider_interfaces.h | 6 ++ .../tensorrt/tensorrt_execution_provider.cc | 32 ++++-- .../tensorrt/tensorrt_execution_provider.h | 2 + onnxruntime/core/session/inference_session.cc | 3 + onnxruntime/core/session/inference_session.h | 7 ++ .../providers/tensorrt/tensorrt_basic_test.cc | 98 +++++++++++-------- onnxruntime/test/shared_lib/test_inference.cc | 4 - 22 files changed, 281 insertions(+), 111 deletions(-) diff --git a/include/onnxruntime/core/framework/allocator.h b/include/onnxruntime/core/framework/allocator.h index 1059c398f6..b66548fee6 100644 --- a/include/onnxruntime/core/framework/allocator.h +++ b/include/onnxruntime/core/framework/allocator.h @@ -22,9 +22,6 @@ namespace onnxruntime { constexpr const char* CPU = "Cpu"; constexpr const char* CUDA = "Cuda"; constexpr const char* CUDA_PINNED = "CudaPinned"; -// TODO: Unify the allocator for CUDA and Tensorrt -constexpr const char* TRT = "Tensorrt"; -constexpr const char* TRT_PINNED = "TensorrtPinned"; constexpr const char* MIGRAPHX = "MIGraphX"; constexpr const char* MIGRAPHX_PINNED = "MIGraphXPinned"; diff --git a/include/onnxruntime/core/framework/execution_provider.h b/include/onnxruntime/core/framework/execution_provider.h index 3c877f2bf0..6069ed4839 100644 --- a/include/onnxruntime/core/framework/execution_provider.h +++ b/include/onnxruntime/core/framework/execution_provider.h @@ -4,7 +4,6 @@ #pragma once #ifndef PROVIDER_BRIDGE_PROVIDER -#include #include #include @@ -25,13 +24,14 @@ class KernelRegistryManager; #include "core/framework/provider_options.h" #include "core/framework/func_api.h" +#include "core/framework/allocatormgr.h" namespace onnxruntime { /** Logical device representation. */ -using AllocatorMap = std::map; +using AllocatorMap = std::unordered_map; using MemoryInfoSet = std::set; // if we are export the fused function to dll, the function will still in the same binary as onnxruntime @@ -167,6 +167,8 @@ class IExecutionProvider { void InsertAllocator(AllocatorPtr allocator); void ReplaceAllocator(AllocatorPtr allocator); + // TODO: temparary sulotion, need to unify the interface in EP and AllocatorManager + void TryInsertAllocator(AllocatorPtr allocator); // creation of a fused node is not supported in a minimal build, so any EP enabled in that scenario must support // compilation via GraphViewer instances. @@ -248,10 +250,17 @@ class IExecutionProvider { @remarks e.g. the TensorRT Execution Provider is used in multiple sessions and the underlying infrastructure caches compiled kernels, so the name must be unique and deterministic across models and sessions. NOTE: Ideally this would be a protected method, but to work across the EP bridge it has to be public and - virtual, and ModelMetadefIdGenerator but be defined in the header as well. + virtual, and ModelMetadefIdGenerator but be defined in the header as well. */ virtual int GenerateMetaDefId(const onnxruntime::GraphViewer& graph_viewer, uint64_t& model_hash) const; + /** + Register allocators used for EP + TODO: Used for CUDA & TRT only for now, will have one more PR to apply this for all EPs. + EPs will have a shared pointer to allocator_manager, allocator_managerall will be the only place for allocators + */ + virtual void RegisterAllocator(std::shared_ptr allocator_manager); + private: const std::string type_; AllocatorMap allocators_; diff --git a/onnxruntime/core/framework/allocator.cc b/onnxruntime/core/framework/allocator.cc index 6403e38960..6967d57f8d 100644 --- a/onnxruntime/core/framework/allocator.cc +++ b/onnxruntime/core/framework/allocator.cc @@ -69,14 +69,6 @@ ORT_API_STATUS_IMPL(OrtApis::CreateMemoryInfo, _In_ const char* name1, enum OrtA *out = new OrtMemoryInfo( onnxruntime::CUDA_PINNED, type, OrtDevice(OrtDevice::CPU, OrtDevice::MemType::CUDA_PINNED, static_cast(id1)), id1, mem_type1); - } else if (strcmp(name1, onnxruntime::TRT) == 0) { - *out = new OrtMemoryInfo( - onnxruntime::TRT, type, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, static_cast(id1)), id1, - mem_type1); - } else if (strcmp(name1, onnxruntime::TRT_PINNED) == 0) { - *out = new OrtMemoryInfo( - onnxruntime::TRT_PINNED, type, OrtDevice(OrtDevice::CPU, OrtDevice::MemType::CUDA_PINNED, static_cast(id1)), - id1, mem_type1); } else { return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Specified device is not supported."); } diff --git a/onnxruntime/core/framework/allocatormgr.cc b/onnxruntime/core/framework/allocatormgr.cc index 00d1b98142..bffce15c96 100644 --- a/onnxruntime/core/framework/allocatormgr.cc +++ b/onnxruntime/core/framework/allocatormgr.cc @@ -13,6 +13,13 @@ namespace onnxruntime { using namespace common; +namespace { +//It assumes max(OrtMemType) <= 1, min(OrtMemType) = -2 +inline int MakeKey(int id, OrtMemType mem_type) { + return id << 2 | (mem_type + 2); +} +} // namespace + AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info) { auto device_allocator = std::unique_ptr(info.device_alloc_factory(info.device_id)); @@ -54,4 +61,34 @@ AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info) { return AllocatorPtr(std::move(device_allocator)); } + +// Update allocator in the provider if already present; ignore if not. +void AllocatorManager::ReplaceAllocator(AllocatorPtr allocator) { + const auto& info = allocator->Info(); + auto ite = mem_info_set_.find(info); + if (ite != mem_info_set_.end()) { + const int key = MakeKey(info.id, info.mem_type); + allocators_[key] = allocator; + } +} + +void AllocatorManager::InsertAllocator(AllocatorPtr allocator) { + const OrtMemoryInfo& info = allocator->Info(); + auto ite = mem_info_set_.find(info); + if (ite != mem_info_set_.end()) { + ORT_THROW("duplicated allocator"); + } + const int key = MakeKey(info.id, info.mem_type); + allocators_.insert({key, allocator}); + mem_info_set_.insert(ite, info); + allocator_list_.push_back(allocator); +} + +AllocatorPtr AllocatorManager::GetAllocator(int id, OrtMemType mem_type) const { + auto iter = allocators_.find(MakeKey(id, mem_type)); + if (iter != allocators_.end()) { + return iter->second; + } + return nullptr; +} } // namespace onnxruntime diff --git a/onnxruntime/core/framework/allocatormgr.h b/onnxruntime/core/framework/allocatormgr.h index 851fbba6f9..1fa54d4f7d 100644 --- a/onnxruntime/core/framework/allocatormgr.h +++ b/onnxruntime/core/framework/allocatormgr.h @@ -6,11 +6,18 @@ #include "core/common/common.h" #include "core/framework/allocator.h" #include "core/session/onnxruntime_c_api.h" +#include namespace onnxruntime { using AllocatorFactory = std::function(OrtDevice::DeviceId)>; +using AllocatorMap = std::unordered_map; +// TODO: update OrtMemoryInfo, use unordered_set instead +using MemoryInfoSet = std::set; + +const int DEFAULT_CPU_ALLOCATOR_DEVICE_ID = 0; + struct AllocatorCreationInfo { AllocatorCreationInfo(AllocatorFactory device_alloc_factory0, OrtDevice::DeviceId device_id0 = 0, @@ -33,4 +40,24 @@ struct AllocatorCreationInfo { // Valid values can be found in onnxruntime_c_api.h. AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info); +// TODO: Only used for TRT and CUDA EP currently, need to add more identifiers to use it across all EPs +class AllocatorManager { + // + public: + AllocatorManager() = default; + void InsertAllocator(AllocatorPtr allocator); + void ReplaceAllocator(AllocatorPtr allocator); + //Get an allocator with specified device id and MemType. Return nullptr if it doesn't exist + AllocatorPtr GetAllocator(int id, OrtMemType mem_type) const; + + private: + AllocatorMap allocators_; + // to ensure only allocators with unique OrtMemoryInfo are registered in the provider. + MemoryInfoSet mem_info_set_; + + // convenience list of the allocators so GetAllocatorList doesn't have to build a new vector each time + // contains the same instances as allocators_ + std::vector allocator_list_; +}; + } // namespace onnxruntime diff --git a/onnxruntime/core/framework/execution_provider.cc b/onnxruntime/core/framework/execution_provider.cc index 094f6756de..45d1297381 100644 --- a/onnxruntime/core/framework/execution_provider.cc +++ b/onnxruntime/core/framework/execution_provider.cc @@ -74,6 +74,20 @@ void IExecutionProvider::InsertAllocator(AllocatorPtr allocator) { allocator_list_.push_back(allocator); } +void IExecutionProvider::TryInsertAllocator(AllocatorPtr allocator) { + const OrtMemoryInfo& info = allocator->Info(); + auto ite = mem_info_set_.find(info); + if (ite != mem_info_set_.end()) { + LOGS_DEFAULT(WARNING) << "duplicated allocator: " << info.ToString(); + return; + } + InsertAllocator(allocator); +} + +void IExecutionProvider::RegisterAllocator(std::shared_ptr ) { + return; +} + #if !defined(ORT_MINIMAL_BUILD) common::Status IExecutionProvider::Compile(const std::vector& /*fused_node*/, std::vector& /*node_compute_funcs*/) { diff --git a/onnxruntime/core/framework/provider_bridge_ort.cc b/onnxruntime/core/framework/provider_bridge_ort.cc index 9e8e55182f..5e6a1f85eb 100644 --- a/onnxruntime/core/framework/provider_bridge_ort.cc +++ b/onnxruntime/core/framework/provider_bridge_ort.cc @@ -203,6 +203,7 @@ struct ProviderHostImpl : ProviderHost { // IExecutionProvider AllocatorPtr IExecutionProvider__GetAllocator(const IExecutionProvider* p, int id, OrtMemType mem_type) override { return p->IExecutionProvider::GetAllocator(id, mem_type); } void IExecutionProvider__InsertAllocator(IExecutionProvider* p, AllocatorPtr allocator) override { return p->IExecutionProvider::InsertAllocator(allocator); } + void IExecutionProvider__TryInsertAllocator(IExecutionProvider* p, AllocatorPtr allocator) override { return p->IExecutionProvider::TryInsertAllocator(allocator); } std::vector> IExecutionProvider__GetCapability(const IExecutionProvider* p, const onnxruntime::GraphViewer& graph_viewer, const std::vector& kernel_registries) override { return p->IExecutionProvider::GetCapability(graph_viewer, kernel_registries); } common::Status IExecutionProvider__Compile(IExecutionProvider* p, const std::vector& fused_nodes, std::vector& node_compute_funcs) override { @@ -221,6 +222,10 @@ struct ProviderHostImpl : ProviderHost { return p->IExecutionProvider::GenerateMetaDefId(graph_viewer, model_hash); } + void IExecutionProvider__RegisterAllocator(IExecutionProvider* p, std::shared_ptr allocator_manager) override { + return p->IExecutionProvider::RegisterAllocator(allocator_manager); + } + // Status std::string Status__ToString(const Status* p) override { return p->ToString(); } @@ -563,7 +568,11 @@ struct ProviderHostImpl : ProviderHost { const TensorShape& Tensor__Shape(const Tensor* p) override { return p->Shape(); } size_t Tensor__SizeInBytes(const Tensor* p) override { return p->SizeInBytes(); } - const OrtMemoryInfo& Tensor__Location(const Tensor* p) override { return p->Location(); } + const OrtMemoryInfo& Tensor__Location(const Tensor* p) override { return p->Location(); } + + // AllocatorManager + void AllocatorManager__InsertAllocator(AllocatorManager* p, AllocatorPtr allocator) override { p->InsertAllocator(allocator); } + AllocatorPtr AllocatorManager__GetAllocator(AllocatorManager* p, int id, OrtMemType mem_type) override { return p->GetAllocator(id, mem_type); }; } provider_host_; diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index 01f1ed9a2e..17cd8c1336 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -107,42 +107,10 @@ CUDAExecutionProvider::CUDAExecutionProvider(const CUDAExecutionProviderInfo& in size_t free = 0; size_t total = 0; CUDA_CALL_THROW(cudaMemGetInfo(&free, &total)); - - AllocatorCreationInfo default_memory_info( - [](OrtDevice::DeviceId device_id) { - return onnxruntime::make_unique(device_id, CUDA); - }, - info_.device_id, - true, - {info_.cuda_mem_limit, - static_cast(info_.arena_extend_strategy), - -1, -1}); - - InsertAllocator(CreateAllocator(default_memory_info)); - - AllocatorCreationInfo pinned_memory_info( - [](OrtDevice::DeviceId device_id) { - return onnxruntime::make_unique(device_id, CUDA_PINNED); - }, - CPU_ALLOCATOR_DEVICE_ID); - - InsertAllocator(CreateAllocator(pinned_memory_info)); - - // TODO: this is actually used for the cuda kernels which explicitly ask for inputs from CPU. - // This will be refactored/removed when allocator and execution provider are decoupled. - AllocatorCreationInfo cpu_memory_info( - [](int device_id) { - return onnxruntime::make_unique( - OrtMemoryInfo("CUDA_CPU", OrtAllocatorType::OrtDeviceAllocator, OrtDevice(), device_id, - OrtMemTypeCPUInput)); - }, - CPU_ALLOCATOR_DEVICE_ID); - - InsertAllocator(CreateAllocator(cpu_memory_info)); } CUDAExecutionProvider::~CUDAExecutionProvider() { - auto cpu_alloc = GetAllocator(CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPU); + auto cpu_alloc = GetAllocator(DEFAULT_CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPU); { std::lock_guard lock(deferred_release_cpu_ptr_mutex_); auto it = deferred_release_cpu_ptr_.begin(); @@ -1991,4 +1959,60 @@ CUDAExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph, return result; } +void CUDAExecutionProvider::RegisterAllocator(std::shared_ptr allocator_manager) { + // Try to get a CUDA allocator from allocator manager first + // Used to allocate CUDA device memory + auto cuda_alloc = allocator_manager->GetAllocator(info_.device_id, OrtMemTypeDefault); + if (nullptr == cuda_alloc) { + AllocatorCreationInfo default_memory_info( + [](OrtDevice::DeviceId device_id) { + return onnxruntime::make_unique(device_id, CUDA); + }, + info_.device_id, + true, + {info_.cuda_mem_limit, + static_cast(info_.arena_extend_strategy), + -1, -1}); + cuda_alloc = CreateAllocator(default_memory_info); + allocator_manager->InsertAllocator(cuda_alloc); + } + TryInsertAllocator(cuda_alloc); + + // OrtMemTypeCPUOutput -- allocated by cudaMallocHost, used to copy CUDA device memory to CPU + // Use pinned memory instead of pageable memory make the data transfer faster + // Used by node MemcpyToHost only + auto cuda_pinned_alloc = allocator_manager->GetAllocator(DEFAULT_CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPUOutput); + if (nullptr == cuda_pinned_alloc) { + AllocatorCreationInfo pinned_memory_info( + [](OrtDevice::DeviceId device_id) { + return onnxruntime::make_unique(device_id, CUDA_PINNED); + }, + DEFAULT_CPU_ALLOCATOR_DEVICE_ID); + + cuda_pinned_alloc = CreateAllocator(pinned_memory_info); + allocator_manager->InsertAllocator(cuda_pinned_alloc); + } + TryInsertAllocator(cuda_pinned_alloc); + + // OrtMemTypeCPUInput -- CUDA op place the input on CPU and will not be accessed by CUDA kernel, no sync issue + auto cuda_cpu_alloc = allocator_manager->GetAllocator(DEFAULT_CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPUInput); + if (nullptr == cuda_cpu_alloc) { + // TODO: this is actually used for the cuda kernels which explicitly ask for inputs from CPU. + // This will be refactored/removed when allocator and execution provider are decoupled. + // Need to move the OrtMemoryType out of Allocator, that's one thing blocking us to share it with CPU EP + // CPUAllocator is OrtMemTypeDefault for CPU EP + AllocatorCreationInfo cpu_memory_info( + [](int device_id) { + return onnxruntime::make_unique( + OrtMemoryInfo("CUDA_CPU", OrtAllocatorType::OrtDeviceAllocator, OrtDevice(), device_id, + OrtMemTypeCPUInput)); + }, + DEFAULT_CPU_ALLOCATOR_DEVICE_ID); + + cuda_cpu_alloc = CreateAllocator(cpu_memory_info); + allocator_manager->InsertAllocator(cuda_cpu_alloc); + } + TryInsertAllocator(cuda_cpu_alloc); +} + } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.h b/onnxruntime/core/providers/cuda/cuda_execution_provider.h index d70ca20be2..2f55ee6bbb 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.h +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.h @@ -18,8 +18,6 @@ namespace onnxruntime { -const int CPU_ALLOCATOR_DEVICE_ID = 0; - // Logical device representation. class CUDAExecutionProvider : public IExecutionProvider { public: @@ -77,6 +75,8 @@ class CUDAExecutionProvider : public IExecutionProvider { return CUDAExecutionProviderInfo::ToProviderOptions(info_); } + void RegisterAllocator(std::shared_ptr allocator_manager) override; + private: CUDAExecutionProviderInfo info_; cudaDeviceProp device_prop_; diff --git a/onnxruntime/core/providers/cuda/cuda_kernel.h b/onnxruntime/core/providers/cuda/cuda_kernel.h index b1351a0964..12fccc9f19 100644 --- a/onnxruntime/core/providers/cuda/cuda_kernel.h +++ b/onnxruntime/core/providers/cuda/cuda_kernel.h @@ -44,7 +44,7 @@ class CudaKernel : public OpKernel { template inline IAllocatorUniquePtr AllocateBufferOnCPUPinned(size_t count_or_bytes) const { - AllocatorPtr allocator = provider_->GetAllocator(CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPU); + AllocatorPtr allocator = provider_->GetAllocator(DEFAULT_CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPU); if (!allocator) return nullptr; return IAllocator::MakeUniquePtr(allocator, count_or_bytes); diff --git a/onnxruntime/core/providers/rocm/rocm_execution_provider.cc b/onnxruntime/core/providers/rocm/rocm_execution_provider.cc index 27efe52a45..d5e0ee42b2 100644 --- a/onnxruntime/core/providers/rocm/rocm_execution_provider.cc +++ b/onnxruntime/core/providers/rocm/rocm_execution_provider.cc @@ -124,7 +124,7 @@ ROCMExecutionProvider::ROCMExecutionProvider(const ROCMExecutionProviderInfo& in [](OrtDevice::DeviceId device_id) { return onnxruntime::make_unique(device_id, CUDA_PINNED); }, - CPU_ALLOCATOR_DEVICE_ID); + DEFAULT_CPU_ALLOCATOR_DEVICE_ID); InsertAllocator(CreateAllocator(pinned_memory_info)); @@ -136,13 +136,13 @@ ROCMExecutionProvider::ROCMExecutionProvider(const ROCMExecutionProviderInfo& in OrtMemoryInfo("HIP_CPU", OrtAllocatorType::OrtDeviceAllocator, OrtDevice(), device_id, OrtMemTypeCPUInput)); }, - CPU_ALLOCATOR_DEVICE_ID); + DEFAULT_CPU_ALLOCATOR_DEVICE_ID); InsertAllocator(CreateAllocator(cpu_memory_info)); } ROCMExecutionProvider::~ROCMExecutionProvider() { - auto cpu_alloc = GetAllocator(CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPU); + auto cpu_alloc = GetAllocator(DEFAULT_CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPU); { std::lock_guard lock(deferred_release_cpu_ptr_mutex_); auto it = deferred_release_cpu_ptr_.begin(); diff --git a/onnxruntime/core/providers/rocm/rocm_execution_provider.h b/onnxruntime/core/providers/rocm/rocm_execution_provider.h index 9352c3144f..8859ec42f3 100644 --- a/onnxruntime/core/providers/rocm/rocm_execution_provider.h +++ b/onnxruntime/core/providers/rocm/rocm_execution_provider.h @@ -18,8 +18,6 @@ namespace onnxruntime { -const int CPU_ALLOCATOR_DEVICE_ID = 0; - // Logical device representation. class ROCMExecutionProvider : public IExecutionProvider { public: diff --git a/onnxruntime/core/providers/rocm/rocm_kernel.h b/onnxruntime/core/providers/rocm/rocm_kernel.h index 1e6ad8850b..f70d7136f4 100644 --- a/onnxruntime/core/providers/rocm/rocm_kernel.h +++ b/onnxruntime/core/providers/rocm/rocm_kernel.h @@ -41,7 +41,7 @@ class RocmKernel : public OpKernel { template inline IAllocatorUniquePtr AllocateBufferOnCPUPinned(size_t count_or_bytes) const { - AllocatorPtr allocator = provider_->GetAllocator(CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPU); + AllocatorPtr allocator = provider_->GetAllocator(DEFAULT_CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPU); if (!allocator) return nullptr; return IAllocator::MakeUniquePtr(allocator, count_or_bytes); diff --git a/onnxruntime/core/providers/shared_library/provider_api.h b/onnxruntime/core/providers/shared_library/provider_api.h index c32e171e83..306c5f945e 100644 --- a/onnxruntime/core/providers/shared_library/provider_api.h +++ b/onnxruntime/core/providers/shared_library/provider_api.h @@ -210,6 +210,9 @@ std::string GetEnvironmentVar(const std::string& var_name); inline AutoPadType StringToAutoPadType(const std::string& str) { return g_host->StringToAutoPadType(str); } +void AllocatorManager__InsertAllocator(AllocatorManager* p, AllocatorPtr allocator); +AllocatorPtr AllocatorManager__GetAllocator(AllocatorManager* p, int id, OrtMemType mem_type); + namespace logging { struct Category { diff --git a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc index f8fd2071c5..aa8540f5b9 100644 --- a/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc +++ b/onnxruntime/core/providers/shared_library/provider_bridge_provider.cc @@ -45,6 +45,14 @@ AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info) { return g_host->CreateAllocator(info); } +void AllocatorManager__InsertAllocator(AllocatorManager* p, AllocatorPtr allocator) { + return g_host->AllocatorManager__InsertAllocator(p, allocator); +} + +AllocatorPtr AllocatorManager__GetAllocator(AllocatorManager* p, int id, OrtMemType mem_type) { + return g_host->AllocatorManager__GetAllocator(p, id, mem_type); +} + template <> MLDataType DataTypeImpl::GetType() { return g_host->DataTypeImpl_GetType_float(); @@ -114,6 +122,10 @@ void IExecutionProvider::InsertAllocator(AllocatorPtr allocator) { g_host->IExecutionProvider__InsertAllocator(this, allocator); } +void IExecutionProvider::TryInsertAllocator(AllocatorPtr allocator) { + g_host->IExecutionProvider__TryInsertAllocator(this, allocator); +} + std::vector> IExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_viewer, const std::vector& kernel_registries) const { return g_host->IExecutionProvider__GetCapability(this, graph_viewer, kernel_registries); @@ -138,6 +150,10 @@ int IExecutionProvider::GenerateMetaDefId(const onnxruntime::GraphViewer& graph_ return g_host->IExecutionProvider__GenerateMetaDefId(this, graph_viewer, model_hash); } +void IExecutionProvider::RegisterAllocator(std::shared_ptr allocator_manager) { + return g_host->IExecutionProvider__RegisterAllocator(this, allocator_manager); +} + #ifdef USE_TENSORRT std::unique_ptr CreateCUDAAllocator(int16_t device_id, const char* name) { return g_host->CreateCUDAAllocator(device_id, name); diff --git a/onnxruntime/core/providers/shared_library/provider_interfaces.h b/onnxruntime/core/providers/shared_library/provider_interfaces.h index 49bc2d8d11..5babda61ce 100644 --- a/onnxruntime/core/providers/shared_library/provider_interfaces.h +++ b/onnxruntime/core/providers/shared_library/provider_interfaces.h @@ -159,6 +159,7 @@ struct ProviderHost { // IExecutionProvider virtual AllocatorPtr IExecutionProvider__GetAllocator(const IExecutionProvider* p, int id, OrtMemType mem_type) = 0; virtual void IExecutionProvider__InsertAllocator(IExecutionProvider* p, AllocatorPtr allocator) = 0; + virtual void IExecutionProvider__TryInsertAllocator(IExecutionProvider* p, AllocatorPtr allocator) = 0; virtual std::vector> IExecutionProvider__GetCapability(const IExecutionProvider* p, const onnxruntime::GraphViewer& graph_viewer, const std::vector& kernel_registries) = 0; virtual common::Status IExecutionProvider__Compile(IExecutionProvider* p, const std::vector& fused_nodes, std::vector& node_compute_funcs) = 0; @@ -167,6 +168,7 @@ struct ProviderHost { virtual int IExecutionProvider__GenerateMetaDefId(const IExecutionProvider* p, const onnxruntime::GraphViewer& graph_viewer, uint64_t& model_hash) = 0; + virtual void IExecutionProvider__RegisterAllocator(IExecutionProvider* p, std::shared_ptr allocator_manager) = 0; // Status virtual std::string Status__ToString(const Status* p) = 0; @@ -483,6 +485,10 @@ struct ProviderHost { virtual const TensorShape& Tensor__Shape(const Tensor* p) = 0; virtual size_t Tensor__SizeInBytes(const Tensor* p) = 0; virtual const OrtMemoryInfo& Tensor__Location(const Tensor* p) = 0; + + // AllocatorManager + virtual void AllocatorManager__InsertAllocator(AllocatorManager* p, AllocatorPtr allocator) = 0; + virtual AllocatorPtr AllocatorManager__GetAllocator(AllocatorManager* p, int id, OrtMemType mem_type) = 0; }; extern ProviderHost* g_host; diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 48a509c68c..1b420ccd66 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -374,15 +374,6 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv : IExecutionProvider{onnxruntime::kTensorrtExecutionProvider, true}, device_id_(info.device_id) { CUDA_CALL_THROW(cudaSetDevice(device_id_)); - AllocatorCreationInfo default_memory_info( - [](int id) { return CreateCUDAAllocator(id, onnxruntime::TRT); }, device_id_); - allocator_ = CreateAllocator(default_memory_info); - InsertAllocator(allocator_); - - AllocatorCreationInfo pinned_allocator_info( - [](int) { return CreateCUDAPinnedAllocator(0, onnxruntime::TRT_PINNED); }, device_id_); - InsertAllocator(CreateAllocator(pinned_allocator_info)); - // Get environment variables const std::string max_partition_iterations_env = onnxruntime::GetEnvironmentVar(tensorrt_env_vars::kMaxPartitionIterations); if (!max_partition_iterations_env.empty()) { @@ -457,6 +448,29 @@ AllocatorPtr TensorrtExecutionProvider::GetAllocator(int id, OrtMemType mem_type } } +void TensorrtExecutionProvider::RegisterAllocator(std::shared_ptr allocator_manager) { + allocator_ = AllocatorManager__GetAllocator(allocator_manager.get(), device_id_, OrtMemTypeDefault); + if (nullptr == allocator_) { + AllocatorCreationInfo default_memory_info( + [](OrtDevice::DeviceId device_id) { return CreateCUDAAllocator(device_id, onnxruntime::CUDA); }, device_id_); + allocator_ = CreateAllocator(default_memory_info); + AllocatorManager__InsertAllocator(allocator_manager.get(), allocator_); + } + TryInsertAllocator(allocator_); + + auto cuda_pinned_alloc = AllocatorManager__GetAllocator(allocator_manager.get(), DEFAULT_CPU_ALLOCATOR_DEVICE_ID, OrtMemTypeCPUOutput); + if (nullptr == cuda_pinned_alloc) { + AllocatorCreationInfo pinned_allocator_info( + [](OrtDevice::DeviceId device_id) { + return CreateCUDAPinnedAllocator(device_id, onnxruntime::CUDA_PINNED); + }, + DEFAULT_CPU_ALLOCATOR_DEVICE_ID); + cuda_pinned_alloc = CreateAllocator(pinned_allocator_info); + AllocatorManager__InsertAllocator(allocator_manager.get(), cuda_pinned_alloc); + } + TryInsertAllocator(cuda_pinned_alloc); +} + std::unique_ptr TensorrtExecutionProvider::GetDataTransfer() const { return onnxruntime::CreateGPUDataTransfer(); } diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h index 87dc5f3b59..7ca9e8c1fe 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h @@ -114,6 +114,8 @@ class TensorrtExecutionProvider : public IExecutionProvider { AllocatorPtr GetAllocator(int id, OrtMemType mem_type) const override; + void RegisterAllocator(std::shared_ptr allocator_manager) override; + private: int max_partition_iterations_ = 1000; int min_subgraph_size_ = 1; diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 8cd255676d..aa89b84c40 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -275,6 +275,7 @@ void InferenceSession::ConstructorCommon(const SessionOptions& session_options, telemetry_ = {}; // a monotonically increasing session id for use in telemetry session_id_ = global_session_id_.fetch_add(1); + allocator_manager_ = std::make_shared(); } InferenceSession::InferenceSession(const SessionOptions& session_options, const Environment& session_env) @@ -397,6 +398,8 @@ common::Status InferenceSession::RegisterExecutionProvider(std::unique_ptrType(); + p_exec_provider->RegisterAllocator(allocator_manager_); + // Some session option values (default or user provided) may not work with some EPs. // Rather than put the onus on the user to know these, make the appropriate change while logging the change. if (provider_type == onnxruntime::kDmlExecutionProvider) { diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h index 50727ca7e1..dbba45d1b8 100644 --- a/onnxruntime/core/session/inference_session.h +++ b/onnxruntime/core/session/inference_session.h @@ -20,6 +20,7 @@ #include "core/optimizer/graph_transformer_mgr.h" #include "core/optimizer/insert_cast_transformer.h" #include "core/framework/session_options.h" +#include "core/framework/allocatormgr.h" #ifdef ENABLE_LANGUAGE_INTEROP_OPS #include "core/language_interop_ops/language_interop_ops.h" #endif @@ -384,6 +385,10 @@ class InferenceSession { */ AllocatorPtr GetAllocator(const OrtMemoryInfo& mem_info) const; + std::shared_ptr GetAllocatorManager() { + return allocator_manager_; + } + /** *Get InferenceSession logger. */ @@ -649,6 +654,8 @@ class InferenceSession { // Longer term we may want to directly refer to offsets in this buffer for initializers so we don't need to copy // those into new OrtValue instances, at which point we won't free them until the InferenceSession goes away. std::vector ort_format_model_bytes_; + + std::shared_ptr allocator_manager_; }; struct SessionIOBinding { diff --git a/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc b/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc index 4f7f5a4d1d..b605428cab 100644 --- a/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc +++ b/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc @@ -62,6 +62,15 @@ TEST(TensorrtExecutionProviderTest, EngineCachingTest) { std::string model_file_name = "trt_execution_provider_enginecaching_test.onnx"; status = onnxruntime::Model::Save(model, model_file_name); + SessionOptions so; + so.session_logid = "TensorrtExecutionProviderTest.EngineCachingTest"; + RunOptions run_options; + run_options.run_tag = so.session_logid; + InferenceSession session_object{so, GetEnvironment()}; + auto allocator_manager = session_object.GetAllocatorManager(); + auto cuda_provider = TestCudaExecutionProvider(); + cuda_provider->RegisterAllocator(allocator_manager); + auto cpu_allocator = cuda_provider->GetAllocator(0, OrtMemTypeCPU); // First run with input shape {1, 3, 2} // TRT engine and profile will be created and cached // Data in profile, @@ -71,11 +80,11 @@ TEST(TensorrtExecutionProviderTest, EngineCachingTest) { std::vector dims_mul_x = {1, 3, 2}; std::vector values_mul_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; OrtValue ml_value_x; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_x); + CreateMLValue(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_x); OrtValue ml_value_y; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_y); + CreateMLValue(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_y); OrtValue ml_value_z; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_z); + CreateMLValue(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_z); NameMLValMap feeds; feeds.insert(std::make_pair("X", ml_value_x)); feeds.insert(std::make_pair("Y", ml_value_y)); @@ -90,11 +99,6 @@ TEST(TensorrtExecutionProviderTest, EngineCachingTest) { std::vector expected_dims_mul_m = {1, 3, 2}; std::vector expected_values_mul_m = {3.0f, 6.0f, 9.0f, 12.0f, 15.0f, 18.0f}; - SessionOptions so; - so.session_logid = "TensorrtExecutionProviderTest.EngineCachingTest"; - RunOptions run_options; - run_options.run_tag = so.session_logid; - InferenceSession session_object{so, GetEnvironment()}; std::unique_ptr execution_provider = DefaultTensorrtExecutionProvider(); EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK()); status = session_object.Load(model_file_name); @@ -114,9 +118,9 @@ TEST(TensorrtExecutionProviderTest, EngineCachingTest) { // Y: 1, 1, 3, 2, 2, 6 // Z: 1, 1, 3, 2, 2, 6 dims_mul_x = {1, 1, 6}; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_x); - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_y); - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_z); + CreateMLValue(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_x); + CreateMLValue(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_y); + CreateMLValue(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_z); feeds.clear(); feeds.insert(std::make_pair("X", ml_value_x)); feeds.insert(std::make_pair("Y", ml_value_y)); @@ -169,14 +173,25 @@ TEST(TensorrtExecutionProviderTest, FunctionTest) { std::string model_file_name = "trt_execution_provider_function_test.onnx"; status = onnxruntime::Model::Save(model, model_file_name); + SessionOptions so; + so.session_logid = "TensorrtExecutionProviderTest.FunctionTest"; + RunOptions run_options; + run_options.run_tag = so.session_logid; + InferenceSession session_object{so, GetEnvironment()}; + + auto allocator_manager = session_object.GetAllocatorManager(); + auto cuda_provider = TestCudaExecutionProvider(); + cuda_provider->RegisterAllocator(allocator_manager); + auto cpu_allocator = cuda_provider->GetAllocator(0, OrtMemTypeCPU); + std::vector dims_mul_x = {1, 3, 2}; std::vector values_mul_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; OrtValue ml_value_x; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_x); + CreateMLValue(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_x); OrtValue ml_value_y; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_y); + CreateMLValue(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_y); OrtValue ml_value_z; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_z); + CreateMLValue(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_z); NameMLValMap feeds; feeds.insert(std::make_pair("X", ml_value_x)); feeds.insert(std::make_pair("Y", ml_value_y)); @@ -191,13 +206,6 @@ TEST(TensorrtExecutionProviderTest, FunctionTest) { std::vector expected_dims_mul_m = {1, 3, 2}; std::vector expected_values_mul_m = {3.0f, 6.0f, 9.0f, 12.0f, 15.0f, 18.0f}; - SessionOptions so; - so.session_logid = "TensorrtExecutionProviderTest.FunctionTest"; - RunOptions run_options; - run_options.run_tag = so.session_logid; - - InferenceSession session_object{so, GetEnvironment()}; - std::unique_ptr execution_provider = DefaultTensorrtExecutionProvider(); EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK()); @@ -277,16 +285,27 @@ TEST(TensorrtExecutionProviderTest, NodeIndexMappingTest) { std::string model_file_name = "trt_execution_provider_nodeindexmapping_test.onnx"; status = onnxruntime::Model::Save(model, model_file_name); + SessionOptions so; + so.session_logid = "TensorrtExecutionProviderTest.NodeIndexMappingTest"; + RunOptions run_options; + run_options.run_tag = so.session_logid; + InferenceSession session_object{so, GetEnvironment()}; + + auto allocator_manager = session_object.GetAllocatorManager(); + auto cuda_provider = TestCudaExecutionProvider(); + cuda_provider->RegisterAllocator(allocator_manager); + auto cpu_allocator = cuda_provider->GetAllocator(0, OrtMemTypeCPU); + std::vector dims_mul_x = {1, 3, 2}; std::vector values_mul_x = {true, false, true, false, true, false}; std::vector dims_mul_y = {1, 3, 2}; std::vector values_mul_y = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; OrtValue ml_value_x; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_x); + CreateMLValue(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_x); OrtValue ml_value_y; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_y, values_mul_y, &ml_value_y); + CreateMLValue(cpu_allocator, dims_mul_y, values_mul_y, &ml_value_y); OrtValue ml_value_z; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_y, values_mul_y, &ml_value_z); + CreateMLValue(cpu_allocator, dims_mul_y, values_mul_y, &ml_value_z); NameMLValMap feeds; feeds.insert(std::make_pair("X", ml_value_x)); feeds.insert(std::make_pair("Y", ml_value_y)); @@ -304,13 +323,6 @@ TEST(TensorrtExecutionProviderTest, NodeIndexMappingTest) { std::vector expected_dims_mul_n = {1, 3, 2}; std::vector expected_values_mul_n = {0, 0, 0, 0, 0, 0}; - SessionOptions so; - so.session_logid = "TensorrtExecutionProviderTest.NodeIndexMappingTest"; - RunOptions run_options; - run_options.run_tag = so.session_logid; - - InferenceSession session_object{so, GetEnvironment()}; - std::unique_ptr execution_provider = DefaultTensorrtExecutionProvider(); EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK()); @@ -394,12 +406,23 @@ TEST(TensorrtExecutionProviderTest, RemoveCycleTest) { std::vector dims_mul_z = {1, 3, 2}; std::vector values_mul_z = {true, false, true, false, true, false}; + SessionOptions so; + so.session_logid = "TensorrtExecutionProviderTest.RemoveCycleTest"; + RunOptions run_options; + run_options.run_tag = so.session_logid; + InferenceSession session_object{so, GetEnvironment()}; + + auto allocator_manager = session_object.GetAllocatorManager(); + auto cuda_provider = TestCudaExecutionProvider(); + cuda_provider->RegisterAllocator(allocator_manager); + auto cpu_allocator = cuda_provider->GetAllocator(0, OrtMemTypeCPU); + OrtValue ml_value_x; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_x); + CreateMLValue(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_x); OrtValue ml_value_y; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_y, values_mul_y, &ml_value_y); + CreateMLValue(cpu_allocator, dims_mul_y, values_mul_y, &ml_value_y); OrtValue ml_value_z; - CreateMLValue(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_y, values_mul_y, &ml_value_z); + CreateMLValue(cpu_allocator, dims_mul_y, values_mul_y, &ml_value_z); NameMLValMap feeds; feeds.insert(std::make_pair("X", ml_value_x)); feeds.insert(std::make_pair("Y", ml_value_y)); @@ -414,13 +437,6 @@ TEST(TensorrtExecutionProviderTest, RemoveCycleTest) { std::vector expected_dims_mul_m = {1, 3, 2}; std::vector expected_values_mul_m = {false, false, false, false, false, true}; - SessionOptions so; - so.session_logid = "TensorrtExecutionProviderTest.RemoveCycleTest"; - RunOptions run_options; - run_options.run_tag = so.session_logid; - - InferenceSession session_object{so, GetEnvironment()}; - std::unique_ptr execution_provider = DefaultTensorrtExecutionProvider(); EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK()); diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index bcffcda1e5..5a85b07432 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -708,11 +708,7 @@ TEST(CApiTest, io_binding_cuda) { #endif Ort::Session session(*ort_env, MODEL_URI, session_options); -#ifdef USE_TENSORRT - Ort::MemoryInfo info_cuda("Tensorrt", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemTypeDefault); -#else Ort::MemoryInfo info_cuda("Cuda", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemTypeDefault); -#endif Ort::Allocator cuda_allocator(session, info_cuda); auto allocator_info = cuda_allocator.GetInfo();