Ifpack2_Details_GaussSeidel.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_GAUSS_SEIDEL_HPP
#define IFPACK2_GAUSS_SEIDEL_HPP

namespace Ifpack2 {
namespace Details {
template <typename Scalar, typename LO, typename GO, typename NT>
struct GaussSeidel {
  using crs_matrix_type          = Tpetra::CrsMatrix<Scalar, LO, GO, NT>;
  using bcrs_matrix_type         = Tpetra::BlockCrsMatrix<Scalar, LO, GO, NT>;
  using row_matrix_type          = Tpetra::RowMatrix<Scalar, LO, GO, NT>;
  using local_matrix_device_type = typename crs_matrix_type::local_matrix_device_type;
  using vector_type              = Tpetra::Vector<Scalar, LO, GO, NT>;
  using multivector_type         = Tpetra::MultiVector<Scalar, LO, GO, NT>;
  using block_multivector_type   = Tpetra::BlockMultiVector<Scalar, LO, GO, NT>;
  using mem_space_t              = typename local_matrix_device_type::memory_space;
  using rowmap_t                 = typename local_matrix_device_type::row_map_type::host_mirror_type;
  using entries_t                = typename local_matrix_device_type::index_type::host_mirror_type;
  using values_t                 = typename local_matrix_device_type::values_type::host_mirror_type;
  using const_rowmap_t           = typename rowmap_t::const_type;
  using const_entries_t          = typename entries_t::const_type;
  using const_values_t           = typename values_t::const_type;
  using Offset                   = typename rowmap_t::non_const_value_type;
  using IST                      = typename crs_matrix_type::impl_scalar_type;
#if KOKKOS_VERSION >= 40799
  using KAT = KokkosKernels::ArithTraits<IST>;
#else
  using KAT = Kokkos::ArithTraits<IST>;
#endif
  // Type of view representing inverse diagonal blocks, and its host_mirror_type.
  using InverseBlocks     = Kokkos::View<IST***, typename bcrs_matrix_type::device_type>;
  using InverseBlocksHost = typename InverseBlocks::host_mirror_type;

  typedef typename crs_matrix_type::nonconst_global_inds_host_view_type nonconst_global_inds_host_view_type;
  typedef typename crs_matrix_type::nonconst_local_inds_host_view_type nonconst_local_inds_host_view_type;
  typedef typename crs_matrix_type::nonconst_values_host_view_type nonconst_values_host_view_type;

  // Setup for CrsMatrix
  GaussSeidel(Teuchos::RCP<const crs_matrix_type> A_, Teuchos::RCP<vector_type>& inverseDiagVec_, Teuchos::ArrayRCP<LO>& applyRows_, Scalar omega_) {
    A              = A_;
    numRows        = A_->getLocalNumRows();
    haveRowMatrix  = false;
    inverseDiagVec = inverseDiagVec_;
    applyRows      = applyRows_;
    blockSize      = 1;
    omega          = omega_;
  }

  GaussSeidel(Teuchos::RCP<const row_matrix_type> A_, Teuchos::RCP<vector_type>& inverseDiagVec_, Teuchos::ArrayRCP<LO>& applyRows_, Scalar omega_) {
    A              = A_;
    numRows        = A_->getLocalNumRows();
    haveRowMatrix  = true;
    inverseDiagVec = inverseDiagVec_;
    applyRows      = applyRows_;
    blockSize      = 1;
    omega          = omega_;
    // Here, need to make a deep CRS copy to avoid slow access via getLocalRowCopy
    rowMatrixRowmap  = rowmap_t(Kokkos::ViewAllocateWithoutInitializing("Arowmap"), numRows + 1);
    rowMatrixEntries = entries_t(Kokkos::ViewAllocateWithoutInitializing("Aentries"), A_->getLocalNumEntries());
    rowMatrixValues  = values_t(Kokkos::ViewAllocateWithoutInitializing("Avalues"), A_->getLocalNumEntries());
    size_t maxDegree = A_->getLocalMaxNumRowEntries();
    nonconst_values_host_view_type rowValues("rowValues", maxDegree);
    nonconst_local_inds_host_view_type rowEntries("rowEntries", maxDegree);
    size_t accum = 0;
    for (LO i = 0; i <= numRows; i++) {
      rowMatrixRowmap(i) = accum;
      if (i == numRows)
        break;
      size_t degree;
      A_->getLocalRowCopy(i, rowEntries, rowValues, degree);
      accum += degree;
      size_t rowBegin = rowMatrixRowmap(i);
      for (size_t j = 0; j < degree; j++) {
        rowMatrixEntries(rowBegin + j) = rowEntries[j];
        rowMatrixValues(rowBegin + j)  = rowValues[j];
      }
    }
  }

