Serial: Allow for distinct execution space instances (kokkos#6441)

* Serial: Allow for distinct execution space instances * Remove redundant test * Make default constructor for NewInstance explicit * Shorten partition_space implementation. Co-authored-by: Damien L-G <dalg24+github@gmail.com> * Fix unused function warning --------- Co-authored-by: Damien L-G <dalg24+github@gmail.com>
masterleinad · Oct 10, 2023 · ebef19b · ebef19b
1 parent ad9eb20
commit ebef19b
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 39 deletions.
diff --git a/core/src/Serial/Kokkos_Serial.cpp b/core/src/Serial/Kokkos_Serial.cpp
@@ -145,6 +145,12 @@ Serial::Serial()
     : m_space_instance(&Impl::SerialInternal::singleton(),
                        [](Impl::SerialInternal*) {}) {}
 
+Serial::Serial(NewInstance)
+    : m_space_instance(new Impl::SerialInternal, [](Impl::SerialInternal* ptr) {
+        ptr->finalize();
+        delete ptr;
+      }) {}
+
 void Serial::print_configuration(std::ostream& os, bool /*verbose*/) const {
   os << "Host Serial Execution Space:\n";
   os << "  KOKKOS_ENABLE_SERIAL: yes\n";

diff --git a/core/src/Serial/Kokkos_Serial.hpp b/core/src/Serial/Kokkos_Serial.hpp
@@ -72,6 +72,10 @@ class SerialInternal {
 };
 }  // namespace Impl
 
+struct NewInstance {
+  explicit NewInstance() = default;
+};
+
 /// \class Serial
 /// \brief Kokkos device for non-parallel execution
 ///
@@ -108,6 +112,8 @@ class Serial {
 
   Serial();
 
+  Serial(NewInstance);
+
   /// \brief True if and only if this method is being called in a
   ///   thread-parallel function.
   ///
@@ -218,6 +224,38 @@ struct MemorySpaceAccess<Kokkos::Serial::memory_space,
 }  // namespace Impl
 }  // namespace Kokkos
 
+namespace Kokkos::Experimental {
+
+template <class... Args>
+std::vector<Serial> partition_space(const Serial&, Args...) {
+  static_assert(
+      (... && std::is_arithmetic_v<Args>),
+      "Kokkos Error: partitioning arguments must be integers or floats");
+  std::vector<Serial> instances;
+  instances.reserve(sizeof...(Args));
+  std::generate_n(std::back_inserter(instances), sizeof...(Args),
+                  []() { return Serial{NewInstance{}}; });
+  return instances;
+}
+
+template <class T>
+std::vector<Serial> partition_space(const Serial&,
+                                    std::vector<T> const& weights) {
+  static_assert(
+      std::is_arithmetic<T>::value,
+      "Kokkos Error: partitioning arguments must be integers or floats");
+
+  // We only care about the number of instances to create and ignore weights
+  // otherwise.
+  std::vector<Serial> instances;
+  instances.reserve(weights.size());
+  std::generate_n(std::back_inserter(instances), weights.size(),
+                  []() { return Serial{NewInstance{}}; });
+  return instances;
+}
+
+}  // namespace Kokkos::Experimental
+
 #include <Serial/Kokkos_Serial_Parallel_Range.hpp>
 #include <Serial/Kokkos_Serial_Parallel_MDRange.hpp>
 #include <Serial/Kokkos_Serial_Parallel_Team.hpp>

diff --git a/core/unit_test/TestExecSpacePartitioning.hpp b/core/unit_test/TestExecSpacePartitioning.hpp
@@ -29,30 +29,35 @@ struct SumFunctor {
 };
 
 template <class ExecSpace>
-void check_distinctive(ExecSpace, ExecSpace) {}
-
+void check_distinctive([[maybe_unused]] ExecSpace exec1,
+                       [[maybe_unused]] ExecSpace exec2) {
+#ifdef KOKKOS_ENABLE_SERIAL
+  if constexpr (std::is_same_v<ExecSpace, Kokkos::Serial>) {
+    ASSERT_NE(exec1, exec2);
+  }
+#endif
+#ifdef KOKKOS_ENABLE_OPENMP
+  if constexpr (std::is_same_v<ExecSpace, Kokkos::OpenMP>) {
+    ASSERT_NE(exec1, exec2);
+  }
+#endif
 #ifdef KOKKOS_ENABLE_CUDA
-void check_distinctive(Kokkos::Cuda exec1, Kokkos::Cuda exec2) {
-  ASSERT_NE(exec1.cuda_stream(), exec2.cuda_stream());
-}
+  if constexpr (std::is_same_v<ExecSpace, Kokkos::Cuda>) {
+    ASSERT_NE(exec1.cuda_stream(), exec2.cuda_stream());
+  }
 #endif
 #ifdef KOKKOS_ENABLE_HIP
-void check_distinctive(Kokkos::HIP exec1, Kokkos::HIP exec2) {
-  ASSERT_NE(exec1.hip_stream(), exec2.hip_stream());
-}
+  if constexpr (std::is_same_v<ExecSpace, Kokkos::HIP>) {
+    ASSERT_NE(exec1.hip_stream(), exec2.hip_stream());
+  }
 #endif
 #ifdef KOKKOS_ENABLE_SYCL
-void check_distinctive(Kokkos::Experimental::SYCL exec1,
-                       Kokkos::Experimental::SYCL exec2) {
-  ASSERT_NE(*exec1.impl_internal_space_instance()->m_queue,
-            *exec2.impl_internal_space_instance()->m_queue);
-}
+  if constexpr (std::is_same_v<ExecSpace, Kokkos::Experimental::SYCL>) {
+    ASSERT_NE(*exec1.impl_internal_space_instance()->m_queue,
+              *exec2.impl_internal_space_instance()->m_queue);
+  }
 #endif
-#ifdef KOKKOS_ENABLE_OPENMP
-void check_distinctive(Kokkos::OpenMP exec1, Kokkos::OpenMP exec2) {
-  ASSERT_NE(exec1, exec2);
 }
-#endif
 }  // namespace
 
 #ifdef KOKKOS_ENABLE_OPENMP
@@ -99,28 +104,6 @@ void test_partitioning(std::vector<TEST_EXECSPACE>& instances) {
       });
   ASSERT_EQ(sum1, sum2);
   ASSERT_EQ(sum1, N * (N - 1) / 2);
-
-#if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) || \
-    defined(KOKKOS_ENABLE_SYCL) || defined(KOKKOS_ENABLE_OPENMP)
-  // Eliminate unused function warning
-  // (i.e. when compiling for Serial and CUDA, during Serial compilation the
-  // Cuda overload is unused ...)
-  if (sum1 != sum2) {
-#ifdef KOKKOS_ENABLE_CUDA
-    check_distinctive(Kokkos::Cuda(), Kokkos::Cuda());
-#endif
-#ifdef KOKKOS_ENABLE_HIP
-    check_distinctive(Kokkos::HIP(), Kokkos::HIP());
-#endif
-#ifdef KOKKOS_ENABLE_SYCL
-    check_distinctive(Kokkos::Experimental::SYCL(),
-                      Kokkos::Experimental::SYCL());
-#endif
-#ifdef KOKKOS_ENABLE_OPENMP
-    check_distinctive(Kokkos::OpenMP(), Kokkos::OpenMP());
-#endif
-  }
-#endif
 }
 
 TEST(TEST_CATEGORY, partitioning_by_args) {