kokkos#5635: Add test for parallel_scan with return value for ThreadV…

…ectorRange
masterleinad · Jul 28, 2023 · afbbdcc · afbbdcc
1 parent 7290d52
commit afbbdcc
Showing 1 changed file with 81 additions and 0 deletions.
diff --git a/core/unit_test/TestTeamVector.hpp b/core/unit_test/TestTeamVector.hpp
@@ -605,6 +605,71 @@ struct functor_vec_scan {
   }
 };
 
+// Temporary: This condition will progressively be reduced when parallel_scan
+// with return value will be implemented for more backends.
+#if defined(KOKKOS_ENABLE_SERIAL) || defined(KOKKOS_ENABLE_OPENMP)
+#if !defined(KOKKOS_ENABLE_CUDA) && !defined(KOKKOS_ENABLE_HIP) &&             \
+    !defined(KOKKOS_ENABLE_OPENACC) && !defined(KOKKOS_ENABLE_SYCL) &&         \
+    !defined(KOKKOS_ENABLE_THREADS) && !defined(KOKKOS_ENABLE_OPENMPTARGET) && \
+    !defined(KOKKOS_ENABLE_HPX)
+template <typename Scalar, class ExecutionSpace>
+struct functor_vec_scan_ret_val {
+  using policy_type     = Kokkos::TeamPolicy<ExecutionSpace>;
+  using execution_space = ExecutionSpace;
+  using view_type       = Kokkos::View<Scalar **, execution_space>;
+
+  Kokkos::View<int, Kokkos::LayoutLeft, ExecutionSpace> flag;
+  view_type return_values;
+  int team_size;
+
+  functor_vec_scan_ret_val(
+      Kokkos::View<int, Kokkos::LayoutLeft, ExecutionSpace> flag_, int nteams,
+      int tsize)
+      : flag(flag_), team_size(tsize) {
+    return_values = view_type("return_values", nteams, tsize);
+  }
+
+  KOKKOS_INLINE_FUNCTION
+  void operator()(typename policy_type::member_type team) const {
+    Scalar return_val;
+
+    Kokkos::parallel_scan(
+        Kokkos::ThreadVectorRange(team, 13),
+        [&](int i, Scalar &val, bool final) {
+          val += i;
+
+          if (final) {
+            Scalar test = 0;
+            for (int k = 0; k <= i; k++) test += k;
+            return_values(team.league_rank(), team.team_rank()) = test;
+
+            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;
+            }
+          }
+        },
+        return_val);
+
+    if (flag() == 0 &&
+        return_values(team.league_rank(), team.team_rank()) != return_val) {
+      Kokkos::printf("FAILED vector_scan_ret_val %i %i %lf %lf\n",
+                     team.league_rank(), team.team_rank(),
+                     static_cast<double>(
+                         return_values(team.league_rank(), team.team_rank())),
+                     static_cast<double>(return_val));
+
+      flag() = 1;
+    }
+  }
+};
+#endif
+#endif
+
 template <typename Scalar, class ExecutionSpace>
 struct functor_reduce {
   using value_type      = double;
@@ -676,6 +741,21 @@ bool test_scalar(int nteams, int team_size, int test) {
     Kokkos::parallel_for(
         "B", Kokkos::TeamPolicy<ExecutionSpace>(nteams, team_size, 8),
         functor_vec_single<Scalar, ExecutionSpace>(d_flag, 4, 13));
+  } else if (test == 12) {
+// Temporary: This condition will progressively be reduced when
+// parallel_scan
+// with return value will be implemented for more backends.
+#if defined(KOKKOS_ENABLE_SERIAL) || defined(KOKKOS_ENABLE_OPENMP)
+#if !defined(KOKKOS_ENABLE_CUDA) && !defined(KOKKOS_ENABLE_HIP) &&             \
+    !defined(KOKKOS_ENABLE_OPENACC) && !defined(KOKKOS_ENABLE_SYCL) &&         \
+    !defined(KOKKOS_ENABLE_THREADS) && !defined(KOKKOS_ENABLE_OPENMPTARGET) && \
+    !defined(KOKKOS_ENABLE_HPX)
+    Kokkos::parallel_for(
+        Kokkos::TeamPolicy<ExecutionSpace>(nteams, team_size, 8),
+        functor_vec_scan_ret_val<Scalar, ExecutionSpace>(d_flag, nteams,
+                                                         team_size));
+#endif
+#endif
   }
 
   Kokkos::deep_copy(h_flag, d_flag);
@@ -965,6 +1045,7 @@ TEST(TEST_CATEGORY, team_vector) {
   ASSERT_TRUE((TestTeamVector::Test<TEST_EXECSPACE>(9)));
   ASSERT_TRUE((TestTeamVector::Test<TEST_EXECSPACE>(10)));
   ASSERT_TRUE((TestTeamVector::Test<TEST_EXECSPACE>(11)));
+  ASSERT_TRUE((TestTeamVector::Test<TEST_EXECSPACE>(12)));
 }
 #endif