Cleanup test and match TestTeam.hpp

Use 32 as a minimum team size like in TestTeam.hpp. Better match layout in the 3 tests here.
tcclevenger · Jan 18, 2024 · 870874c · 870874c
1 parent 9beb5ea
commit 870874c
Showing 1 changed file with 26 additions and 13 deletions.
diff --git a/core/unit_test/TestReducers.hpp b/core/unit_test/TestReducers.hpp
@@ -378,29 +378,42 @@ struct TestReducers {
     Kokkos::deep_copy(sum_view, Scalar(1));
 
     // Test team policy reduction
-    constexpr int num_teams = get_num_teams();
-    TeamSumFunctor tf;
+    {
+      constexpr int num_teams = get_num_teams();
+      TeamSumFunctor tf;
+      // FIXME_OPENMPTARGET temporary restriction for team size to be at least
+      // 32
 #ifdef KOKKOS_ENABLE_OPENMPTARGET
-    auto team_pol = Kokkos::TeamPolicy<ExecSpace>(num_teams, Kokkos::AUTO);
+      int team_size =
+          std::is_same<ExecSpace, Kokkos::Experimental::OpenMPTarget>::value
+              ? 32
+              : 1;
 #else
-    auto team_pol = Kokkos::TeamPolicy<ExecSpace>(num_teams, 1);
+      int team_size         = 1;
 #endif
-    Kokkos::parallel_reduce(team_pol, tf, sum_view);
-    Kokkos::deep_copy(sum_scalar, sum_view);
-    ASSERT_EQ(sum_scalar, Scalar{num_teams}) << "num_teams: " << num_teams;
+      auto team_pol = Kokkos::TeamPolicy<ExecSpace>(num_teams, team_size);
+      Kokkos::parallel_reduce(team_pol, tf, sum_view);
+      Kokkos::deep_copy(sum_scalar, sum_view);
+      ASSERT_EQ(sum_scalar, Scalar{num_teams}) << "num_teams: " << num_teams;
+    }
 
     // Test TeamThreadRange level reduction with 0 work produces 0 result
     {
       const int league_size = 1;
       Kokkos::View<Scalar*, ExecSpace> result("result", league_size);
       TeamSumNestedFunctor tnf(f, league_size, 0, result);
-      Kokkos::parallel_for(
+      // FIXME_OPENMPTARGET temporary restriction for team size to be at least
+      // 32
 #ifdef KOKKOS_ENABLE_OPENMPTARGET
-          Kokkos::TeamPolicy<ExecSpace>(1, Kokkos::AUTO),
+      int team_size =
+          std::is_same<ExecSpace, Kokkos::Experimental::OpenMPTarget>::value
+              ? 32
+              : 1;
 #else
-          Kokkos::TeamPolicy<ExecSpace>(1, 1),
+      int team_size         = 1;
 #endif
-          tnf);
+      auto team_pol = Kokkos::TeamPolicy<ExecSpace>(1, team_size);
+      Kokkos::parallel_for(team_pol, tnf);
       auto result_h =
           Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), result);
       ASSERT_EQ(result_h(0), Scalar{0}) << "N: " << N;
@@ -424,8 +437,8 @@ struct TestReducers {
           Kokkos::TeamPolicy<ExecSpace>(league_size, initial_team_size)
               .team_size_max(tnf, Kokkos::ParallelForTag());
       auto team_size = std::min(team_size_max, TEST_EXECSPACE().concurrency());
-      Kokkos::parallel_for(
-          Kokkos::TeamPolicy<ExecSpace>(league_size, team_size), tnf);
+      auto team_pol  = Kokkos::TeamPolicy<ExecSpace>(league_size, team_size);
+      Kokkos::parallel_for(team_pol, tnf);
       auto result_h =
           Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), result);
       for (int i = 0; i < result_h.extent_int(0); ++i) {