40 #ifndef TPETRA_MATRIX_APPLY_HELPER_HPP
41 #define TPETRA_MATRIX_APPLY_HELPER_HPP
43 #include "Kokkos_Core.hpp"
44 #include "KokkosSparse_spmv_handle.hpp"
46 #if KOKKOSKERNELS_VERSION < 40299
47 #error "Tpetra::Details::MatrixApplyHelper can only be included with KokkosKernels 4.3+"
60 template<
typename LocalMatrix,
typename IntLocalMatrix,
typename MultiVectorLocalView>
63 using Rowptrs =
typename LocalMatrix::row_map_type;
64 using IntRowptrs =
typename IntLocalMatrix::row_map_type;
65 using XVectorType =
typename MultiVectorLocalView::const_type;
66 using YVectorType = MultiVectorLocalView;
67 using SPMVHandle = KokkosSparse::SPMVHandle<typename LocalMatrix::device_type, LocalMatrix, XVectorType, YVectorType>;
68 using SPMVHandleInt = KokkosSparse::SPMVHandle<typename LocalMatrix::device_type, IntLocalMatrix, XVectorType, YVectorType>;
70 MatrixApplyHelper(
size_t nnz_,
const typename LocalMatrix::row_map_type& rowptrs, KokkosSparse::SPMVAlgorithm algo = KokkosSparse::SPMV_DEFAULT)
71 : nnz(nnz_), handle(algo), handle_int(algo)
73 if(shouldUseIntRowptrs())
74 fillRowptrsInt(rowptrs);
78 bool shouldUseIntRowptrs()
const
83 return nnz <= size_t(INT_MAX) && !std::is_same_v<int, typename Rowptrs::non_const_value_type>;
87 IntLocalMatrix getIntRowptrMatrix(
const LocalMatrix& A)
89 if constexpr(KokkosSparse::is_crs_matrix_v<LocalMatrix>)
91 return IntLocalMatrix(
"", A.numRows(), A.numCols(), A.nnz(),
92 A.values, rowptrs_int, A.graph.entries);
96 return IntLocalMatrix(
"", A.numRows(), A.numCols(), A.nnz(),
97 A.values, rowptrs_int, A.graph.entries, A.blockDim());
101 void fillRowptrsInt(
const Rowptrs& rowptrs)
105 typename IntRowptrs::non_const_type rowptrs_temp(
106 Kokkos::view_alloc(Kokkos::WithoutInitializing,
"rowptrs_int"), rowptrs.extent(0));
107 Kokkos::parallel_for(Kokkos::RangePolicy<typename LocalMatrix::execution_space>(0, rowptrs.extent(0)),
110 rowptrs_temp(i) = int(rowptrs(i));
112 rowptrs_int = rowptrs_temp;
119 SPMVHandleInt handle_int;
120 IntRowptrs rowptrs_int;
125 #endif // TPETRA_MATRIX_APPLY_HELPER_HPP