Ifpack2_Details_InverseDiagonalKernel_def.hpp

// @HEADER
// *****************************************************************************
//       Ifpack2: Templated Object-Oriented Algebraic Preconditioner Package
//
// Copyright 2009 NTESS and the Ifpack2 contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER
 
#ifndef IFPACK2_DETAILS_INVERSEDIAGONALKERNEL_DEF_HPP
#define IFPACK2_DETAILS_INVERSEDIAGONALKERNEL_DEF_HPP
 
#include "Tpetra_CrsMatrix.hpp"
#include "Tpetra_MultiVector.hpp"
#include "Tpetra_Operator.hpp"
#include "Tpetra_Vector.hpp"
#include "Tpetra_Export_decl.hpp"
#include "Tpetra_Import_decl.hpp"
#include "Kokkos_ArithTraits.hpp"
#include "Teuchos_Assert.hpp"
#include <type_traits>
#include "KokkosSparse_spmv_impl.hpp"
 
namespace Ifpack2 {
namespace Details {
namespace Impl {

template<class DVector,
         class AMatrix,
         class DiagOffsetType,
         bool do_L1,
         bool fix_tiny>
struct InverseDiagonalWithExtraction {
 
  using execution_space = typename AMatrix::execution_space;
  using LO = typename AMatrix::non_const_ordinal_type;
  using value_type = typename AMatrix::non_const_value_type;
  using team_policy = typename Kokkos::TeamPolicy<execution_space>;
  using team_member = typename team_policy::member_type;
  using ATV = Kokkos::ArithTraits<value_type>;
  // using IST = typename vector_type::impl_scalar_type;
  using magnitude_type = typename ATV::mag_type;
  using MATV = Kokkos::ArithTraits<magnitude_type>;
 
  DVector m_d;
  AMatrix m_A;
  DiagOffsetType m_offsets;
  magnitude_type m_L1Eta;
  magnitude_type m_MinDiagonalValue;
 
  InverseDiagonalWithExtraction (const DVector& m_d_,
                                 const AMatrix& m_A_,
                                 const DiagOffsetType& m_offsets_,
                                 const magnitude_type m_L1Eta_,
                                 const magnitude_type m_MinDiagonalValue_) :
    m_d (m_d_),
    m_A (m_A_),
    m_offsets (m_offsets_),
    m_L1Eta (m_L1Eta_),
    m_MinDiagonalValue (m_MinDiagonalValue_)
  {
    const size_t numRows = m_A.numRows ();
 
    TEUCHOS_ASSERT( numRows == size_t (m_d.extent (0)) );
    TEUCHOS_ASSERT( numRows == size_t (m_offsets.extent (0)) );
  }
 
