WIP: Tests (not working)

nliber · Aug 14, 2024 · 6c2e88f · 6c2e88f
1 parent f87e719
commit 6c2e88f
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 15 deletions.
diff --git a/core/unit_test/CMakeLists.txt b/core/unit_test/CMakeLists.txt
@@ -130,6 +130,7 @@ foreach(Tag Threads;Serial;OpenMP;Cuda;HPX;OpenMPTarget;OpenACC;HIP;SYCL)
     # file. That then exceeded the shell command line max length.
     set(${Tag}_SOURCES1A)
     foreach(Name
+        TeamVector
         Abort
         ArrayOps
         AtomicOperations_complexdouble

diff --git a/core/unit_test/TestTeamVector.hpp b/core/unit_test/TestTeamVector.hpp
@@ -22,40 +22,74 @@
 #include <cstdint>
 #include <cinttypes>
 #include <TestNonTrivialScalarTypes.hpp>
+#include "Kokkos_Macros.hpp"
 #include "Kokkos_ReductionIdentity.hpp"
 
 namespace TestTeamVector {
 
-struct RIInt {
+struct SumNotEqualToValueInit {
   KOKKOS_DEFAULTED_FUNCTION
-  constexpr RIInt() = default;
+  constexpr SumNotEqualToValueInit() = default;
 
   KOKKOS_FUNCTION
-  constexpr RIInt(int i_) : i{i_} {}
-  KOKKOS_FUNCTION constexpr RIInt &operator=(int i_) {
+  constexpr SumNotEqualToValueInit(int i_) : i{i_} {}
+
+  KOKKOS_FUNCTION constexpr SumNotEqualToValueInit &operator=(int i_) {
     i = i_;
     return *this;
   }
 
   KOKKOS_FUNCTION
-  friend constexpr RIInt &operator+=(RIInt &lhs, int i_) {
-    lhs.i = i_;
+  friend constexpr SumNotEqualToValueInit &operator+=(
+      SumNotEqualToValueInit &lhs, int i_) {
+    lhs.i += i_;
     return lhs;
   }
   KOKKOS_FUNCTION
-  friend constexpr RIInt &operator+=(RIInt &lhs, RIInt const &rhs) {
+  friend constexpr SumNotEqualToValueInit &operator+=(
+      SumNotEqualToValueInit &lhs, SumNotEqualToValueInit const &rhs) {
     return lhs += rhs.i;
   }
 
-  KOKKOS_FUNCTION friend constexpr bool operator==(RIInt const &lhs,
-                                                   RIInt const &rhs) {
+  KOKKOS_FUNCTION
+  friend constexpr SumNotEqualToValueInit &operator*=(
+      SumNotEqualToValueInit &lhs, int i_) {
+    lhs.i *= i_;
+    return lhs;
+  }
+  KOKKOS_FUNCTION
+  friend constexpr SumNotEqualToValueInit &operator*=(
+      SumNotEqualToValueInit &lhs, SumNotEqualToValueInit const &rhs) {
+    return lhs *= rhs.i;
+  }
+
+  KOKKOS_FUNCTION friend constexpr bool operator==(
+      SumNotEqualToValueInit const &lhs, SumNotEqualToValueInit const &rhs) {
     return lhs.i == rhs.i;
   }
-  KOKKOS_FUNCTION friend constexpr bool operator!=(RIInt const &lhs,
-                                                   RIInt const &rhs) {
+  KOKKOS_FUNCTION friend constexpr bool operator!=(
+      SumNotEqualToValueInit const &lhs, SumNotEqualToValueInit const &rhs) {
     return !(lhs == rhs);
   }
 
+  KOKKOS_FUNCTION friend constexpr bool operator==(
+      SumNotEqualToValueInit const &lhs, int i_) {
+    return lhs.i == i_;
+  }
+  KOKKOS_FUNCTION friend constexpr bool operator!=(
+      SumNotEqualToValueInit const &lhs, int i_) {
+    return !(lhs == i_);
+  }
+
+  KOKKOS_FUNCTION friend constexpr bool operator==(
+      int i_, SumNotEqualToValueInit const &rhs) {
+    return i_ == rhs.i;
+  }
+  KOKKOS_FUNCTION friend constexpr bool operator!=(
+      int i_, SumNotEqualToValueInit const &rhs) {
+    return !(i_ == rhs);
+  }
+
   KOKKOS_FUNCTION constexpr int get() const { return i; }
   KOKKOS_FUNCTION constexpr operator int() const { return get(); }
 
@@ -65,14 +99,16 @@ struct RIInt {
 }  // namespace TestTeamVector
 
 template <>
-struct Kokkos::reduction_identity<TestTeamVector::RIInt> {
-  using value_type = TestTeamVector::RIInt;
+struct Kokkos::reduction_identity<TestTeamVector::SumNotEqualToValueInit> {
+  using value_type = TestTeamVector::SumNotEqualToValueInit;
 
   KOKKOS_FUNCTION constexpr static value_type sum() { return value_type{0}; }
+  KOKKOS_FUNCTION constexpr static value_type prod() { return value_type{1}; }
 };
 
-static_assert(Kokkos::reduction_identity<TestTeamVector::RIInt>::sum() !=
-              TestTeamVector::RIInt{});
+// static_assert(
+// Kokkos::reduction_identity<TestTeamVector::SumNotEqualToValueInit>::sum() !=
+// TestTeamVector::SumNotEqualToValueInit{});
 
 namespace TestTeamVector {
 
@@ -656,6 +692,41 @@ struct functor_vec_scan {
   }
 };
 
+// Tests specifically that reduction_identity<Scalar>::sum() is
+// used inside parallel_scan
+template <typename Scalar, class ExecutionSpace>
+struct functor_vec_scan_sum {
+  using policy_type     = Kokkos::TeamPolicy<ExecutionSpace>;
+  using execution_space = ExecutionSpace;
+
+  Kokkos::View<int, Kokkos::LayoutLeft, ExecutionSpace> flag;
+  functor_vec_scan_sum(
+      Kokkos::View<int, Kokkos::LayoutLeft, ExecutionSpace> flag_)
+      : flag(flag_) {}
+
+  KOKKOS_FUNCTION
+  void operator()(typename policy_type::member_type team) const {
+    Kokkos::parallel_scan(Kokkos::ThreadVectorRange(team, 13), [&](int i,
+                                                                   Scalar &val,
+                                                                   bool final) {
+      val += i;
+
+      if (final) {
+        Scalar test = Kokkos::reduction_identity<Scalar>::sum();
+        for (int k = 0; k <= i; k++) test += k;
+
+        if (test != val) {
+          Kokkos::printf("FAILED vector_par_scan %i %i %lf %lf\n",
+                         team.league_rank(), team.team_rank(),
+                         static_cast<double>(test), static_cast<double>(val));
+
+          flag() = 1;
+        }
+      }
+    });
+  }
+};
+
 // Temporary: This condition will progressively be reduced when parallel_scan
 // with return value will be implemented for more backends.
 #if !defined(KOKKOS_ENABLE_OPENACC)
@@ -790,6 +861,10 @@ bool test_scalar(int nteams, int team_size, int test) {
         Kokkos::TeamPolicy<ExecutionSpace>(nteams, team_size, 8),
         functor_vec_scan_ret_val<Scalar, ExecutionSpace>(d_flag, team_size));
 #endif
+  } else if (test == 13) {
+    Kokkos::parallel_for(
+        Kokkos::TeamPolicy<ExecutionSpace>(nteams, team_size, 8),
+        functor_vec_scan_sum<Scalar, ExecutionSpace>(d_flag));
   }
 
   Kokkos::deep_copy(h_flag, d_flag);
@@ -826,6 +901,23 @@ bool Test(int test) {
   return passed;
 }
 
+template <class ExecutionSpace>
+bool TestParallelScanUsesSum() {
+  bool passed = true;
+// With SYCL 33*8 exceeds the maximum work group size
+#ifdef KOKKOS_ENABLE_SYCL
+  int team_size = 31;
+#else
+  int team_size = 33;
+#endif
+  int const concurrency = ExecutionSpace().concurrency();
+  if (team_size > concurrency) team_size = concurrency;
+  passed = passed && test_scalar<SumNotEqualToValueInit, ExecutionSpace>(
+                         317, team_size, 13);
+
+  return passed;
+}
+
 }  // namespace TestTeamVector
 
 namespace Test {
@@ -1077,6 +1169,9 @@ TEST(TEST_CATEGORY, team_vector) {
   ASSERT_TRUE((TestTeamVector::Test<TEST_EXECSPACE>(10)));
   ASSERT_TRUE((TestTeamVector::Test<TEST_EXECSPACE>(11)));
   ASSERT_TRUE((TestTeamVector::Test<TEST_EXECSPACE>(12)));
+  ASSERT_TRUE((TestTeamVector::Test<TEST_EXECSPACE>(13)));
+
+  ASSERT_TRUE((TestTeamVector::TestParallelScanUsesSum<TEST_EXECSPACE>()));
 }
 
 TEST(TEST_CATEGORY, triple_nested_parallelism) {