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

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";
 

include/onnxruntime/core/framework/execution_provider.h

@@ -4,7 +4,6 @@
 #pragma once
 
 #ifndef PROVIDER_BRIDGE_PROVIDER
-#include <map>
 #include <unordered_map>
 #include <unordered_set>
 
@@ -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<int, AllocatorPtr>;
+using AllocatorMap = std::unordered_map<int, AllocatorPtr>;
 using MemoryInfoSet = std::set<OrtMemoryInfo>;
 
 // 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<AllocatorManager> allocator_manager);
+
  private:
   const std::string type_;
   AllocatorMap allocators_;

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<OrtDevice::DeviceId>(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<OrtDevice::DeviceId>(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<OrtDevice::DeviceId>(id1)),
-        id1, mem_type1);
   } else {
     return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Specified device is not supported.");
   }

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<IAllocator>(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

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 <unordered_map>
 
 namespace onnxruntime {
 
 using AllocatorFactory = std::function<std::unique_ptr<IAllocator>(OrtDevice::DeviceId)>;
 
+using AllocatorMap = std::unordered_map<int, AllocatorPtr>;
+// TODO: update OrtMemoryInfo, use unordered_set instead
+using MemoryInfoSet = std::set<OrtMemoryInfo>;
+
+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<AllocatorPtr> allocator_list_;
+};
+
 }  // namespace onnxruntime

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<AllocatorManager> ) {
+  return;
+}
+
 #if !defined(ORT_MINIMAL_BUILD)
 common::Status IExecutionProvider::Compile(const std::vector<onnxruntime::Node*>& /*fused_node*/,
                                            std::vector<NodeComputeInfo>& /*node_compute_funcs*/) {

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<std::unique_ptr<ComputeCapability>> IExecutionProvider__GetCapability(const IExecutionProvider* p, const onnxruntime::GraphViewer& graph_viewer,
                                                                                     const std::vector<const KernelRegistry*>& kernel_registries) override { return p->IExecutionProvider::GetCapability(graph_viewer, kernel_registries); }
   common::Status IExecutionProvider__Compile(IExecutionProvider* p, const std::vector<onnxruntime::Node*>& fused_nodes, std::vector<NodeComputeInfo>& 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<AllocatorManager> 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_;
 

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<CUDAAllocator>(device_id, CUDA);
-      },
-      info_.device_id,
-      true,
-      {info_.cuda_mem_limit,
-       static_cast<int>(info_.arena_extend_strategy),
-       -1, -1});
-
-  InsertAllocator(CreateAllocator(default_memory_info));
-
-  AllocatorCreationInfo pinned_memory_info(
-      [](OrtDevice::DeviceId device_id) {
-        return onnxruntime::make_unique<CUDAPinnedAllocator>(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<CPUAllocator>(
-            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<OrtMutex> 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<AllocatorManager> 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<CUDAAllocator>(device_id, CUDA);
+        },
+        info_.device_id,
+        true,
+        {info_.cuda_mem_limit,
+         static_cast<int>(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<CUDAPinnedAllocator>(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<CPUAllocator>(
+              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

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<AllocatorManager> allocator_manager) override;
+
  private:
   CUDAExecutionProviderInfo info_;
   cudaDeviceProp device_prop_;

onnxruntime/core/providers/cuda/cuda_kernel.h

@@ -44,7 +44,7 @@ class CudaKernel : public OpKernel {
 
   template <typename T>
   inline IAllocatorUniquePtr<T> 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<T>(allocator, count_or_bytes);

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<ROCMPinnedAllocator>(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<OrtMutex> lock(deferred_release_cpu_ptr_mutex_);
     auto it = deferred_release_cpu_ptr_.begin();

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:

onnxruntime/core/providers/rocm/rocm_kernel.h

@@ -41,7 +41,7 @@ class RocmKernel : public OpKernel {
 
   template <typename T>
   inline IAllocatorUniquePtr<T> 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<T>(allocator, count_or_bytes);

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 {

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<float>() {
   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<std::unique_ptr<ComputeCapability>> IExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph_viewer,
                                                                                   const std::vector<const KernelRegistry*>& 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<AllocatorManager> allocator_manager) {
+  return g_host->IExecutionProvider__RegisterAllocator(this, allocator_manager);
+}
+
 #ifdef USE_TENSORRT
 std::unique_ptr<IAllocator> CreateCUDAAllocator(int16_t device_id, const char* name) {
   return g_host->CreateCUDAAllocator(device_id, name);

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<std::unique_ptr<ComputeCapability>> IExecutionProvider__GetCapability(const IExecutionProvider* p, const onnxruntime::GraphViewer& graph_viewer,
                                                                                             const std::vector<const KernelRegistry*>& kernel_registries) = 0;
   virtual common::Status IExecutionProvider__Compile(IExecutionProvider* p, const std::vector<onnxruntime::Node*>& fused_nodes, std::vector<NodeComputeInfo>& 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<AllocatorManager> 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;

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<AllocatorManager> 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<IDataTransfer> TensorrtExecutionProvider::GetDataTransfer() const {
   return onnxruntime::CreateGPUDataTransfer();
 }

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<AllocatorManager> allocator_manager) override;
+
  private:
   int max_partition_iterations_ = 1000;
   int min_subgraph_size_ = 1;

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<onnxruntime::AllocatorManager>();
 }
 
 InferenceSession::InferenceSession(const SessionOptions& session_options, const Environment& session_env)
@@ -397,6 +398,8 @@ common::Status InferenceSession::RegisterExecutionProvider(std::unique_ptr<IExec
 
   const std::string& provider_type = p_exec_provider->Type();
 
+  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) {

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<onnxruntime::AllocatorManager> 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<uint8_t> ort_format_model_bytes_;
+  
+  std::shared_ptr<onnxruntime::AllocatorManager> allocator_manager_;
 };
 
 struct SessionIOBinding {

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<int64_t> dims_mul_x = {1, 3, 2};
   std::vector<float> values_mul_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
   OrtValue ml_value_x;
-  CreateMLValue<float>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_x);
+  CreateMLValue<float>(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_x);
   OrtValue ml_value_y;
-  CreateMLValue<float>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_y);
+  CreateMLValue<float>(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_y);
   OrtValue ml_value_z;
-  CreateMLValue<float>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_z);
+  CreateMLValue<float>(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<int64_t> expected_dims_mul_m = {1, 3, 2};
   std::vector<float> 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<IExecutionProvider> 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<float>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_x);
-  CreateMLValue<float>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_y);
-  CreateMLValue<float>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_z);
+  CreateMLValue<float>(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_x);
+  CreateMLValue<float>(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_y);
+  CreateMLValue<float>(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<int64_t> dims_mul_x = {1, 3, 2};
   std::vector<float> values_mul_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
   OrtValue ml_value_x;
-  CreateMLValue<float>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_x);
+  CreateMLValue<float>(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_x);
   OrtValue ml_value_y;
-  CreateMLValue<float>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_y);
+  CreateMLValue<float>(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_y);
   OrtValue ml_value_z;
-  CreateMLValue<float>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_z);
+  CreateMLValue<float>(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<int64_t> expected_dims_mul_m = {1, 3, 2};
   std::vector<float> 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<IExecutionProvider> 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<int64_t> dims_mul_x = {1, 3, 2};
   std::vector<bool> values_mul_x = {true, false, true, false, true, false};
   std::vector<int64_t> dims_mul_y = {1, 3, 2};
   std::vector<float> values_mul_y = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
   OrtValue ml_value_x;
-  CreateMLValue<bool>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_x);
+  CreateMLValue<bool>(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_x);
   OrtValue ml_value_y;
-  CreateMLValue<float>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_y, values_mul_y, &ml_value_y);
+  CreateMLValue<float>(cpu_allocator, dims_mul_y, values_mul_y, &ml_value_y);
   OrtValue ml_value_z;
-  CreateMLValue<float>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_y, values_mul_y, &ml_value_z);
+  CreateMLValue<float>(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<int64_t> expected_dims_mul_n = {1, 3, 2};
   std::vector<float> 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<IExecutionProvider> execution_provider = DefaultTensorrtExecutionProvider();
   EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK());
 
@@ -394,12 +406,23 @@ TEST(TensorrtExecutionProviderTest, RemoveCycleTest) {
   std::vector<int64_t> dims_mul_z = {1, 3, 2};
   std::vector<bool> 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<bool>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_x, values_mul_x, &ml_value_x);
+  CreateMLValue<bool>(cpu_allocator, dims_mul_x, values_mul_x, &ml_value_x);
   OrtValue ml_value_y;
-  CreateMLValue<bool>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_y, values_mul_y, &ml_value_y);
+  CreateMLValue<bool>(cpu_allocator, dims_mul_y, values_mul_y, &ml_value_y);
   OrtValue ml_value_z;
-  CreateMLValue<bool>(TestCudaExecutionProvider()->GetAllocator(0, OrtMemTypeCPU), dims_mul_y, values_mul_y, &ml_value_z);
+  CreateMLValue<bool>(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<int64_t> expected_dims_mul_m = {1, 3, 2};
   std::vector<bool> 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<IExecutionProvider> execution_provider = DefaultTensorrtExecutionProvider();
   EXPECT_TRUE(session_object.RegisterExecutionProvider(std::move(execution_provider)).IsOK());
 

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