  GaussSeidel(Teuchos::RCP<const bcrs_matrix_type> A_, const InverseBlocks& inverseBlockDiag_, Teuchos::ArrayRCP<LO>& applyRows_, Scalar omega_) {
    A             = A_;
    numRows       = A_->getLocalNumRows();
    haveRowMatrix = false;
    // note: next 2 lines are no-ops if inverseBlockDiag_ is already host-accessible
    inverseBlockDiag = Kokkos::create_mirror_view(inverseBlockDiag_);
    Kokkos::deep_copy(inverseBlockDiag, inverseBlockDiag_);
    applyRows = applyRows_;
    omega     = omega_;
    blockSize = A_->getBlockSize();
  }

  template <bool useApplyRows, bool multipleRHS, bool omegaNotOne>
  void applyImpl(const const_values_t& Avalues, const const_rowmap_t& Arowmap, const const_entries_t& Aentries, multivector_type& x, const multivector_type& b, Tpetra::ESweepDirection direction) {
    // note: direction is either Forward or Backward (Symmetric is handled in apply())
    LO numApplyRows = useApplyRows ? (LO)applyRows.size() : numRows;
    // note: inverseDiagMV always has only one column
    auto inverseDiag = Kokkos::subview(inverseDiagVec->getLocalViewHost(Tpetra::Access::ReadOnly), Kokkos::ALL(), 0);
    bool forward     = direction == Tpetra::Forward;
    if (multipleRHS) {
      LO numVecs = x.getNumVectors();
      Teuchos::Array<IST> accum(numVecs);
      auto xlcl = x.get2dViewNonConst();
      auto blcl = b.get2dView();
      for (LO i = 0; i < numApplyRows; i++) {
        LO row;
        if (forward)
          row = useApplyRows ? applyRows[i] : i;
        else
          row = useApplyRows ? applyRows[numApplyRows - 1 - i] : numApplyRows - 1 - i;
        for (LO k = 0; k < numVecs; k++) {
          accum[k] = KAT::zero();
        }
        Offset rowBegin = Arowmap(row);
        Offset rowEnd   = Arowmap(row + 1);
        for (Offset j = rowBegin; j < rowEnd; j++) {
          LO col  = Aentries(j);
          IST val = Avalues(j);
          for (LO k = 0; k < numVecs; k++) {
            accum[k] += val * IST(xlcl[k][col]);
          }
        }
        // Update x
        IST dinv = inverseDiag(row);
        for (LO k = 0; k < numVecs; k++) {
          if (omegaNotOne)
            xlcl[k][row] += Scalar(omega * dinv * (IST(blcl[k][row]) - accum[k]));
          else
            xlcl[k][row] += Scalar(dinv * (IST(blcl[k][row]) - accum[k]));
        }
      }
    } else {
      auto xlcl = Kokkos::subview(x.getLocalViewHost(Tpetra::Access::ReadWrite), Kokkos::ALL(), 0);
      auto blcl = Kokkos::subview(b.getLocalViewHost(Tpetra::Access::ReadOnly), Kokkos::ALL(), 0);
      auto dlcl = Kokkos::subview(inverseDiagVec->getLocalViewHost(Tpetra::Access::ReadOnly), Kokkos::ALL(), 0);
      for (LO i = 0; i < numApplyRows; i++) {
        LO row;
        if (forward)
          row = useApplyRows ? applyRows[i] : i;
        else
          row = useApplyRows ? applyRows[numApplyRows - 1 - i] : numApplyRows - 1 - i;
        IST accum       = KAT::zero();
        Offset rowBegin = Arowmap(row);
        Offset rowEnd   = Arowmap(row + 1);
        for (Offset j = rowBegin; j < rowEnd; j++) {
          accum += Avalues(j) * xlcl(Aentries(j));
        }
        // Update x
        IST dinv = dlcl(row);
        if (omegaNotOne)
          xlcl(row) += omega * dinv * (blcl(row) - accum);
        else
          xlcl(row) += dinv * (blcl(row) - accum);
      }
    }
  }

