Set the device id explicitly for CUDA API calls in impl_initialize

japlews · Nov 7, 2023 · d4a517f · d4a517f
1 parent 403c34f
commit d4a517f
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/core/src/Cuda/Kokkos_Cuda_Instance.cpp b/core/src/Cuda/Kokkos_Cuda_Instance.cpp
@@ -685,12 +685,10 @@ void Cuda::impl_initialize(InitializationSettings const &settings) {
   const auto &dev_info     = Impl::CudaInternalDevices::singleton();
 
   const struct cudaDeviceProp &cudaProp = dev_info.m_cudaProp[cuda_device_id];
+  Impl::CudaInternal::m_deviceProp      = cudaProp;
 
-  m_cudaDev                        = cuda_device_id;
-  Impl::CudaInternal::m_deviceProp = cudaProp;
-
-  Kokkos::Impl::cuda_device_synchronize(
-      "Kokkos::CudaInternal::initialize: Fence on space initialization");
+  KOKKOS_IMPL_CUDA_SAFE_CALL(cudaSetDevice(cuda_device_id));
+  KOKKOS_IMPL_CUDA_SAFE_CALL(cudaDeviceSynchronize());
 
   // Query what compute capability architecture a kernel executes:
   Impl::CudaInternal::m_cudaArch = Impl::cuda_kernel_arch(cuda_device_id);
@@ -758,6 +756,7 @@ Kokkos::Cuda::initialize WARNING: Cuda is allocating into UVMSpace by default
   //----------------------------------
 
   cudaStream_t singleton_stream;
+  KOKKOS_IMPL_CUDA_SAFE_CALL(cudaSetDevice(cuda_device_id));
   KOKKOS_IMPL_CUDA_SAFE_CALL(cudaStreamCreate(&singleton_stream));
 
   Impl::CudaInternal::singleton().initialize(singleton_stream,