Let Impl::get_gpu return std::optional and delegate device selection …

…when appropriate
masterleinad · Feb 7, 2024 · 1327c37 · 1327c37
1 parent 442e4d4
commit 1327c37
Show file tree

Hide file tree

Showing 7 changed files with 13 additions and 9 deletions.
diff --git a/core/src/Cuda/Kokkos_Cuda_Instance.cpp b/core/src/Cuda/Kokkos_Cuda_Instance.cpp
@@ -532,7 +532,7 @@ int Cuda::impl_is_initialized() {
 }
 
 void Cuda::impl_initialize(InitializationSettings const &settings) {
-  const int cuda_device_id = Impl::get_gpu(settings);
+  const int cuda_device_id = Impl::get_gpu(settings).value_or(0);
 
   cudaDeviceProp cudaProp;
   KOKKOS_IMPL_CUDA_SAFE_CALL(

diff --git a/core/src/HIP/Kokkos_HIP.cpp b/core/src/HIP/Kokkos_HIP.cpp
@@ -42,7 +42,7 @@ int HIP::impl_is_initialized() {
 }
 
 void HIP::impl_initialize(InitializationSettings const& settings) {
-  const int hip_device_id = Impl::get_gpu(settings);
+  const int hip_device_id = Impl::get_gpu(settings).value_or(0);
 
   Impl::HIPInternal::m_hipDev = hip_device_id;
   KOKKOS_IMPL_HIP_SAFE_CALL(

diff --git a/core/src/OpenACC/Kokkos_OpenACC.cpp b/core/src/OpenACC/Kokkos_OpenACC.cpp
@@ -59,7 +59,7 @@ void Kokkos::Experimental::OpenACC::impl_initialize(
         acc_get_device_num(acc_device_host);
   } else {
     using Kokkos::Impl::get_gpu;
-    int const dev_num = get_gpu(settings);
+    int const dev_num = get_gpu(settings).value_or(0);
     acc_set_device_num(dev_num, Impl::OpenACC_Traits::dev_type);
     Impl::OpenACCInternal::m_acc_device_num = dev_num;
   }

diff --git a/core/src/OpenMPTarget/Kokkos_OpenMPTarget_Instance.cpp b/core/src/OpenMPTarget/Kokkos_OpenMPTarget_Instance.cpp
@@ -179,7 +179,7 @@ void OpenMPTarget::impl_static_fence(const std::string& name) {
 
 void OpenMPTarget::impl_initialize(InitializationSettings const& settings) {
   using Kokkos::Impl::get_gpu;
-  const int device_num = get_gpu(settings);
+  const int device_num = get_gpu(settings).value_or(0);
   omp_set_default_device(device_num);
 
   Impl::OpenMPTargetInternal::impl_singleton()->impl_initialize();

diff --git a/core/src/SYCL/Kokkos_SYCL.cpp b/core/src/SYCL/Kokkos_SYCL.cpp
@@ -154,7 +154,7 @@ void SYCL::impl_initialize(InitializationSettings const& settings) {
     return;
   }
 #endif
-  const auto id = ::Kokkos::Impl::get_gpu(settings);
+  const auto id = ::Kokkos::Impl::get_gpu(settings).value_or(0);
   Impl::SYCLInternal::singleton().initialize(gpu_devices[id]);
   Impl::SYCLInternal::m_syclDev = id;
 }

diff --git a/core/src/impl/Kokkos_Core.cpp b/core/src/impl/Kokkos_Core.cpp
@@ -16,6 +16,7 @@
 
 #ifndef KOKKOS_IMPL_PUBLIC_INCLUDE
 #define KOKKOS_IMPL_PUBLIC_INCLUDE
+#include <optional>
 #endif
 
 #include <Kokkos_Core.hpp>
@@ -377,7 +378,8 @@ std::vector<int> Kokkos::Impl::get_visible_devices(int device_count) {
   return visible_devices;
 }
 
-int Kokkos::Impl::get_gpu(const InitializationSettings& settings) {
+std::optional<int> Kokkos::Impl::get_gpu(
+    const InitializationSettings& settings) {
   std::vector<int> visible_devices = get_visible_devices(get_device_count());
   int const num_devices            = visible_devices.size();
   // device_id is provided
@@ -425,14 +427,15 @@ int Kokkos::Impl::get_gpu(const InitializationSettings& settings) {
 
   int const mpi_local_rank = mpi_local_rank_on_node();
 
-  // use first GPU available for execution if unable to detect local MPI rank
+  // if unable to detect local MPI rank return nullopt to delegate device
+  // selection to the backend
   if (mpi_local_rank < 0) {
     if (settings.has_map_device_id_by()) {
       std::cerr << "Warning: unable to detect local MPI rank."
                 << " Falling back to the first GPU available for execution."
                 << " Raised by Kokkos::initialize()." << std::endl;
     }
-    return visible_devices[0];
+    return std::nullopt;
   }
 
   // use device assigned by CTest when resource allocation is activated

diff --git a/core/src/impl/Kokkos_DeviceManagement.hpp b/core/src/impl/Kokkos_DeviceManagement.hpp
@@ -17,12 +17,13 @@
 #ifndef KOKKOS_DEVICE_MANAGEMENT_HPP
 #define KOKKOS_DEVICE_MANAGEMENT_HPP
 
+#include <optional>
 #include <vector>
 
 namespace Kokkos {
 class InitializationSettings;
 namespace Impl {
-int get_gpu(const Kokkos::InitializationSettings& settings);
+std::optional<int> get_gpu(const Kokkos::InitializationSettings& settings);
 // This declaration is provided for testing purposes only
 int get_ctest_gpu(int local_rank);
 std::vector<int> get_visible_devices(int device_count);  // test-only