42 #ifndef TPETRA_DETAILS_LEFTSCALELOCALCRSMATRIX_HPP
43 #define TPETRA_DETAILS_LEFTSCALELOCALCRSMATRIX_HPP
52 #include "TpetraCore_config.h"
53 #include "Kokkos_Core.hpp"
54 #include "Kokkos_ArithTraits.hpp"
55 #include <type_traits>
68 template<
class LocalSparseMatrixType,
69 class ScalingFactorsViewType,
74 typename std::remove_const<typename LocalSparseMatrixType::value_type>::type;
75 using mag_type =
typename ScalingFactorsViewType::non_const_value_type;
76 static_assert (ScalingFactorsViewType::rank == 1,
77 "scalingFactors must be a rank-1 Kokkos::View.");
78 using device_type =
typename LocalSparseMatrixType::device_type;
79 using LO =
typename LocalSparseMatrixType::ordinal_type;
80 using policy_type = Kokkos::TeamPolicy<typename device_type::execution_space, LO>;
91 const ScalingFactorsViewType& scalingFactors,
92 const bool assumeSymmetric) :
94 scalingFactors_ (scalingFactors),
95 assumeSymmetric_ (assumeSymmetric)
98 KOKKOS_INLINE_FUNCTION
void
99 operator () (
const typename policy_type::member_type & team)
const
101 using KAM = Kokkos::ArithTraits<mag_type>;
103 const LO lclRow = team.league_rank();
104 const mag_type curRowNorm = scalingFactors_(lclRow);
107 const mag_type scalingFactor = assumeSymmetric_ ?
108 KAM::sqrt (curRowNorm) : curRowNorm;
109 auto curRow = A_lcl_.row (lclRow);
110 const LO numEnt = curRow.length;
111 Kokkos::parallel_for(Kokkos::TeamThreadRange(team, numEnt), [&](
const LO k) {
113 curRow.value (k) = curRow.value(k) / scalingFactor;
116 curRow.value (k) = curRow.value(k) * scalingFactor;
122 LocalSparseMatrixType A_lcl_;
123 typename ScalingFactorsViewType::const_type scalingFactors_;
124 bool assumeSymmetric_;
140 template<
class LocalSparseMatrixType,
class ScalingFactorsViewType>
143 const ScalingFactorsViewType& scalingFactors,
144 const bool assumeSymmetric,
145 const bool divide =
true)
147 using device_type =
typename LocalSparseMatrixType::device_type;
148 using execution_space =
typename device_type::execution_space;
149 using LO =
typename LocalSparseMatrixType::ordinal_type;
150 using policy_type = Kokkos::TeamPolicy<execution_space, LO>;
152 const LO lclNumRows = A_lcl.numRows ();
156 typename ScalingFactorsViewType::const_type,
true>;
157 functor_type functor (A_lcl, scalingFactors, assumeSymmetric);
158 Kokkos::parallel_for (
"leftScaleLocalCrsMatrix",
159 policy_type (lclNumRows, Kokkos::AUTO), functor);
164 typename ScalingFactorsViewType::const_type,
false>;
165 functor_type functor (A_lcl, scalingFactors, assumeSymmetric);
166 Kokkos::parallel_for (
"leftScaleLocalCrsMatrix",
167 policy_type (lclNumRows, Kokkos::AUTO), functor);
174 #endif // TPETRA_DETAILS_LEFTSCALELOCALCRSMATRIX_HPP
void leftScaleLocalCrsMatrix(const LocalSparseMatrixType &A_lcl, const ScalingFactorsViewType &scalingFactors, const bool assumeSymmetric, const bool divide=true)
Left-scale a KokkosSparse::CrsMatrix.
Kokkos::parallel_for functor that left-scales a KokkosSparse::CrsMatrix.
LeftScaleLocalCrsMatrix(const LocalSparseMatrixType &A_lcl, const ScalingFactorsViewType &scalingFactors, const bool assumeSymmetric)