Add Kokkos::num_threads() and Kokkos::device_id()

masterleinad · Feb 3, 2023 · a4af6f7 · a4af6f7
1 parent 2aa2576
commit a4af6f7
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/core/src/Kokkos_Core.hpp b/core/src/Kokkos_Core.hpp
@@ -99,6 +99,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;