Add test

masterleinad · Apr 18, 2024 · 34d0db2 · 34d0db2
1 parent e7b486f
commit 34d0db2
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/core/unit_test/TestExecSpacePartitioning.hpp b/core/unit_test/TestExecSpacePartitioning.hpp
@@ -28,6 +28,17 @@ struct SumFunctor {
   void operator()(int i, int& lsum) const { lsum += i; }
 };
 
+template <class ExecSpace>
+void check_space_member_for_policies(const ExecSpace& exec) {
+  Kokkos::RangePolicy<ExecSpace> range_policy(exec, 0, 1);
+  ASSERT_EQ(range_policy.space(), exec);
+  Kokkos::MDRangePolicy<ExecSpace, Kokkos::Rank<2>> mdrange_policy(exec, {0, 0},
+                                                                   {1, 1});
+  ASSERT_EQ(mdrange_policy.space(), exec);
+  Kokkos::TeamPolicy<ExecSpace> team_policy(exec, 1, 1);
+  ASSERT_EQ(team_policy.space(), exec);
+}
+
 template <class ExecSpace>
 void check_distinctive([[maybe_unused]] ExecSpace exec1,
                        [[maybe_unused]] ExecSpace exec2) {
@@ -89,6 +100,9 @@ void run_threaded_test(const Lambda1 l1, const Lambda2 l2) {
 
 void test_partitioning(std::vector<TEST_EXECSPACE>& instances) {
   check_distinctive(instances[0], instances[1]);
+  check_space_member_for_policies(instances[0]);
+  check_space_member_for_policies(instances[1]);
+
   int sum1, sum2;
   int N = 3910;
   run_threaded_test(