10 #ifndef IFPACK2_DETAILS_INVERSEDIAGONALKERNEL_DEF_HPP
11 #define IFPACK2_DETAILS_INVERSEDIAGONALKERNEL_DEF_HPP
13 #include "Tpetra_CrsMatrix.hpp"
14 #include "Tpetra_MultiVector.hpp"
15 #include "Tpetra_Operator.hpp"
16 #include "Tpetra_Vector.hpp"
17 #include "Tpetra_Export_decl.hpp"
18 #include "Tpetra_Import_decl.hpp"
19 #include "Kokkos_ArithTraits.hpp"
20 #include "Teuchos_Assert.hpp"
21 #include <type_traits>
22 #include "KokkosSparse_spmv_impl.hpp"
31 template <
class DVector,
37 using execution_space =
typename AMatrix::execution_space;
38 using LO =
typename AMatrix::non_const_ordinal_type;
39 using value_type =
typename AMatrix::non_const_value_type;
40 using team_policy =
typename Kokkos::TeamPolicy<execution_space>;
41 using team_member =
typename team_policy::member_type;
42 using ATV = Kokkos::ArithTraits<value_type>;
44 using magnitude_type =
typename ATV::mag_type;
45 using MATV = Kokkos::ArithTraits<magnitude_type>;
49 DiagOffsetType m_offsets;
50 magnitude_type m_L1Eta;
51 magnitude_type m_MinDiagonalValue;
55 const DiagOffsetType& m_offsets_,
56 const magnitude_type m_L1Eta_,
57 const magnitude_type m_MinDiagonalValue_)
60 , m_offsets(m_offsets_)
62 , m_MinDiagonalValue(m_MinDiagonalValue_) {
63 const size_t numRows = m_A.numRows();
69 KOKKOS_INLINE_FUNCTION
70 void operator()(
const LO lclRow)
const {
71 const size_t INV = Tpetra::Details::OrdinalTraits<size_t>::invalid();
72 const value_type one = ATV::one();
75 m_d(lclRow, 0) = ATV::zero();
77 if (m_offsets(lclRow) != INV) {
78 auto curRow = m_A.rowConst(lclRow);
79 value_type d = curRow.value(m_offsets(lclRow));
90 const magnitude_type half = MATV::one() / (MATV::one() + MATV::one());
91 const LO numRows =
static_cast<LO
>(m_A.numRows());
92 const LO row_length =
static_cast<LO
>(curRow.length);
93 magnitude_type diagonal_boost = MATV::zero();
94 for (LO iEntry = 0; iEntry < row_length; iEntry++) {
95 if (curRow.colidx(iEntry) >= numRows)
96 diagonal_boost += ATV::magnitude(curRow.value(iEntry));
98 diagonal_boost *= half;
99 if (ATV::magnitude(d) < m_L1Eta * diagonal_boost)
106 if (ATV::magnitude(d) <= m_MinDiagonalValue)
107 d = m_MinDiagonalValue;
111 m_d(lclRow, 0) = one / d;
118 template <
class TpetraOperatorType>
124 template <
class TpetraOperatorType>
125 void InverseDiagonalKernel<TpetraOperatorType>::
127 if (A_op_.get() != A.
get()) {
130 using Teuchos::rcp_dynamic_cast;
131 A_crs_ = rcp_dynamic_cast<
const crs_matrix_type>(A);
134 "Ifpack2::Details::InverseDiagonalKernel: operator A must be a Tpetra::CrsMatrix.");
136 const size_t lclNumRows = A_crs_->getRowMap()->getLocalNumElements();
138 if (offsets_.extent(0) < lclNumRows) {
139 using Kokkos::view_alloc;
140 using Kokkos::WithoutInitializing;
141 using offsets_view_type = Kokkos::View<size_t*, device_type>;
143 offsets_ = offsets_view_type();
144 auto howAlloc = view_alloc(
"offsets", WithoutInitializing);
145 offsets_ = offsets_view_type(howAlloc, lclNumRows);
148 A_crs_->getCrsGraph()->getLocalDiagOffsets(offsets_);
152 template <
class TpetraOperatorType>
153 void InverseDiagonalKernel<TpetraOperatorType>::
154 compute(vector_type& D_inv,
155 bool do_l1, magnitude_type L1Eta,
156 bool fixTinyDiagEntries, magnitude_type MinDiagonalValue) {
158 using d_type =
typename vector_type::dual_view_type::t_dev;
160 using d_matrix_type =
typename crs_matrix_type::local_matrix_device_type;
162 const char kernel_label[] =
"inverse_diagonal_kernel";
163 using execution_space =
typename NT::execution_space;
164 using range_type = Kokkos::RangePolicy<execution_space, LO>;
165 const size_t lclNumRows = A_crs_->getRowMap()->getLocalNumElements();
166 auto policy = range_type(0, lclNumRows);
168 d_type d = D_inv.getLocalViewDevice(Tpetra::Access::OverwriteAll);
169 d_matrix_type a = A_crs_->getLocalMatrixDevice();
172 constexpr
bool do_l1_template =
true;
173 if (fixTinyDiagEntries) {
174 constexpr
bool fix_tiny_template =
true;
176 Impl::InverseDiagonalWithExtraction<d_type,
181 functor_type func(d, a, offsets_, L1Eta, MinDiagonalValue);
182 Kokkos::parallel_for(kernel_label, policy, func);
184 constexpr
bool fix_tiny_template =
false;
186 Impl::InverseDiagonalWithExtraction<d_type,
191 functor_type func(d, a, offsets_, L1Eta, MinDiagonalValue);
192 Kokkos::parallel_for(kernel_label, policy, func);
195 constexpr
bool do_l1_template =
false;
196 if (fixTinyDiagEntries) {
197 constexpr
bool fix_tiny_template =
true;
199 Impl::InverseDiagonalWithExtraction<d_type,
204 functor_type func(d, a, offsets_, L1Eta, MinDiagonalValue);
205 Kokkos::parallel_for(kernel_label, policy, func);
207 constexpr
bool fix_tiny_template =
false;
209 Impl::InverseDiagonalWithExtraction<d_type,
214 functor_type func(d, a, offsets_, L1Eta, MinDiagonalValue);
215 Kokkos::parallel_for(kernel_label, policy, func);
223 #define IFPACK2_DETAILS_INVERSEDIAGONALKERNEL_INSTANT(SC, LO, GO, NT) \
224 template class Ifpack2::Details::InverseDiagonalKernel<Tpetra::Operator<SC, LO, GO, NT> >;
226 #endif // IFPACK2_DETAILS_INVERSEDIAGONALKERNEL_DEF_HPP
#define TEUCHOS_TEST_FOR_EXCEPTION(throw_exception_test, Exception, msg)
Compute scaled damped residual for Chebyshev.
Definition: Ifpack2_Details_InverseDiagonalKernel_decl.hpp:45
#define TEUCHOS_ASSERT(assertion_test)