From eab6d5d32a7ec09901cdbc26b590c59312c2bd84 Mon Sep 17 00:00:00 2001
From: Brian Kelley <bmkelle@sandia.gov>
Date: Tue, 5 Sep 2023 17:58:54 -0600
Subject: [PATCH] Stokhos: fix KokkosKernels #1959

PR #12190 actually failed to fix KokkosSparse::spmv for Sacado scalar
types, when building with Kokkos/KokkosKernels develop branch.

This actually fixes that issue (tested with develop and master KokkosKernels)
and is quite a bit cleaner (though it uses version macros that can be
taken out for the 4.2 release).
---
 .../pce/linalg/Kokkos_CrsMatrix_UQ_PCE.hpp    | 131 +++++-------------
 .../linalg/Kokkos_CrsMatrix_MP_Vector.hpp     | 131 +++++-------------
 2 files changed, 68 insertions(+), 194 deletions(-)

diff --git a/packages/stokhos/src/sacado/kokkos/pce/linalg/Kokkos_CrsMatrix_UQ_PCE.hpp b/packages/stokhos/src/sacado/kokkos/pce/linalg/Kokkos_CrsMatrix_UQ_PCE.hpp
index 714b05c5e17e..6683ee7689bc 100644
--- a/packages/stokhos/src/sacado/kokkos/pce/linalg/Kokkos_CrsMatrix_UQ_PCE.hpp
+++ b/packages/stokhos/src/sacado/kokkos/pce/linalg/Kokkos_CrsMatrix_UQ_PCE.hpp
@@ -1467,7 +1467,11 @@ class MeanMultiply< KokkosSparse::CrsMatrix< Sacado::UQ::PCE<MatrixStorage>,
 
 namespace KokkosSparse {
 
-template <typename AlphaType,
+template <
+#if KOKKOSKERNELS_VERSION >= 40199
+          typename ExecutionSpace,
+#endif
+          typename AlphaType,
           typename BetaType,
           typename MatrixType,
           typename InputType,
@@ -1479,6 +1483,10 @@ typename std::enable_if<
   Kokkos::is_view_uq_pce< Kokkos::View< OutputType, OutputP... > >::value
   >::type
 spmv(
+#if KOKKOSKERNELS_VERSION >= 40199
+  const ExecutionSpace& space,
+#endif
+  KokkosKernels::Experimental::Controls,
   const char mode[],
   const AlphaType& a,
   const MatrixType& A,
@@ -1494,6 +1502,12 @@ spmv(
   typedef Stokhos::MeanMultiply<MatrixType, typename InputVectorType::const_type,
                                 OutputVectorType> mean_multiply_type;
 
+#if KOKKOSKERNELS_VERSION >= 40199
+  if(space != ExecutionSpace()) {
+    Kokkos::Impl::raise_error(
+      "Stokhos spmv not implemented for non-default execution space instance");
+  }
+#endif
   if(mode[0]!='N') {
     Kokkos::Impl::raise_error(
       "Stokhos spmv not implemented for transposed or conjugated matrix-vector multiplies");
@@ -1514,7 +1528,11 @@ spmv(
                           Sacado::Value<BetaType>::eval(b) );
 }
 
-template <typename AlphaType,
+template <
+#if KOKKOSKERNELS_VERSION >= 40199
+          typename ExecutionSpace,
+#endif
+          typename AlphaType,
           typename BetaType,
           typename MatrixType,
           typename InputType,
@@ -1526,30 +1544,10 @@ typename std::enable_if<
   Kokkos::is_view_uq_pce< Kokkos::View< OutputType, OutputP... > >::value
   >::type
 spmv(
+#if KOKKOSKERNELS_VERSION >= 40199
+  const ExecutionSpace& space,
+#endif
   KokkosKernels::Experimental::Controls,
-  const char mode[],
-  const AlphaType& a,
-  const MatrixType& A,
-  const Kokkos::View< InputType, InputP... >& x,
-  const BetaType& b,
-  const Kokkos::View< OutputType, OutputP... >& y,
-  const RANK_ONE)
-{
-  spmv(mode, a, A, x, b, y, RANK_ONE());
-}
-
-template <typename AlphaType,
-          typename BetaType,
-          typename MatrixType,
-          typename InputType,
-          typename ... InputP,
-          typename OutputType,
-          typename ... OutputP>
-typename std::enable_if<
-  Kokkos::is_view_uq_pce< Kokkos::View< InputType, InputP... > >::value &&
-  Kokkos::is_view_uq_pce< Kokkos::View< OutputType, OutputP... > >::value
-  >::type
-spmv(
   const char mode[],
   const AlphaType& a,
   const MatrixType& A,
@@ -1558,6 +1556,12 @@ spmv(
   const Kokkos::View< OutputType, OutputP... >& y,
   const RANK_TWO)
 {
+#if KOKKOSKERNELS_VERSION >= 40199
+  if(space != ExecutionSpace()) {
+    Kokkos::Impl::raise_error(
+      "Stokhos spmv not implemented for non-default execution space instance");
+  }
+#endif
   if(mode[0]!='N') {
     Kokkos::Impl::raise_error(
       "Stokhos spmv not implemented for transposed or conjugated matrix-vector multiplies");
@@ -1565,7 +1569,11 @@ spmv(
   if (y.extent(1) == 1) {
     auto y_1D = subview(y, Kokkos::ALL(), 0);
     auto x_1D = subview(x, Kokkos::ALL(), 0);
-    spmv(mode, a, A, x_1D, b, y_1D, RANK_ONE());
+#if KOKKOSKERNELS_VERSION >= 40199
+    spmv(space, KokkosKernels::Experimental::Controls(), mode, a, A, x_1D, b, y_1D, RANK_ONE());
+#else
+    spmv(KokkosKernels::Experimental::Controls(), mode, a, A, x_1D, b, y_1D, RANK_ONE());
+#endif
   }
   else {
     typedef Kokkos::View< OutputType, OutputP... > OutputVectorType;
@@ -1595,77 +1603,6 @@ spmv(
   }
 }
 
-template <typename AlphaType,
-          typename BetaType,
-          typename MatrixType,
-          typename InputType,
-          typename ... InputP,
-          typename OutputType,
-          typename ... OutputP>
-typename std::enable_if<
-  Kokkos::is_view_uq_pce< Kokkos::View< InputType, InputP... > >::value &&
-  Kokkos::is_view_uq_pce< Kokkos::View< OutputType, OutputP... > >::value
-  >::type
-spmv(
-  KokkosKernels::Experimental::Controls,
-  const char mode[],
-  const AlphaType& a,
-  const MatrixType& A,
-  const Kokkos::View< InputType, InputP... >& x,
-  const BetaType& b,
-  const Kokkos::View< OutputType, OutputP... >& y,
-  const RANK_TWO)
-{
-  spmv(mode, a, A, x, b, y, RANK_TWO());
-}
-
-template <typename AlphaType,
-          typename BetaType,
-          typename MatrixType,
-          typename InputType,
-          typename ... InputP,
-          typename OutputType,
-          typename ... OutputP>
-std::enable_if_t<
-  Kokkos::is_view_uq_pce< Kokkos::View< InputType, InputP... > >::value &&
-  Kokkos::is_view_uq_pce< Kokkos::View< OutputType, OutputP... > >::value &&
-  KokkosSparse::is_crs_matrix<MatrixType>::value>
-spmv(
-  KokkosKernels::Experimental::Controls controls,
-  const char mode[],
-  const AlphaType& a,
-  const MatrixType& A,
-  const Kokkos::View< InputType, InputP... >& x,
-  const BetaType& b,
-  const Kokkos::View< OutputType, OutputP... >& y)
-{
-  using RANK_SPECIALISE = std::conditional_t<std::decay_t<decltype(x)>::rank == 2, RANK_TWO, RANK_ONE>;
-  spmv(controls, mode, a, A, x, b, y, RANK_SPECIALISE());
-}
-
-template <typename AlphaType,
-          typename BetaType,
-          typename MatrixType,
-          typename InputType,
-          typename ... InputP,
-          typename OutputType,
-          typename ... OutputP>
-std::enable_if_t<
-  Kokkos::is_view_uq_pce< Kokkos::View< InputType, InputP... > >::value &&
-  Kokkos::is_view_uq_pce< Kokkos::View< OutputType, OutputP... > >::value &&
-  KokkosSparse::is_crs_matrix<MatrixType>::value>
-spmv(
-  const char mode[],
-  const AlphaType& a,
-  const MatrixType& A,
-  const Kokkos::View< InputType, InputP... >& x,
-  const BetaType& b,
-  const Kokkos::View< OutputType, OutputP... >& y)
-{
-  using RANK_SPECIALISE = std::conditional_t<std::decay_t<decltype(x)>::rank == 2, RANK_TWO, RANK_ONE>;
-  spmv(mode, a, A, x, b, y, RANK_SPECIALISE());
-}
-
 }
 
 #endif /* #ifndef KOKKOS_CRSMATRIX_UQ_PCE_HPP */
diff --git a/packages/stokhos/src/sacado/kokkos/vector/linalg/Kokkos_CrsMatrix_MP_Vector.hpp b/packages/stokhos/src/sacado/kokkos/vector/linalg/Kokkos_CrsMatrix_MP_Vector.hpp
index eb6aa8239653..e1c828a60740 100644
--- a/packages/stokhos/src/sacado/kokkos/vector/linalg/Kokkos_CrsMatrix_MP_Vector.hpp
+++ b/packages/stokhos/src/sacado/kokkos/vector/linalg/Kokkos_CrsMatrix_MP_Vector.hpp
@@ -529,7 +529,11 @@ class Multiply< KokkosSparse::CrsMatrix< Sacado::MP::Vector<MatrixStorage>,
 
 namespace KokkosSparse {
 
-template <typename AlphaType,
+template <
+#if KOKKOSKERNELS_VERSION >= 40199
+          typename ExecutionSpace,
+#endif
+          typename AlphaType,
           typename BetaType,
           typename MatrixType,
           typename InputType,
@@ -541,6 +545,10 @@ typename std::enable_if<
   Kokkos::is_view_mp_vector< Kokkos::View< OutputType, OutputP... > >::value
   >::type
 spmv(
+#if KOKKOSKERNELS_VERSION >= 40199
+  const ExecutionSpace& space,
+#endif
+  KokkosKernels::Experimental::Controls,
   const char mode[],
   const AlphaType& a,
   const MatrixType& A,
@@ -554,6 +562,12 @@ spmv(
   using input_vector_type = const_type_t<InputVectorType>;
   typedef typename InputVectorType::array_type::non_const_value_type value_type;
 
+#if KOKKOSKERNELS_VERSION >= 40199
+  if(space != ExecutionSpace()) {
+    Kokkos::Impl::raise_error(
+      "Stokhos spmv not implemented for non-default execution space instance");
+  }
+#endif
   if(mode[0]!='N') {
     Kokkos::Impl::raise_error(
       "Stokhos spmv not implemented for transposed or conjugated matrix-vector multiplies");
@@ -612,7 +626,11 @@ spmv(
   }
 }
 
-template <typename AlphaType,
+template <
+#if KOKKOSKERNELS_VERSION >= 40199
+          typename ExecutionSpace,
+#endif
+          typename AlphaType,
           typename BetaType,
           typename MatrixType,
           typename InputType,
@@ -624,30 +642,10 @@ typename std::enable_if<
   Kokkos::is_view_mp_vector< Kokkos::View< OutputType, OutputP... > >::value
   >::type
 spmv(
+#if KOKKOSKERNELS_VERSION >= 40199
+  const ExecutionSpace& space,
+#endif
   KokkosKernels::Experimental::Controls,
-  const char mode[],
-  const AlphaType& a,
-  const MatrixType& A,
-  const Kokkos::View< InputType, InputP... >& x,
-  const BetaType& b,
-  const Kokkos::View< OutputType, OutputP... >& y,
-  const RANK_ONE)
-{
-  spmv(mode, a, A, x, b, y, RANK_ONE());
-}
-
-template <typename AlphaType,
-          typename BetaType,
-          typename MatrixType,
-          typename InputType,
-          typename ... InputP,
-          typename OutputType,
-          typename ... OutputP>
-typename std::enable_if<
-  Kokkos::is_view_mp_vector< Kokkos::View< InputType, InputP... > >::value &&
-  Kokkos::is_view_mp_vector< Kokkos::View< OutputType, OutputP... > >::value
-  >::type
-spmv(
   const char mode[],
   const AlphaType& a,
   const MatrixType& A,
@@ -656,6 +654,12 @@ spmv(
   const Kokkos::View< OutputType, OutputP... >& y,
   const RANK_TWO)
 {
+#if KOKKOSKERNELS_VERSION >= 40199
+  if(space != ExecutionSpace()) {
+    Kokkos::Impl::raise_error(
+      "Stokhos spmv not implemented for non-default execution space instance");
+  }
+#endif
   if(mode[0]!='N') {
     Kokkos::Impl::raise_error(
       "Stokhos spmv not implemented for transposed or conjugated matrix-vector multiplies");
@@ -663,7 +667,11 @@ spmv(
   if (y.extent(1) == 1) {
     auto y_1D = subview(y, Kokkos::ALL(), 0);
     auto x_1D = subview(x, Kokkos::ALL(), 0);
-    spmv(mode, a, A, x_1D, b, y_1D, RANK_ONE());
+#if KOKKOSKERNELS_VERSION >= 40199
+    spmv(space, KokkosKernels::Experimental::Controls(), mode, a, A, x_1D, b, y_1D, RANK_ONE());
+#else
+    spmv(KokkosKernels::Experimental::Controls(), mode, a, A, x_1D, b, y_1D, RANK_ONE());
+#endif
   }
   else {
     typedef Kokkos::View< OutputType, OutputP... > OutputVectorType;
@@ -725,77 +733,6 @@ spmv(
   }
 }
 
-template <typename AlphaType,
-          typename BetaType,
-          typename MatrixType,
-          typename InputType,
-          typename ... InputP,
-          typename OutputType,
-          typename ... OutputP>
-typename std::enable_if<
-  Kokkos::is_view_mp_vector< Kokkos::View< InputType, InputP... > >::value &&
-  Kokkos::is_view_mp_vector< Kokkos::View< OutputType, OutputP... > >::value
-  >::type
-spmv(
-  KokkosKernels::Experimental::Controls,
-  const char mode[],
-  const AlphaType& a,
-  const MatrixType& A,
-  const Kokkos::View< InputType, InputP... >& x,
-  const BetaType& b,
-  const Kokkos::View< OutputType, OutputP... >& y,
-  const RANK_TWO)
-{
-  spmv(mode, a, A, x, b, y, RANK_TWO());
-}
-
-template <typename AlphaType,
-          typename BetaType,
-          typename MatrixType,
-          typename InputType,
-          typename ... InputP,
-          typename OutputType,
-          typename ... OutputP>
-std::enable_if_t<
-  Kokkos::is_view_mp_vector< Kokkos::View< InputType, InputP... > >::value &&
-  Kokkos::is_view_mp_vector< Kokkos::View< OutputType, OutputP... > >::value && 
-  KokkosSparse::is_crs_matrix<MatrixType>::value>
-spmv(
-  KokkosKernels::Experimental::Controls controls,
-  const char mode[],
-  const AlphaType& a,
-  const MatrixType& A,
-  const Kokkos::View< InputType, InputP... >& x,
-  const BetaType& b,
-  const Kokkos::View< OutputType, OutputP... >& y)
-{
-  using RANK_SPECIALISE = std::conditional_t<std::decay_t<decltype(x)>::rank == 2, RANK_TWO, RANK_ONE>;
-  spmv(controls, mode, a, A, x, b, y, RANK_SPECIALISE());
-}
-
-template <typename AlphaType,
-          typename BetaType,
-          typename MatrixType,
-          typename InputType,
-          typename ... InputP,
-          typename OutputType,
-          typename ... OutputP>
-std::enable_if_t<
-  Kokkos::is_view_mp_vector< Kokkos::View< InputType, InputP... > >::value &&
-  Kokkos::is_view_mp_vector< Kokkos::View< OutputType, OutputP... > >::value && 
-  KokkosSparse::is_crs_matrix<MatrixType>::value>
-spmv(
-  const char mode[],
-  const AlphaType& a,
-  const MatrixType& A,
-  const Kokkos::View< InputType, InputP... >& x,
-  const BetaType& b,
-  const Kokkos::View< OutputType, OutputP... >& y)
-{
-  using RANK_SPECIALISE = std::conditional_t<std::decay_t<decltype(x)>::rank == 2, RANK_TWO, RANK_ONE>;
-  spmv(mode, a, A, x, b, y, RANK_SPECIALISE());
-}
-
 }
 
 #endif /* #ifndef KOKKOS_CRSMATRIX_MP_VECTOR_HPP */