42 #ifndef IFPACK2_DETAILS_INVERSEDIAGONALKERNEL_DEF_HPP 43 #define IFPACK2_DETAILS_INVERSEDIAGONALKERNEL_DEF_HPP 45 #include "Tpetra_CrsMatrix.hpp" 46 #include "Tpetra_MultiVector.hpp" 47 #include "Tpetra_Operator.hpp" 48 #include "Tpetra_Vector.hpp" 49 #include "Tpetra_Export_decl.hpp" 50 #include "Tpetra_Import_decl.hpp" 51 #include "Kokkos_ArithTraits.hpp" 52 #include "Teuchos_Assert.hpp" 53 #include <type_traits> 54 #include "KokkosSparse_spmv_impl.hpp" 63 template<
class DVector,
70 using execution_space =
typename AMatrix::execution_space;
71 using LO =
typename AMatrix::non_const_ordinal_type;
72 using value_type =
typename AMatrix::non_const_value_type;
73 using team_policy =
typename Kokkos::TeamPolicy<execution_space>;
74 using team_member =
typename team_policy::member_type;
75 using ATV = Kokkos::ArithTraits<value_type>;
77 using magnitude_type =
typename ATV::mag_type;
78 using MATV = Kokkos::ArithTraits<magnitude_type>;
82 DiagOffsetType m_offsets;
83 magnitude_type m_L1Eta;
84 magnitude_type m_MinDiagonalValue;
88 const DiagOffsetType& m_offsets_,
89 const magnitude_type m_L1Eta_,
90 const magnitude_type m_MinDiagonalValue_) :
93 m_offsets (m_offsets_),
95 m_MinDiagonalValue (m_MinDiagonalValue_)
97 const size_t numRows = m_A.numRows ();
99 TEUCHOS_ASSERT( numRows ==
size_t (m_d.extent (0)) );
100 TEUCHOS_ASSERT( numRows ==
size_t (m_offsets.extent (0)) );
103 KOKKOS_INLINE_FUNCTION
104 void operator() (
const LO lclRow)
const 106 const size_t INV = Tpetra::Details::OrdinalTraits<size_t>::invalid ();
107 const value_type one = ATV::one();
110 m_d(lclRow,0) = ATV::zero();
112 if (m_offsets(lclRow) != INV) {
113 auto curRow = m_A.rowConst (lclRow);
114 value_type d = curRow.value(m_offsets(lclRow));
125 const magnitude_type half = MATV::one () / (MATV::one () + MATV::one ());
126 const LO numRows =
static_cast<LO
> (m_A.numRows ());
127 const LO row_length =
static_cast<LO
> (curRow.length);
128 magnitude_type diagonal_boost = MATV::zero();
129 for (LO iEntry = 0; iEntry < row_length; iEntry++) {
130 if (curRow.colidx(iEntry) >= numRows)
131 diagonal_boost += ATV::magnitude(curRow.value(iEntry));
133 diagonal_boost *= half;
134 if (ATV::magnitude(d) < m_L1Eta * diagonal_boost)
141 if (ATV::magnitude(d) <= m_MinDiagonalValue)
142 d = m_MinDiagonalValue;
146 m_d(lclRow,0) = one / d;
155 template<
class TpetraOperatorType>
162 template<
class TpetraOperatorType>
164 InverseDiagonalKernel<TpetraOperatorType>::
165 setMatrix (
const Teuchos::RCP<const operator_type>& A)
167 if (A_op_.get () != A.get ()) {
170 using Teuchos::rcp_dynamic_cast;
171 A_crs_ = rcp_dynamic_cast<
const crs_matrix_type> (A);
173 const size_t lclNumRows = A_crs_->getRowMap ()->getNodeNumElements ();
175 if (offsets_.extent (0) < lclNumRows) {
176 using Kokkos::view_alloc;
177 using Kokkos::WithoutInitializing;
178 using offsets_view_type = Kokkos::View<size_t*, device_type>;
180 offsets_ = offsets_view_type ();
181 auto howAlloc = view_alloc (
"offsets", WithoutInitializing);
182 offsets_ = offsets_view_type (howAlloc, lclNumRows);
185 A_crs_->getCrsGraph ()->getLocalDiagOffsets (offsets_);
189 template<
class TpetraOperatorType>
191 InverseDiagonalKernel<TpetraOperatorType>::
192 compute (vector_type& D_inv,
193 bool do_l1, magnitude_type L1Eta,
194 bool fixTinyDiagEntries, magnitude_type MinDiagonalValue)
199 using d_type =
typename vector_type::dual_view_type::t_dev;
201 using d_matrix_type =
typename crs_matrix_type::local_matrix_device_type;
203 const char kernel_label[] =
"inverse_diagonal_kernel";
204 using execution_space =
typename NT::execution_space;
205 using range_type = Kokkos::RangePolicy<execution_space, LO>;
206 const size_t lclNumRows = A_crs_->getRowMap ()->getNodeNumElements ();
207 auto policy = range_type(0, lclNumRows);
209 d_type d = D_inv.getLocalViewDevice(Tpetra::Access::OverwriteAll);
210 d_matrix_type a = A_crs_->getLocalMatrixDevice();
213 constexpr
bool do_l1_template =
true;
214 if (fixTinyDiagEntries) {
215 constexpr
bool fix_tiny_template =
true;
217 Impl::InverseDiagonalWithExtraction<d_type,
222 functor_type func (d, a, offsets_, L1Eta, MinDiagonalValue);
223 Kokkos::parallel_for (kernel_label, policy, func);
225 constexpr
bool fix_tiny_template =
false;
227 Impl::InverseDiagonalWithExtraction<d_type,
232 functor_type func (d, a, offsets_, L1Eta, MinDiagonalValue);
233 Kokkos::parallel_for (kernel_label, policy, func);
236 constexpr
bool do_l1_template =
false;
237 if (fixTinyDiagEntries) {
238 constexpr
bool fix_tiny_template =
true;
240 Impl::InverseDiagonalWithExtraction<d_type,
245 functor_type func (d, a, offsets_, L1Eta, MinDiagonalValue);
246 Kokkos::parallel_for (kernel_label, policy, func);
248 constexpr
bool fix_tiny_template =
false;
250 Impl::InverseDiagonalWithExtraction<d_type,
255 functor_type func (d, a, offsets_, L1Eta, MinDiagonalValue);
256 Kokkos::parallel_for (kernel_label, policy, func);
264 #define IFPACK2_DETAILS_INVERSEDIAGONALKERNEL_INSTANT(SC,LO,GO,NT) \ 265 template class Ifpack2::Details::InverseDiagonalKernel<Tpetra::Operator<SC, LO, GO, NT> >; 267 #endif // IFPACK2_DETAILS_INVERSEDIAGONALKERNEL_DEF_HPP Ifpack2 implementation details.
Compute scaled damped residual for Chebyshev.
Definition: Ifpack2_Details_InverseDiagonalKernel_decl.hpp:77
Preconditioners and smoothers for Tpetra sparse matrices.
Definition: Ifpack2_AdditiveSchwarz_decl.hpp:73