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"
#if KOKKOS_VERSION >= 40799
#include "KokkosKernels_ArithTraits.hpp"
#else
#include "Kokkos_ArithTraits.hpp"
#endif
#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;
#if KOKKOS_VERSION >= 40799
  using ATV = KokkosKernels::ArithTraits<value_type>;
#else
  using ATV  = Kokkos::ArithTraits<value_type>;
#endif
  // using IST = typename vector_type::impl_scalar_type;
  using magnitude_type = typename ATV::mag_type;
#if KOKKOS_VERSION >= 40799
  using MATV = KokkosKernels::ArithTraits<magnitude_type>;
#else
  using MATV = Kokkos::ArithTraits<magnitude_type>;
#endif

  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