  KOKKOS_INLINE_FUNCTION
  void operator() (const LO lclRow) const
  {
    const size_t INV = Tpetra::Details::OrdinalTraits<size_t>::invalid ();
    const value_type one = ATV::one();
 
    // In case the row has no entries
    m_d(lclRow,0) = ATV::zero();
 
    if (m_offsets(lclRow) != INV) {
      auto curRow = m_A.rowConst (lclRow);
      value_type d = curRow.value(m_offsets(lclRow));
 
      if (do_L1) {
        // Setup for L1 Methods.
        // Here we add half the value of the off-processor entries in the row,
        // but only if diagonal isn't sufficiently large.
        //
        // This follows from Equation (6.5) in: Baker, Falgout, Kolev and
        // Yang.  "Multigrid Smoothers for Ultraparallel Computing."  SIAM
        // J. Sci. Comput., Vol. 33, No. 5. (2011), pp. 2864-2887.
 
        const magnitude_type half = MATV::one () / (MATV::one () + MATV::one ());
        const LO numRows = static_cast<LO> (m_A.numRows ());
        const LO row_length = static_cast<LO> (curRow.length);
        magnitude_type diagonal_boost = MATV::zero();
        for (LO iEntry = 0; iEntry < row_length; iEntry++) {
          if (curRow.colidx(iEntry) >= numRows)
            diagonal_boost += ATV::magnitude(curRow.value(iEntry));
        }
        diagonal_boost *= half;
        if (ATV::magnitude(d) < m_L1Eta * diagonal_boost)
          d += diagonal_boost;
      }
 
      if (fix_tiny) {
        // Replace diagonal entries that are too small.
 
        if (ATV::magnitude(d) <= m_MinDiagonalValue)
          d = m_MinDiagonalValue;
      }
 
      // invert diagonal entries
      m_d(lclRow,0) = one / d;
    }
  }
 
};
 
} // namespace Impl
 
 
template<class TpetraOperatorType>
InverseDiagonalKernel<TpetraOperatorType>::
InverseDiagonalKernel (const Teuchos::RCP<const operator_type>& A)
{
  setMatrix (A);
}
 
template<class TpetraOperatorType>
void
InverseDiagonalKernel<TpetraOperatorType>::
setMatrix (const Teuchos::RCP<const operator_type>& A)
{
  if (A_op_.get () != A.get ()) {
    A_op_ = A;
 
    using Teuchos::rcp_dynamic_cast;
    A_crs_ = rcp_dynamic_cast<const crs_matrix_type> (A);
 
    TEUCHOS_TEST_FOR_EXCEPTION
      (A_crs_.is_null(), std::logic_error,
       "Ifpack2::Details::InverseDiagonalKernel: operator A must be a Tpetra::CrsMatrix.");
 
    const size_t lclNumRows = A_crs_->getRowMap ()->getLocalNumElements ();
 
    if (offsets_.extent (0) < lclNumRows) {
      using Kokkos::view_alloc;
      using Kokkos::WithoutInitializing;
      using offsets_view_type = Kokkos::View<size_t*, device_type>;
 
      offsets_ = offsets_view_type (); // clear 1st to save mem
      auto howAlloc = view_alloc ("offsets", WithoutInitializing);
      offsets_ = offsets_view_type (howAlloc, lclNumRows);
    }
 
    A_crs_->getCrsGraph ()->getLocalDiagOffsets (offsets_);
  }
}
 
template<class TpetraOperatorType>
void
InverseDiagonalKernel<TpetraOperatorType>::
compute (vector_type& D_inv,
         bool do_l1, magnitude_type L1Eta,
         bool fixTinyDiagEntries, magnitude_type MinDiagonalValue)
{
 
 
    // Canonicalize template arguments to avoid redundant instantiations.
  using d_type = typename vector_type::dual_view_type::t_dev;
  //  using h_matrix_type = typename crs_matrix_type::local_matrix_host_type;
  using d_matrix_type = typename crs_matrix_type::local_matrix_device_type;
 
  const char kernel_label[] = "inverse_diagonal_kernel";
  using execution_space = typename NT::execution_space;
  using range_type = Kokkos::RangePolicy<execution_space, LO>;
  const size_t lclNumRows = A_crs_->getRowMap ()->getLocalNumElements ();
  auto policy = range_type(0, lclNumRows);
 
  d_type d = D_inv.getLocalViewDevice(Tpetra::Access::OverwriteAll);
  d_matrix_type a = A_crs_->getLocalMatrixDevice();
 
  if (do_l1) {
    constexpr bool do_l1_template = true;
    if (fixTinyDiagEntries) {
      constexpr bool fix_tiny_template = true;
      using functor_type =
        Impl::InverseDiagonalWithExtraction<d_type,
                               d_matrix_type,
                               offset_type,
                               do_l1_template,
                               fix_tiny_template>;
      functor_type func (d, a, offsets_, L1Eta, MinDiagonalValue);
      Kokkos::parallel_for (kernel_label, policy, func);
    } else {
      constexpr bool fix_tiny_template = false;
      using functor_type =
        Impl::InverseDiagonalWithExtraction<d_type,
                               d_matrix_type,
                               offset_type,
                               do_l1_template,
                               fix_tiny_template>;
      functor_type func (d, a, offsets_, L1Eta, MinDiagonalValue);
      Kokkos::parallel_for (kernel_label, policy, func);
    }
  } else {
    constexpr bool do_l1_template = false;
    if (fixTinyDiagEntries) {
      constexpr bool fix_tiny_template = true;
      using functor_type =
        Impl::InverseDiagonalWithExtraction<d_type,
                               d_matrix_type,
                               offset_type,
                               do_l1_template,
                               fix_tiny_template>;
      functor_type func (d, a, offsets_, L1Eta, MinDiagonalValue);
      Kokkos::parallel_for (kernel_label, policy, func);
    } else {
      constexpr bool fix_tiny_template = false;
      using functor_type =
        Impl::InverseDiagonalWithExtraction<d_type,
                               d_matrix_type,
                               offset_type,
                               do_l1_template,
                               fix_tiny_template>;
      functor_type func (d, a, offsets_, L1Eta, MinDiagonalValue);
      Kokkos::parallel_for (kernel_label, policy, func);
    }
  }
}
 
} // namespace Details
} // namespace Ifpack2
 
#define IFPACK2_DETAILS_INVERSEDIAGONALKERNEL_INSTANT(SC,LO,GO,NT) \
  template class Ifpack2::Details::InverseDiagonalKernel<Tpetra::Operator<SC, LO, GO, NT> >;
 
#endif // IFPACK2_DETAILS_INVERSEDIAGONALKERNEL_DEF_HPP