  void applyBlock(block_multivector_type& x, const block_multivector_type& b, Tpetra::ESweepDirection direction) {
    if (direction == Tpetra::Symmetric) {
      applyBlock(x, b, Tpetra::Forward);
      applyBlock(x, b, Tpetra::Backward);
      return;
    }
    auto Abcrs = Teuchos::rcp_dynamic_cast<const bcrs_matrix_type>(A);
    if (Abcrs.is_null())
      throw std::runtime_error("Ifpack2::Details::GaussSeidel::applyBlock: A must be a BlockCrsMatrix");
    auto Avalues   = Abcrs->getValuesHost();
    auto AlclGraph = Abcrs->getCrsGraph().getLocalGraphHost();
    auto Arowmap   = AlclGraph.row_map;
    auto Aentries  = AlclGraph.entries;
    // Number of scalars in Avalues per block entry.
    Offset bs2 = blockSize * blockSize;
    LO numVecs = x.getNumVectors();
    Kokkos::View<IST**, Kokkos::LayoutLeft, Kokkos::HostSpace> accum(
        Kokkos::ViewAllocateWithoutInitializing("BlockGaussSeidel Accumulator"), blockSize, numVecs);
    Kokkos::View<IST**, Kokkos::LayoutLeft, Kokkos::HostSpace> dinv_accum(
        Kokkos::ViewAllocateWithoutInitializing("BlockGaussSeidel A_ii^-1*Accumulator"), blockSize, numVecs);
    bool useApplyRows = !applyRows.is_null();
    bool forward      = direction == Tpetra::Forward;
    LO numApplyRows   = useApplyRows ? applyRows.size() : numRows;
    for (LO i = 0; i < numApplyRows; i++) {
      LO row;
      if (forward)
        row = useApplyRows ? applyRows[i] : i;
      else
        row = useApplyRows ? applyRows[numApplyRows - 1 - i] : numApplyRows - 1 - i;
      for (LO v = 0; v < numVecs; v++) {
        auto bRow = b.getLocalBlockHost(row, v, Tpetra::Access::ReadOnly);
        for (LO k = 0; k < blockSize; k++) {
          accum(k, v) = KAT::zero();
        }
      }
      Offset rowBegin = Arowmap(row);
      Offset rowEnd   = Arowmap(row + 1);
      for (Offset j = rowBegin; j < rowEnd; j++) {
        LO col         = Aentries(j);
        const IST* blk = &Avalues(j * bs2);
        for (LO v = 0; v < numVecs; v++) {
          auto xCol = x.getLocalBlockHost(col, v, Tpetra::Access::ReadOnly);
          for (LO br = 0; br < blockSize; br++) {
            for (LO bc = 0; bc < blockSize; bc++) {
              IST Aval = blk[br * blockSize + bc];
              accum(br, v) += Aval * xCol(bc);
            }
          }
        }
      }
      // Update x: term is omega * Aii^-1 * accum, where omega is scalar, Aii^-1 is bs*bs, and accum is bs*nv
      auto invBlock = Kokkos::subview(inverseBlockDiag, row, Kokkos::ALL(), Kokkos::ALL());
      Kokkos::deep_copy(dinv_accum, KAT::zero());
      for (LO v = 0; v < numVecs; v++) {
        auto bRow = b.getLocalBlockHost(row, v, Tpetra::Access::ReadOnly);
        for (LO br = 0; br < blockSize; br++) {
          accum(br, v) = bRow(br) - accum(br, v);
        }
      }
      for (LO v = 0; v < numVecs; v++) {
        for (LO br = 0; br < blockSize; br++) {
          for (LO bc = 0; bc < blockSize; bc++) {
            dinv_accum(br, v) += invBlock(br, bc) * accum(bc, v);
          }
        }
      }
      // Update x
      for (LO v = 0; v < numVecs; v++) {
        auto xRow = x.getLocalBlockHost(row, v, Tpetra::Access::ReadWrite);
        for (LO k = 0; k < blockSize; k++) {
          xRow(k) += omega * dinv_accum(k, v);
        }
      }
    }
  }

