update memory info and data transfer in TRT EP's factor to accommodate mutiple GPU devices

Author: chilo-ms

plugin_execution_providers/tensorrt/cuda_allocator.h

@@ -13,22 +13,13 @@ constexpr const char* CUDA_PINNED_ALLOCATOR = "CudaPinned";
 using DeviceId = int16_t;
 
 struct CUDAAllocator : OrtAllocator {
-  CUDAAllocator(const OrtMemoryInfo* mem_info, const char* name = CUDA_ALLOCATOR) {
+  CUDAAllocator(const OrtMemoryInfo* mem_info, DeviceId device_id) : mem_info_(mem_info), device_id_(device_id) {
     OrtAllocator::version = ORT_API_VERSION;
-    OrtAllocator::Alloc = [](OrtAllocator* this_, size_t size) { return static_cast<CUDAAllocator*>(this_)->Alloc(size); };
+    OrtAllocator::Alloc = [](OrtAllocator* this_, size_t size) {
+      return static_cast<CUDAAllocator*>(this_)->Alloc(size);
+    };
     OrtAllocator::Free = [](OrtAllocator* this_, void* p) { static_cast<CUDAAllocator*>(this_)->Free(p); };
     OrtAllocator::Info = [](const OrtAllocator* this_) { return static_cast<const CUDAAllocator*>(this_)->Info(); };
-
-    mem_info_ = mem_info;
-
-    device_id_ = device_id;
-
-    const OrtApi* api = OrtGetApiBase()->GetApi(ORT_API_VERSION);
-    api->CreateMemoryInfo(name,
-                          OrtAllocatorType::OrtDeviceAllocator,
-                          static_cast<int>(device_id),
-                          OrtMemType::OrtMemTypeDefault,
-                          &mem_info_);
   }
   // TODO: Handle destructor
   //~CUDAAllocator();

plugin_execution_providers/tensorrt/tensorrt_execution_provider.cc

@@ -654,7 +654,7 @@ OrtStatusPtr BindContextOutput(Ort::KernelContext& ctx,
 }
 
 OrtStatusPtr BindKernelOutput(Ort::KernelContext& ctx,
-                              OrtMemoryInfo* /*mem_info*/,
+                              const OrtMemoryInfo* /*mem_info*/,
                               DDSOutputAllocatorMap& allocator_map,
                               char const* output_name,
                               size_t output_index,
@@ -1416,6 +1416,7 @@ OrtStatus* TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(OrtEp* this
     tactics = GetTacticSourceFromString(tactic_sources_);
   }
   *compute_state = {
+                 static_cast<uint32_t>(device_id_),
                  fused_node_name,
                  builder_.get(),
                  &parsers_[fused_node_name],
@@ -2281,6 +2282,7 @@ OrtStatus* TRTEpNodeComputeInfo::ComputeImpl(OrtNodeComputeInfo* this_ptr, void*
   std::unordered_map<std::string, std::vector<int64_t>>
       shape_tensor_values_int64;  // same as above but for int64 shape tensor input
 
+  uint16_t device_id = trt_state->device_id;
   auto max_workspace_size = trt_state->max_workspace_size;
   auto trt_builder = trt_state->builder;
   auto trt_engine = trt_state->engine->get();
@@ -2317,7 +2319,7 @@ OrtStatus* TRTEpNodeComputeInfo::ComputeImpl(OrtNodeComputeInfo* this_ptr, void*
 
   // Get default OrtMemoryInfo from factory
   // Get allocator from OrtKernelContext
-  OrtMemoryInfo* mem_info = ep.factory_.GetDefaultMemInfo();
+  const OrtMemoryInfo* mem_info = ep.factory_.GetDefaultGpuMemInfoForDeviceId(device_id);
   OrtAllocator* alloc = nullptr;
   ep.GetAllocator(&alloc);
   if (alloc == nullptr) {

plugin_execution_providers/tensorrt/tensorrt_execution_provider.h

@@ -153,6 +153,7 @@ class OutputAllocator : public nvinfer1::IOutputAllocator {
 };
 
 struct TensorrtComputeState {
+  uint32_t device_id;
   std::string fused_node_name;
   nvinfer1::IBuilder* builder;
   tensorrt_ptr::unique_pointer<nvonnxparser::IParser>* parser = nullptr;
@@ -207,6 +208,7 @@ struct TensorrtComputeState {
 
 // Minimum information to construct kernel function state for direct engine load code path
 struct TensorrtComputeStateForEPContext {
+  uint32_t device_id;
   std::string fused_node_name;
   std::unique_ptr<nvinfer1::ICudaEngine>* engine = nullptr;
   std::unique_ptr<nvinfer1::IExecutionContext>* context = nullptr;

plugin_execution_providers/tensorrt/tensorrt_execution_provider_data_transfer.cc

@@ -9,14 +9,21 @@
 void CUDA_RETURN_IF_ERROR(cudaError_t res);
 
 /*static*/
-bool ORT_API_CALL TRTEpDataTransfer::CanCopyImpl(void* this_ptr,
-                                                   const OrtMemoryDevice* src_memory_device,
-                                                   const OrtMemoryDevice* dst_memory_device) noexcept {
+bool ORT_API_CALL TRTEpDataTransfer::CanCopyImpl(void* this_ptr, const OrtMemoryDevice* src_memory_device,
+                                                 const OrtMemoryDevice* dst_memory_device) noexcept {
   auto& impl = *static_cast<TRTEpDataTransfer*>(this_ptr);
-  bool src_is_our_device = impl.ep_api.MemoryDevice_AreEqual(src_memory_device, impl.device_mem_info);
-  bool dst_is_our_device = impl.ep_api.MemoryDevice_AreEqual(dst_memory_device, impl.device_mem_info);
 
-  return src_is_our_device || dst_is_our_device;
+  auto it = std::find_if(impl.cuda_gpu_mem_devices_.begin(), impl.cuda_gpu_mem_devices_.end(),
+                         [&impl, &src_memory_device, &dst_memory_device](const OrtMemoryDevice* memory_device) {
+                           bool src_is_our_device = impl.ep_api.MemoryDevice_AreEqual(src_memory_device, memory_device);
+                           bool dst_is_our_device = impl.ep_api.MemoryDevice_AreEqual(dst_memory_device, memory_device);
+                           return src_is_our_device || dst_is_our_device;
+                         });
+
+  if (it != impl.cuda_gpu_mem_devices_.end()) {
+    return true;
+  }
+  return false;
 }
 
 // function to copy one or more tensors.

plugin_execution_providers/tensorrt/tensorrt_execution_provider_data_transfer.h

@@ -6,9 +6,9 @@
 #include "tensorrt_execution_provider_utils.h"
 
 struct TRTEpDataTransfer : OrtDataTransferImpl, ApiPtrs {
-  TRTEpDataTransfer(ApiPtrs api_ptrs, const OrtMemoryDevice* device_mem_info_,
-                      const OrtMemoryDevice* shared_mem_info_ = nullptr)
-      : ApiPtrs(api_ptrs), device_mem_info{device_mem_info_}, shared_mem_info{shared_mem_info_} {
+  TRTEpDataTransfer(ApiPtrs api_ptrs, std::vector<const OrtMemoryDevice*> device_mem_infos,
+                      std::vector<const OrtMemoryDevice*> shared_mem_infos)
+      : ApiPtrs(api_ptrs), cuda_gpu_mem_devices_{device_mem_infos}, cuda_pinned_mem_devices_{shared_mem_infos} {
     CanCopy = CanCopyImpl;
     CopyTensors = CopyTensorsImpl;
     Release = ReleaseImpl;
@@ -25,6 +25,6 @@ struct TRTEpDataTransfer : OrtDataTransferImpl, ApiPtrs {
   static void ORT_API_CALL ReleaseImpl(void* this_ptr) noexcept;
 
  private:
-  const OrtMemoryDevice* device_mem_info;
-  const OrtMemoryDevice* shared_mem_info;
+  std::vector<const OrtMemoryDevice*> cuda_gpu_mem_devices_;
+  std::vector<const OrtMemoryDevice*> cuda_pinned_mem_devices_;
 };

plugin_execution_providers/tensorrt/tensorrt_provider_factory.cc

@@ -28,32 +28,6 @@ TensorrtExecutionProviderFactory::TensorrtExecutionProviderFactory(const char* e
   ReleaseAllocator = ReleaseAllocatorImpl;
 
   CreateDataTransfer = CreateDataTransferImpl;
-
-  // Default GPU allocator OrtMemoryInfo 
-  OrtMemoryInfo* mem_info = nullptr;
-  auto* status = ort_api.CreateMemoryInfo_V2("Cuda", OrtMemoryInfoDeviceType_GPU,
-                                       /*vendor*/ 0x10DE, /* device_id */ 0, OrtDeviceMemoryType_DEFAULT,
-                                       /*alignment*/ 0, OrtAllocatorType::OrtDeviceAllocator, &mem_info);
-  assert(status == nullptr);  // should never fail.
-  default_gpu_memory_info_ = MemoryInfoUniquePtr(mem_info, ort_api.ReleaseMemoryInfo);
-
-  // CUDA PINNED allocator OrtMemoryInfo
-  // HOST_ACCESSIBLE memory should use the non-CPU device type
-  mem_info = nullptr;
-  status = ort_api.CreateMemoryInfo_V2("CudaPinned", OrtMemoryInfoDeviceType_GPU,
-                                       /*vendor*/ 0x10DE, /* device_id */ 0, OrtDeviceMemoryType_HOST_ACCESSIBLE,
-                                       /*alignment*/ 0, OrtAllocatorType::OrtDeviceAllocator, &mem_info);
-  assert(status == nullptr);  // should never fail.
-  host_accessible_gpu_memory_info_ = MemoryInfoUniquePtr(mem_info, ort_api.ReleaseMemoryInfo);
-
-  // Create gpu data transfer
-  data_transfer_impl_ = std::make_unique<TRTEpDataTransfer>(
-      apis,
-      ep_api.MemoryInfo_GetMemoryDevice(default_gpu_memory_info_.get()),         // device memory
-      ep_api.MemoryInfo_GetMemoryDevice(host_accessible_gpu_memory_info_.get())  // shared memory
-  );
-
-  data_transfer_impl_.reset();  // but we're CPU only so we return nullptr for the IDataTransfer.
 }
 
 const char* ORT_API_CALL TensorrtExecutionProviderFactory::GetNameImpl(const OrtEpFactory* this_ptr) noexcept {
@@ -76,6 +50,9 @@ OrtStatus* ORT_API_CALL TensorrtExecutionProviderFactory::GetSupportedDevicesImp
   size_t& num_ep_devices = *p_num_ep_devices;
   auto* factory = static_cast<TensorrtExecutionProviderFactory*>(this_ptr);
 
+  std::vector<const OrtMemoryDevice*> cuda_gpu_mem_devices;
+  std::vector<const OrtMemoryDevice*> cuda_pinned_mem_devices;
+
   for (size_t i = 0; i < num_devices && num_ep_devices < max_ep_devices; ++i) {
     // C API
     const OrtHardwareDevice& device = *devices[i];
@@ -88,7 +65,7 @@ OrtStatus* ORT_API_CALL TensorrtExecutionProviderFactory::GetSupportedDevicesImp
 
       // The ep options can be provided here as default values.
       // Users can also call SessionOptionsAppendExecutionProvider_V2 C API with provided ep options to override.
-      factory->ort_api.AddKeyValuePair(ep_metadata, "version", "0.1"); // random example using made up values
+      factory->ort_api.AddKeyValuePair(ep_metadata, "gpu_type", "data center"); // random example using made up values
       factory->ort_api.AddKeyValuePair(ep_options, "trt_builder_optimization_level", "3");
 
       // OrtEpDevice copies ep_metadata and ep_options.
@@ -103,25 +80,60 @@ OrtStatus* ORT_API_CALL TensorrtExecutionProviderFactory::GetSupportedDevicesImp
         return status;
       }
 
-      // register the allocator info required by the EP.
-      RETURN_IF_ERROR(factory->ep_api.EpDevice_AddAllocatorInfo(ep_device, factory->default_gpu_memory_info_.get()));
-      RETURN_IF_ERROR(factory->ep_api.EpDevice_AddAllocatorInfo(ep_device, factory->host_accessible_gpu_memory_info_.get()));
+      uint32_t vendor_id = factory->ort_api.HardwareDevice_VendorId(&device);
+      uint32_t device_id = factory->ort_api.HardwareDevice_DeviceId(&device);
+      
+      // CUDA allocator OrtMemoryInfo
+      OrtMemoryInfo* mem_info = nullptr;
+      status = factory->ort_api.CreateMemoryInfo_V2("Cuda", OrtMemoryInfoDeviceType_GPU, vendor_id, device_id, OrtDeviceMemoryType_DEFAULT,
+                                               /*alignment*/ 0, OrtAllocatorType::OrtDeviceAllocator, &mem_info);
+
+      assert(status == nullptr);  // should never fail.
+      MemoryInfoUniquePtr cuda_gpu_memory_info = MemoryInfoUniquePtr(mem_info, factory->ort_api.ReleaseMemoryInfo);
+
+      // CUDA PINNED allocator OrtMemoryInfo
+      // HOST_ACCESSIBLE memory should use the non-CPU device type.
+      mem_info = nullptr;
+      status = factory->ort_api.CreateMemoryInfo_V2("CudaPinned", OrtMemoryInfoDeviceType_GPU, vendor_id, device_id, OrtDeviceMemoryType_HOST_ACCESSIBLE,
+                                           /*alignment*/ 0, OrtAllocatorType::OrtDeviceAllocator, &mem_info);
+
+      assert(status == nullptr);  // should never fail.
+      MemoryInfoUniquePtr cuda_pinned_memory_info = MemoryInfoUniquePtr(mem_info, factory->ort_api.ReleaseMemoryInfo);
+
+      // Register the allocator info required by TRT EP.
+      RETURN_IF_ERROR(factory->ep_api.EpDevice_AddAllocatorInfo(ep_device, cuda_gpu_memory_info.get()));
+      RETURN_IF_ERROR(factory->ep_api.EpDevice_AddAllocatorInfo(ep_device, cuda_pinned_memory_info.get()));
+
+      // Get memory device from memory info for gpu data transfer
+      cuda_gpu_mem_devices.push_back(factory->ep_api.MemoryInfo_GetMemoryDevice(cuda_gpu_memory_info.get()));
+      cuda_pinned_mem_devices.push_back(factory->ep_api.MemoryInfo_GetMemoryDevice(cuda_pinned_memory_info.get()));
+
+      factory->SetDefaultGpuMemInfo(std::move(cuda_gpu_memory_info), device_id);
+      factory->SetHostAccessibleMemInfo(std::move(cuda_pinned_memory_info), device_id);
 
       ep_devices[num_ep_devices++] = ep_device;
     }
 
-    // C++ API equivalent. Throws on error.
-    //{
-    //  Ort::ConstHardwareDevice device(devices[i]);
-    //  if (device.Type() == OrtHardwareDeviceType::OrtHardwareDeviceType_GPU) {
-    //    Ort::KeyValuePairs ep_metadata;
-    //    Ort::KeyValuePairs ep_options;
-    //    ep_metadata.Add("version", "0.1");
-    //    ep_options.Add("trt_builder_optimization_level", "3");
-    //    Ort::EpDevice ep_device{*this_ptr, device, ep_metadata.GetConst(), ep_options.GetConst()};
-    //    ep_devices[num_ep_devices++] = ep_device.release();
-    //  }
-    //}
+  // Create gpu data transfer
+  auto data_transfer_impl = std::make_unique<TRTEpDataTransfer>(
+      static_cast<const ApiPtrs&>(*factory),
+      cuda_gpu_mem_devices,    // device memory
+      cuda_pinned_mem_devices  // shared memory
+  );
+
+  
+  // C++ API equivalent. Throws on error.
+  //{
+  //  Ort::ConstHardwareDevice device(devices[i]);
+  //  if (device.Type() == OrtHardwareDeviceType::OrtHardwareDeviceType_GPU) {
+  //    Ort::KeyValuePairs ep_metadata;
+  //    Ort::KeyValuePairs ep_options;
+  //    ep_metadata.Add("version", "0.1");
+  //    ep_options.Add("trt_builder_optimization_level", "3");
+  //    Ort::EpDevice ep_device{*this_ptr, device, ep_metadata.GetConst(), ep_options.GetConst()};
+  //    ep_devices[num_ep_devices++] = ep_device.release();
+  //  }
+  //}
   }
 
   return nullptr;
@@ -181,11 +193,14 @@ OrtStatus* ORT_API_CALL TensorrtExecutionProviderFactory::CreateAllocatorImpl(
 
   // NOTE: The OrtMemoryInfo pointer should only ever be coming straight from an OrtEpDevice, and pointer based
   // matching should work.
-  if (memory_info == factory.default_gpu_memory_info_.get()) {
+  
+  uint32_t device_id = 0;
+
+  if (factory.GetDeviceIdForDefaultGpuMemInfo(memory_info, &device_id)) {
     // create a CUDA allocator
-    auto cuda_allocator = std::make_unique<CUDAAllocator>(memory_info);
+    auto cuda_allocator = std::make_unique<CUDAAllocator>(memory_info, static_cast<uint16_t>(device_id));
     *allocator = cuda_allocator.release();
-  } else if (memory_info == factory.host_accessible_gpu_memory_info_.get()) {
+  } else if (factory.GetDeviceIdForHostAccessibleMemInfo(memory_info, &device_id)) {
     // create a CUDA PINNED allocator
     auto cuda_pinned_allocator = std::make_unique<CUDAPinnedAllocator>(memory_info);
     *allocator = cuda_pinned_allocator.release();
@@ -212,8 +227,50 @@ OrtStatus* ORT_API_CALL TensorrtExecutionProviderFactory::CreateDataTransferImpl
   return nullptr;
 }
 
-OrtMemoryInfo* TensorrtExecutionProviderFactory::GetDefaultMemInfo() const {
-  return default_gpu_memory_info_.get();
+bool TensorrtExecutionProviderFactory::GetDeviceIdForDefaultGpuMemInfo(const OrtMemoryInfo* mem_info, uint32_t* device_id) const {
+  auto iter = cuda_gpu_memory_info_to_device_id_map_.find(mem_info);
+  if (iter != cuda_gpu_memory_info_to_device_id_map_.end()) {
+    *device_id = iter->second;
+    return true;
+  }
+  return false;
+}
+
+const OrtMemoryInfo* TensorrtExecutionProviderFactory::GetDefaultGpuMemInfoForDeviceId(uint32_t device_id) const {
+  auto iter = device_id_to_cuda_gpu_memory_info_map_.find(device_id);
+  if (iter != device_id_to_cuda_gpu_memory_info_map_.end()) {
+    return iter->second;
+  }
+  return nullptr;
+}
+
+void TensorrtExecutionProviderFactory::SetDefaultGpuMemInfo(MemoryInfoUniquePtr mem_info, uint32_t device_id) {
+  cuda_gpu_memory_info_to_device_id_map_[mem_info.get()] = device_id;
+  device_id_to_cuda_gpu_memory_info_map_[device_id] = mem_info.get();
+  cuda_gpu_memory_infos_.push_back(std::move(mem_info));
+}
+
+bool TensorrtExecutionProviderFactory::GetDeviceIdForHostAccessibleMemInfo(const OrtMemoryInfo* mem_info, uint32_t* device_id) const {
+  auto iter = cuda_pinned_memory_info_to_device_id_map_.find(mem_info);
+  if (iter != cuda_pinned_memory_info_to_device_id_map_.end()) {
+    *device_id = iter->second;
+    return true;
+  }
+  return false;
+}
+
+const OrtMemoryInfo* TensorrtExecutionProviderFactory::GetHostAccessibleMemInfoForDeviceId(uint32_t device_id) const {
+  auto iter = device_id_to_cuda_pinned_memory_info_map_.find(device_id);
+  if (iter != device_id_to_cuda_pinned_memory_info_map_.end()) {
+    return iter->second;
+  }
+  return nullptr;
+}
+
+void TensorrtExecutionProviderFactory::SetHostAccessibleMemInfo(MemoryInfoUniquePtr mem_info, uint32_t device_id) {
+  cuda_pinned_memory_info_to_device_id_map_[mem_info.get()] = device_id;
+  device_id_to_cuda_pinned_memory_info_map_[device_id] = mem_info.get();
+  cuda_pinned_memory_infos_.push_back(std::move(mem_info));
 }
 
 // To make symbols visible on macOS/iOS

plugin_execution_providers/tensorrt/tensorrt_provider_factory.h

@@ -3,13 +3,18 @@
 #include "tensorrt_execution_provider_utils.h"
 #include "tensorrt_execution_provider_data_transfer.h"
 
+using MemoryInfoUniquePtr = std::unique_ptr<OrtMemoryInfo, std::function<void(OrtMemoryInfo*)>>;
+
 ///
 /// Plugin TensorRT EP factory that can create an OrtEp and return information about the supported hardware devices.
 ///
 struct TensorrtExecutionProviderFactory : public OrtEpFactory, public ApiPtrs {
  public:
   TensorrtExecutionProviderFactory(const char* ep_name, ApiPtrs apis);
-  OrtMemoryInfo* GetDefaultMemInfo() const;
+
+  const OrtMemoryInfo* GetDefaultGpuMemInfoForDeviceId(uint32_t device_id) const;
+
+  const OrtMemoryInfo* GetHostAccessibleMemInfoForDeviceId(uint32_t device_id) const;
 
  private:
   static const char* ORT_API_CALL GetNameImpl(const OrtEpFactory* this_ptr) noexcept;
@@ -37,17 +42,30 @@ struct TensorrtExecutionProviderFactory : public OrtEpFactory, public ApiPtrs {
   static OrtStatus* ORT_API_CALL CreateDataTransferImpl(OrtEpFactory* this_ptr,
                                                         OrtDataTransferImpl** data_transfer) noexcept;
 
+  bool GetDeviceIdForDefaultGpuMemInfo(const OrtMemoryInfo* mem_info, uint32_t* device_id) const;
+
+  void SetDefaultGpuMemInfo(MemoryInfoUniquePtr mem_info, uint32_t device_id);
+
+  bool GetDeviceIdForHostAccessibleMemInfo(const OrtMemoryInfo* mem_info, uint32_t* device_id) const;
+
+  void SetHostAccessibleMemInfo(MemoryInfoUniquePtr mem_info, uint32_t device_id);
+
   const std::string ep_name_;           // EP name
   const std::string vendor_{"Nvidia"};  // EP vendor name
 
-  // CPU allocator so we can control the arena behavior. optional as ORT always provides a CPU allocator if needed.
-  using MemoryInfoUniquePtr = std::unique_ptr<OrtMemoryInfo, std::function<void(OrtMemoryInfo*)>>;
-  //MemoryInfoUniquePtr cpu_memory_info_;
+  // OrtMemoryInfo for allocators and data transfer.
+  
+  // CUDA gpu memory and CUDA pinned memory are required for allocator and data transfer, these are the OrtMemoryInfo instance required for that.
+  // Current TRT EP implementation uses one default OrtMemoryInfo and one host accessible OrtMemoryInfo per ep device.
+  std::unordered_map<const OrtMemoryInfo*, uint32_t> cuda_gpu_memory_info_to_device_id_map_;   // OrtMemoryInfo -> device id
+  std::unordered_map<const OrtMemoryInfo*, uint32_t> cuda_pinned_memory_info_to_device_id_map_;
+  std::unordered_map<uint32_t, const OrtMemoryInfo*> device_id_to_cuda_gpu_memory_info_map_;   // device id -> OrtMemoryInfo
+  std::unordered_map<uint32_t, const OrtMemoryInfo*> device_id_to_cuda_pinned_memory_info_map_;
+  std::vector<MemoryInfoUniquePtr> cuda_gpu_memory_infos_;
+  std::vector<MemoryInfoUniquePtr> cuda_pinned_memory_infos_;
 
-  // GPU memory and pinned/shared memory are required for data transfer, these are the
-  // OrtMemoryInfo instance required for that.
-  MemoryInfoUniquePtr default_gpu_memory_info_;
-  MemoryInfoUniquePtr host_accessible_gpu_memory_info_;
+  // CPU allocator so we can control the arena behavior. optional as ORT always provides a CPU allocator if needed.
+  // MemoryInfoUniquePtr cpu_memory_info_;
 
   std::unique_ptr<TRTEpDataTransfer> data_transfer_impl_;  // data transfer implementation for this factory
 };