Merge pull request #12190 from brian-kelley/FixKK1959

Stokhos: update for new KokkosSparse::spmv overloads
trilinos · Sep 1, 2023 · 73dcc21 · 73dcc21
2 parents 004eee5 + db1f57b
commit 73dcc21
Show file tree

Hide file tree

Showing 2 changed files with 94 additions and 0 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
@@ -1619,6 +1619,53 @@ spmv(
   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
@@ -749,6 +749,53 @@ spmv(
   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 */