From 739c9d4875c91a069f246d4546f84d3e7a715ebc Mon Sep 17 00:00:00 2001 From: KeDengMS Date: Thu, 9 Apr 2020 18:54:58 -0700 Subject: [PATCH] Always call cudaSetDevice at the beginning of session::Run (#3475) This is required for running multithreaded with multi-GPUs. Without it, when running in a work thread it would default to GPU 0, while CUDAExecutionProvider is assigned on other GPUs. That might cause CUDA crash when some CUDA resources is from GPU 0, while being used in GPU N>0. --- onnxruntime/core/providers/cuda/cuda_execution_provider.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index 514d6d50d6..d4e5d39b15 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -198,6 +198,8 @@ void CUDAExecutionProvider::AddDeferredReleaseCPUPtr(void* p) { } Status CUDAExecutionProvider::OnRunStart() { + // always set CUDA device when session::Run() in case it runs in a worker thread + CUDA_RETURN_IF_ERROR(cudaSetDevice(GetDeviceId())); auto cpu_alloc = GetAllocator(0, OrtMemTypeCPU); // check if cudaEvents has passed for deferred release // note that we need to take a mutex in case of multi-threaded Run()