  // Version of apply for CrsMatrix/RowMatrix (for BlockCrsMatrix, call applyBlock)
  void apply(multivector_type& x, const multivector_type& b, Tpetra::ESweepDirection direction) {
    if (direction == Tpetra::Symmetric) {
      apply(x, b, Tpetra::Forward);
      apply(x, b, Tpetra::Backward);
      return;
    }
    const_values_t Avalues;
    const_rowmap_t Arowmap;
    const_entries_t Aentries;
    if (haveRowMatrix) {
      Avalues  = rowMatrixValues;
      Arowmap  = rowMatrixRowmap;
      Aentries = rowMatrixEntries;
    } else {
      auto Acrs = Teuchos::rcp_dynamic_cast<const crs_matrix_type>(A);
      if (Acrs.is_null())
        throw std::runtime_error("Ifpack2::Details::GaussSeidel::apply: either haveRowMatrix, or A is CrsMatrix");
      auto Alcl = Acrs->getLocalMatrixHost();
      Avalues   = Alcl.values;
      Arowmap   = Alcl.graph.row_map;
      Aentries  = Alcl.graph.entries;
    }
    if (applyRows.is_null()) {
      if (x.getNumVectors() > 1)
        this->template applyImpl<false, true, true>(Avalues, Arowmap, Aentries, x, b, direction);
      else {
        // Optimize for the all-rows, single vector, omega = 1 case
        if (omega == KAT::one())
          this->template applyImpl<false, false, false>(Avalues, Arowmap, Aentries, x, b, direction);
        else
          this->template applyImpl<false, false, true>(Avalues, Arowmap, Aentries, x, b, direction);
      }
    } else {
      this->template applyImpl<true, true, true>(Avalues, Arowmap, Aentries, x, b, direction);
    }
  }

  Teuchos::RCP<const row_matrix_type> A;
  // For efficiency, if input is a RowMatrix, keep a persistent copy of the CRS formatted local matrix.
  bool haveRowMatrix;
  values_t rowMatrixValues;
  rowmap_t rowMatrixRowmap;
  entries_t rowMatrixEntries;
  LO numRows;
  IST omega;
  // If set up with BlockCrsMatrix, the block size. Otherwise 1.
  LO blockSize;
  // If using a non-block matrix, the inverse diagonal.
  Teuchos::RCP<vector_type> inverseDiagVec;
  // If using a block matrix, the inverses of all diagonal blocks.
  InverseBlocksHost inverseBlockDiag;
  // If null, apply over all rows in natural order. Otherwise, apply for each row listed in order.
  Teuchos::ArrayRCP<LO> applyRows;
};
}  // namespace Details
}  // namespace Ifpack2

#endif