[Mult-GPU inferencing] Add new API to get/set device id. Set correct device id in cuda allocator. (#6592)

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -1163,6 +1163,17 @@ struct OrtApi {
    */
   ORT_API2_STATUS(SessionOptionsAppendExecutionProvider_TensorRT,
                   _In_ OrtSessionOptions* options, _In_ const OrtTensorRTProviderOptions* tensorrt_options);
+
+  /**
+  * Set the current device id of the GPU execution provider (cuda/tensorrt/rocm). The device id should be less
+  * than the total number of devices available. Using this API makes sense only when doing multi-GPU inferencing.
+  */
+  ORT_API2_STATUS(SetCurrentGpuDeviceId, _In_ int device_id);
+
+  /**
+   * Get the current device id of the GPU execution provider (cuda/tensorrt/rocm).
+   */
+  ORT_API2_STATUS(GetCurrentGpuDeviceId, _In_ int* device_id);
 };
 
 /*

onnxruntime/core/providers/cuda/cuda_allocator.cc

@@ -32,7 +32,23 @@ void CUDAAllocator::CheckDevice(bool throw_when_fail) const {
 #endif
 }
 
+void CUDAAllocator::SetDevice(bool throw_when_fail) const {
+  int current_device;
+  auto cuda_err = cudaGetDevice(&current_device);
+  if (cuda_err == cudaSuccess) {
+    int allocator_device_id = Info().id;
+    if (current_device != allocator_device_id) {
+      cuda_err = cudaSetDevice(allocator_device_id);
+    }
+  }
+
+  if (cuda_err != cudaSuccess && throw_when_fail) {
+    CUDA_CALL_THROW(cuda_err);
+  }
+}
+
 void* CUDAAllocator::Alloc(size_t size) {
+  SetDevice(true);
   CheckDevice(true);
   void* p = nullptr;
   if (size > 0) {
@@ -43,6 +59,7 @@ void* CUDAAllocator::Alloc(size_t size) {
 }
 
 void CUDAAllocator::Free(void* p) {
+  SetDevice(false);
   CheckDevice(false);  // ignore CUDA failure when free
   cudaFree(p);         // do not throw error since it's OK for cudaFree to fail during shutdown
 }

onnxruntime/core/providers/cuda/cuda_allocator.h

@@ -20,6 +20,7 @@ class CUDAAllocator : public IAllocator {
 
  private:
   void CheckDevice(bool throw_when_fail) const;
+  void SetDevice(bool throw_when_fail) const;
 };
 
 //TODO: add a default constructor

onnxruntime/core/providers/cuda/cuda_provider_factory.cc

@@ -63,3 +63,31 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider_CUDA,
 
   return nullptr;
 }
+
+ORT_API_STATUS_IMPL(OrtApis::SetCurrentGpuDeviceId, _In_ int device_id) {
+  int num_devices;
+  auto cuda_err = cudaGetDeviceCount(&num_devices);
+  if (cuda_err != cudaSuccess) {
+    return CreateStatus(ORT_FAIL, "Failed to set device id since cudaGetDeviceCount failed.");
+  }
+
+  if (device_id >= num_devices) {
+    std::ostringstream ostr;
+    ostr << "Invalid device id. Device id should be less than total number of devices (" << num_devices << ")";
+    return CreateStatus(ORT_INVALID_ARGUMENT, ostr.str().c_str());
+  }
+
+  cuda_err = cudaSetDevice(device_id);
+  if (cuda_err != cudaSuccess) {
+    return CreateStatus(ORT_FAIL, "Failed to set device id.");
+  }
+  return nullptr;
+}
+
+ORT_API_STATUS_IMPL(OrtApis::GetCurrentGpuDeviceId, _In_ int* device_id) {
+  auto cuda_err = cudaGetDevice(device_id);
+  if (cuda_err != cudaSuccess) {
+    return CreateStatus(ORT_FAIL, "Failed to get device id.");
+  }
+  return nullptr;
+}

onnxruntime/core/session/onnxruntime_c_api.cc

@@ -1823,6 +1823,16 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider_CUDA,
   ORT_UNUSED_PARAMETER(cuda_options);
   return CreateStatus(ORT_FAIL, "CUDA execution provider is not enabled.");
 }
+
+ORT_API_STATUS_IMPL(OrtApis::SetCurrentGpuDeviceId, _In_ int device_id) {
+  ORT_UNUSED_PARAMETER(device_id);
+  return CreateStatus(ORT_FAIL, "CUDA execution provider is not enabled.");
+}
+
+ORT_API_STATUS_IMPL(OrtApis::GetCurrentGpuDeviceId, _In_ int* device_id) {
+  ORT_UNUSED_PARAMETER(device_id);
+  return CreateStatus(ORT_FAIL, "CUDA execution provider is not enabled.");
+}
 #endif
 
 #if defined(ORT_MINIMAL_BUILD)
@@ -2093,7 +2103,10 @@ static constexpr OrtApi ort_api_1_to_7 = {
 
     // Version 7 - In development, feel free to add/remove/rearrange here
     &OrtApis::ModelMetadataGetGraphDescription,
+
     &OrtApis::SessionOptionsAppendExecutionProvider_TensorRT,
+    &OrtApis::SetCurrentGpuDeviceId,
+    &OrtApis::GetCurrentGpuDeviceId,
 };
 
 // Assert to do a limited check to ensure Version 1 of OrtApi never changes (will detect an addition or deletion but not if they cancel out each other)

onnxruntime/core/session/ort_apis.h

@@ -257,4 +257,6 @@ ORT_API_STATUS_IMPL(CreateArenaCfg, _In_ size_t max_mem, int arena_extend_strate
 ORT_API(void, ReleaseArenaCfg, _Frees_ptr_opt_ OrtArenaCfg*);
 ORT_API_STATUS_IMPL(SessionOptionsAppendExecutionProvider_TensorRT,
                     _In_ OrtSessionOptions* options, _In_ const OrtTensorRTProviderOptions* tensorrt_options);
+ORT_API_STATUS_IMPL(SetCurrentGpuDeviceId, _In_ int device_id);
+ORT_API_STATUS_IMPL(GetCurrentGpuDeviceId, _In_ int* device_id);
 }  // namespace OrtApis