use execution space in SpmvMergeFlat

kokkos · Feb 28, 2023 · 76d8ed5 · 76d8ed5
1 parent 373fc5e
commit 76d8ed5
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/sparse/impl/KokkosSparse_spmv_impl.hpp b/sparse/impl/KokkosSparse_spmv_impl.hpp
@@ -207,8 +207,12 @@ struct SpmvMergeFlat {
   static void spmv(const char mode[], const y_value_type& alpha,
                    const AMatrix& A, const XVector& x, const y_value_type& beta,
                    const YVector& y) {
+
+    using execution_space = typename YVector::execution_space;
     const int numThreads = exec_space::concurrency();
-    Kokkos::RangePolicy<Kokkos::Schedule<Kokkos::Static>> policy(0, numThreads);
+
+    // TODO use the correct execution space
+    Kokkos::RangePolicy<Kokkos::Schedule<Kokkos::Static>, execution_space> policy(execution_space(), 0, numThreads);
 
     const size_t threadPathLength = (A.numRows() + A.nnz() + exec_space::concurrency() - 1) /
                                   exec_space::concurrency();
@@ -1028,7 +1032,6 @@ static void spmv_beta(const KokkosKernels::Experimental::Controls& controls,
                       const AMatrix& A, const XVector& x,
                       typename YVector::const_value_type& beta,
                       const YVector& y) {
-
   using execution_space = typename YVector::execution_space;
 
   // use merge-based algorithm if requested

diff --git a/sparse/unit_test/Test_Sparse_spmv.hpp b/sparse/unit_test/Test_Sparse_spmv.hpp
@@ -514,11 +514,11 @@ void test_spmv_algorithms(lno_t numRows, size_type nnz, lno_t bandwidth,
     test_spmv<scalar_t, lno_t, size_type, Device>(controls, numRows, nnz, bandwidth, row_size_variance, heavy);
   }
 
-  // {
-  //   Controls controls;
-  //   controls.setParameter("algorithm", "merge");
-  //   test_spmv<scalar_t, lno_t, size_type, Device>(controls, numRows, nnz, bandwidth, row_size_variance, heavy);
-  // }
+  {
+    Controls controls;
+    controls.setParameter("algorithm", "merge");
+    test_spmv<scalar_t, lno_t, size_type, Device>(controls, numRows, nnz, bandwidth, row_size_variance, heavy);
+  }
 
 }