From 2aaa14e77d82496bdc625cf1175a6ea06f155422 Mon Sep 17 00:00:00 2001
From: Carl Pearson <cwpears@sandia.gov>
Date: Thu, 9 May 2024 10:23:30 -0600
Subject: [PATCH] Run unit tests in correct execution space

---
 sparse/unit_test/Test_Sparse_spmv.hpp     |  9 +++++++--
 sparse/unit_test/Test_Sparse_spmv_bsr.hpp | 12 ++++++++----
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/sparse/unit_test/Test_Sparse_spmv.hpp b/sparse/unit_test/Test_Sparse_spmv.hpp
index d69b75acfc..79875e6c37 100644
--- a/sparse/unit_test/Test_Sparse_spmv.hpp
+++ b/sparse/unit_test/Test_Sparse_spmv.hpp
@@ -416,6 +416,11 @@ void test_spmv(KokkosSparse::SPMVAlgorithm algo, lno_t numRows, size_type nnz,
   using handle_t =
       KokkosSparse::SPMVHandle<Device, crsMat_t, x_vector_type, y_vector_type>;
 
+  using ExecSpace     = typename crsMat_t::execution_space;
+  using my_exec_space = Kokkos::RangePolicy<ExecSpace>;
+  using y_policy = Kokkos::RangePolicy<typename y_vector_type::execution_space>;
+  using x_policy = Kokkos::RangePolicy<typename x_vector_type::execution_space>;
+
   constexpr mag_t max_x   = static_cast<mag_t>(1);
   constexpr mag_t max_y   = static_cast<mag_t>(1);
   constexpr mag_t max_val = static_cast<mag_t>(1);
@@ -448,13 +453,13 @@ void test_spmv(KokkosSparse::SPMVAlgorithm algo, lno_t numRows, size_type nnz,
   Kokkos::deep_copy(input_y_nans, input_y);
   Kokkos::deep_copy(input_yt_nans, input_yt);
   Kokkos::parallel_for(
-      input_y_nans.extent(0), KOKKOS_LAMBDA(const size_t i) {
+      y_policy(0, input_y_nans.extent(0)), KOKKOS_LAMBDA(const size_t i) {
         if (0 == (i % 19)) {
           input_y_nans(i) = NAN;
         }
       });
   Kokkos::parallel_for(
-      input_yt_nans.extent(0), KOKKOS_LAMBDA(const size_t i) {
+      y_policy(0, input_yt_nans.extent(0)), KOKKOS_LAMBDA(const size_t i) {
         if (0 == (i % 23)) {
           input_yt_nans(i) = Kokkos::nan("");
         }
diff --git a/sparse/unit_test/Test_Sparse_spmv_bsr.hpp b/sparse/unit_test/Test_Sparse_spmv_bsr.hpp
index f33c253adc..273d6f4636 100644
--- a/sparse/unit_test/Test_Sparse_spmv_bsr.hpp
+++ b/sparse/unit_test/Test_Sparse_spmv_bsr.hpp
@@ -334,6 +334,8 @@ auto random_vecs_for_spmv(const char *mode, const Bsr &a,
   using scalar_type     = typename Bsr::non_const_value_type;
   using vector_type     = typename VectorTypeFor<Bsr>::type;
   using execution_space = typename Bsr::execution_space;
+  using policy_type =
+      Kokkos::RangePolicy<typename vector_type::execution_space>;
 
   size_t nx = a.numCols() * a.blockDim();
   size_t ny = a.numRows() * a.blockDim();
@@ -349,13 +351,13 @@ auto random_vecs_for_spmv(const char *mode, const Bsr &a,
 
   if (nans) {
     Kokkos::parallel_for(
-        x.extent(0), KOKKOS_LAMBDA(size_t i) {
+        policy_type(0, x.extent(0)), KOKKOS_LAMBDA(size_t i) {
           if (0 == (i % 17)) {
             x(i) = Kokkos::nan("");
           }
         });
     Kokkos::parallel_for(
-        y.extent(0), KOKKOS_LAMBDA(size_t i) {
+        policy_type(0, y.extent(0)), KOKKOS_LAMBDA(size_t i) {
           if (0 == (i % 17)) {
             y(i) = Kokkos::nan("");
           }
@@ -592,6 +594,8 @@ auto random_multivecs_for_spm_mv(const char *mode, const Bsr &a,
   using scalar_type     = typename Bsr::non_const_value_type;
   using vector_type     = typename MultiVectorTypeFor<Layout, Bsr>::type;
   using execution_space = typename Bsr::execution_space;
+  using policy_type =
+      Kokkos::RangePolicy<typename vector_type::execution_space>;
 
   size_t nx = a.numCols() * a.blockDim();
   size_t ny = a.numRows() * a.blockDim();
@@ -608,7 +612,7 @@ auto random_multivecs_for_spm_mv(const char *mode, const Bsr &a,
   // sprinkle some "random" NaNs in
   if (nans) {
     Kokkos::parallel_for(
-        x.extent(0), KOKKOS_LAMBDA(size_t i) {
+        policy_type(0, x.extent(0)), KOKKOS_LAMBDA(size_t i) {
           for (size_t j = 0; j < x.extent(1); ++j) {
             if (0 == ((i * x.extent(1) + j) % 13)) {
               x(i, j) = Kokkos::nan("");
@@ -616,7 +620,7 @@ auto random_multivecs_for_spm_mv(const char *mode, const Bsr &a,
           }
         });
     Kokkos::parallel_for(
-        y.extent(0), KOKKOS_LAMBDA(size_t i) {
+        policy_type(0, y.extent(0)), KOKKOS_LAMBDA(size_t i) {
           for (size_t j = 0; j < y.extent(1); ++j) {
             if (0 == ((i * y.extent(1) + j) % 17)) {
               y(i, j) = Kokkos::nan("");