Merge pull request #5855 from dalg24/num_threads_and_device_id

Add runtime functions to query the number of threads and the device ID
sandialabs · Apr 10, 2023 · 8352a11 · 8352a11
2 parents b24dcb4 + d8d9c58
commit 8352a11
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 9 deletions.
diff --git a/core/src/Kokkos_Core.hpp b/core/src/Kokkos_Core.hpp
@@ -100,6 +100,9 @@ void declare_configuration_metadata(const std::string& category,
 [[nodiscard]] bool is_initialized() noexcept;
 [[nodiscard]] bool is_finalized() noexcept;
 
+[[nodiscard]] int device_id() noexcept;
+[[nodiscard]] int num_threads() noexcept;
+
 bool show_warnings() noexcept;
 bool tune_internals() noexcept;
 

diff --git a/core/src/impl/Kokkos_Core.cpp b/core/src/impl/Kokkos_Core.cpp
@@ -165,6 +165,26 @@ bool is_valid_map_device_id_by(std::string const& x) {
 
 }  // namespace
 
+[[nodiscard]] int Kokkos::device_id() noexcept {
+#if defined(KOKKOS_ENABLE_CUDA)
+  return Cuda().cuda_device();
+#elif defined(KOKKOS_ENABLE_HIP)
+  return HIP().hip_device();
+#elif defined(KOKKOS_ENABLE_OPENACC)
+  return Experimental::OpenACC().acc_device_number();
+#elif defined(KOKKOS_ENABLE_OPENMPTARGET)
+  return omp_get_default_device();  // FIXME_OPENMPTARGET
+#elif defined(KOKKOS_ENABLE_SYCL)
+  return Experimental::Impl::SYCLInternal::m_syclDev;
+#else
+  return -1;
+#endif
+}
+
+[[nodiscard]] int Kokkos::num_threads() noexcept {
+  return DefaultHostExecutionSpace().concurrency();
+}
+
 Kokkos::Impl::ExecSpaceManager& Kokkos::Impl::ExecSpaceManager::get_instance() {
   static ExecSpaceManager space_initializer = {};
   return space_initializer;

diff --git a/core/unit_test/UnitTest_DeviceAndThreads.cpp b/core/unit_test/UnitTest_DeviceAndThreads.cpp
@@ -38,21 +38,24 @@ int get_device_count() {
 }
 
 int get_device_id() {
+  int device_id;
 #if defined(KOKKOS_ENABLE_CUDA)
-  int device;
-  KOKKOS_IMPL_CUDA_SAFE_CALL(cudaGetDevice(&device));
-  return device;
+  KOKKOS_IMPL_CUDA_SAFE_CALL(cudaGetDevice(&device_id));
 #elif defined(KOKKOS_ENABLE_HIP)
-  int device_id;
   KOKKOS_IMPL_HIP_SAFE_CALL(hipGetDevice(&device_id));
-  return device_id;
 #elif defined(KOKKOS_ENABLE_OPENMPTARGET)
-  return omp_get_device_num();
+  device_id = omp_get_device_num();
 #elif defined(KOKKOS_ENABLE_OPENACC)
-  return acc_get_device_num(acc_get_device_type());
+  device_id = acc_get_device_num(acc_get_device_type());
+#elif defined(KOKKOS_ENABLE_SYCL)
+  // FIXME_SYCL ?
+  assert(false);
+  return -2;
 #else
-  return -1;
+  device_id = -1;
 #endif
+  assert(device_id == Kokkos::device_id());
+  return device_id;
 }
 
 int get_max_threads() {
@@ -66,7 +69,9 @@ int get_max_threads() {
 }
 
 int get_num_threads() {
-  return Kokkos::DefaultHostExecutionSpace().concurrency();
+  int const num_threads = Kokkos::DefaultHostExecutionSpace().concurrency();
+  assert(num_threads == Kokkos::num_threads());
+  return num_threads;
 }
 
 int get_disable_warnings() { return !Kokkos::show_warnings(); }