Ifpack2_Container_decl.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_CONTAINER_DECL_HPP
#define IFPACK2_CONTAINER_DECL_HPP


#include "Ifpack2_ConfigDefs.hpp"
#include "Tpetra_RowMatrix.hpp"
#include "Teuchos_Describable.hpp"
#include <Tpetra_Map.hpp>
#include <Tpetra_BlockCrsMatrix.hpp>
#include <Teuchos_ParameterList.hpp>
#include <iostream>

#ifndef DOXYGEN_SHOULD_SKIP_THIS
namespace Teuchos {
// Forward declaration to avoid include.
class ParameterList;
}  // namespace Teuchos
#endif  // DOXYGEN_SHOULD_SKIP_THIS

namespace Ifpack2 {

template <class MatrixType>
class Container : public Teuchos::Describable {
 protected:
  using scalar_type           = typename MatrixType::scalar_type;
  using local_ordinal_type    = typename MatrixType::local_ordinal_type;
  using global_ordinal_type   = typename MatrixType::global_ordinal_type;
  using node_type             = typename MatrixType::node_type;
  using SC                    = scalar_type;
  using LO                    = local_ordinal_type;
  using GO                    = global_ordinal_type;
  using NO                    = node_type;
  using import_type           = Tpetra::Import<LO, GO, NO>;
  using mv_type               = Tpetra::MultiVector<SC, LO, GO, NO>;
  using vector_type           = Tpetra::Vector<SC, LO, GO, NO>;
  using map_type              = Tpetra::Map<LO, GO, NO>;
  using crs_matrix_type       = Tpetra::CrsMatrix<SC, LO, GO, NO>;
  using block_crs_matrix_type = Tpetra::BlockCrsMatrix<SC, LO, GO, NO>;
  using row_matrix_type       = Tpetra::RowMatrix<SC, LO, GO, NO>;

  static_assert(std::is_same<MatrixType, row_matrix_type>::value,
                "Ifpack2::Container: Please use MatrixType = Tpetra::RowMatrix.");

#if KOKKOS_VERSION >= 40799
  using ISC = typename KokkosKernels::ArithTraits<SC>::val_type;
#else
  using ISC  = typename Kokkos::ArithTraits<SC>::val_type;
#endif

  using HostView      = typename mv_type::dual_view_type::t_host;
  using ConstHostView = typename HostView::const_type;

 public:
  Container(const Teuchos::RCP<const row_matrix_type>& matrix,
            const Teuchos::Array<Teuchos::Array<LO> >& partitions,
            bool pointIndexed);

  virtual ~Container();

  Teuchos::ArrayView<const LO> getBlockRows(int blockIndex) const;

  virtual void initialize() = 0;

  void setBlockSizes(const Teuchos::Array<Teuchos::Array<LO> >& partitions);

  void getMatDiag() const;

  bool isInitialized() const;

  bool isComputed() const;

  virtual void compute() = 0;

  void DoJacobi(ConstHostView X, HostView Y, SC dampingFactor) const;
  void DoOverlappingJacobi(ConstHostView X, HostView Y, ConstHostView W, SC dampingFactor, bool nonsymScaling) const;
  void DoGaussSeidel(ConstHostView X, HostView Y, HostView Y2, SC dampingFactor) const;
  void DoSGS(ConstHostView X, HostView Y, HostView Y2, SC dampingFactor) const;

  virtual void setParameters(const Teuchos::ParameterList& List) = 0;

  virtual void
  apply(ConstHostView X,
        HostView Y,
        int blockIndex,
        Teuchos::ETransp mode = Teuchos::NO_TRANS,
        SC alpha              = Teuchos::ScalarTraits<SC>::one(),
        SC beta               = Teuchos::ScalarTraits<SC>::zero()) const = 0;

  virtual void
  weightedApply(ConstHostView X,
                HostView Y,
                ConstHostView D,
                int blockIndex,
                Teuchos::ETransp mode = Teuchos::NO_TRANS,
                SC alpha              = Teuchos::ScalarTraits<SC>::one(),
                SC beta               = Teuchos::ScalarTraits<SC>::zero()) const = 0;

  // The underlying container implements the splitting <tt>A = D + R</tt>. Only
  // it can have efficient access to D and R, as these are constructed in the
  // symbolic and numeric phases.
  //
  // This is the first performance-portable implementation of a block
  // relaxation, and it is supported currently only by BlockTriDiContainer.
  virtual void applyInverseJacobi(const mv_type& /* X */, mv_type& /* Y */,
                                  SC dampingFactor,
                                  bool /* zeroStartingSolution = false */,
                                  int /* numSweeps = 1 */) const = 0;

  virtual void applyMV(const mv_type& X, mv_type& Y) const;

  virtual void weightedApplyMV(const mv_type& X,
                               mv_type& Y,
                               vector_type& W) const;

  virtual void clearBlocks();

  virtual std::ostream& print(std::ostream& os) const = 0;

  static std::string getName();

 protected:
  virtual void DoGSBlock(ConstHostView X, HostView Y, HostView Y2, HostView Resid,
                         SC dampingFactor, LO i) const;

  Teuchos::RCP<const row_matrix_type> inputMatrix_;

  Teuchos::RCP<const crs_matrix_type> inputCrsMatrix_;

  Teuchos::RCP<const block_crs_matrix_type> inputBlockMatrix_;

  int numBlocks_;
  Teuchos::Array<LO> blockRows_;  // size: total # of local rows (in all local blocks)
  Teuchos::Array<LO> blockSizes_;  // size: # of blocks
  Teuchos::Array<LO> blockOffsets_;
  mutable Teuchos::RCP<vector_type> Diag_;
  bool IsParallel_;
  LO NumLocalRows_;
  GO NumGlobalRows_;
  GO NumGlobalNonzeros_;
  bool hasBlockCrs_;
  int bcrsBlockSize_;
  bool pointIndexed_;
  LO scalarsPerRow_;

  bool IsInitialized_;
  bool IsComputed_;

  LO maxBlockSize_;
};

namespace Details {
template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
struct StridedRowView;
}


template <class MatrixType, class LocalScalarType>
class ContainerImpl : public Container<MatrixType> {

 protected:
  using local_scalar_type = LocalScalarType;
  using SC                = typename Container<MatrixType>::scalar_type;
  using LO                = typename Container<MatrixType>::local_ordinal_type;
  using GO                = typename Container<MatrixType>::global_ordinal_type;
  using NO                = typename Container<MatrixType>::node_type;
  using StridedRowView    = Details::StridedRowView<SC, LO, GO, NO>;
  using typename Container<MatrixType>::import_type;
  using typename Container<MatrixType>::row_matrix_type;
  using typename Container<MatrixType>::crs_matrix_type;
  using typename Container<MatrixType>::block_crs_matrix_type;
  using typename Container<MatrixType>::mv_type;
  using typename Container<MatrixType>::vector_type;
  using typename Container<MatrixType>::map_type;
  using typename Container<MatrixType>::ISC;
  using LSC = LocalScalarType;
#if KOKKOS_VERSION >= 40799
  using LISC = typename KokkosKernels::ArithTraits<LSC>::val_type;
#else
  using LISC = typename Kokkos::ArithTraits<LSC>::val_type;
#endif

  using local_mv_type = Tpetra::MultiVector<LSC, LO, GO, NO>;

  using typename Container<MatrixType>::HostView;
  using typename Container<MatrixType>::ConstHostView;
  using HostViewLocal         = typename local_mv_type::dual_view_type::t_host;
  using HostSubviewLocal      = Kokkos::View<LISC**, Kokkos::LayoutStride, typename HostViewLocal::memory_space>;
  using ConstHostSubviewLocal = Kokkos::View<const LISC**, Kokkos::LayoutStride, typename HostViewLocal::memory_space>;

  static_assert(std::is_same<MatrixType, row_matrix_type>::value,
                "Ifpack2::Container: Please use MatrixType = Tpetra::RowMatrix.");
  //

  LO translateRowToCol(LO row);

  Details::StridedRowView<SC, LO, GO, NO> getInputRowView(LO row) const;

  virtual void extract() = 0;

 public:
  ContainerImpl(const Teuchos::RCP<const row_matrix_type>& matrix,
                const Teuchos::Array<Teuchos::Array<LO> >& partitions,
                bool pointIndexed);

  virtual ~ContainerImpl();

  virtual void initialize() = 0;

  virtual void compute() = 0;

  virtual void setParameters(const Teuchos::ParameterList& List);

  virtual void
  apply(ConstHostView X,
        HostView Y,
        int blockIndex,
        Teuchos::ETransp mode = Teuchos::NO_TRANS,
        SC alpha              = Teuchos::ScalarTraits<SC>::one(),
        SC beta               = Teuchos::ScalarTraits<SC>::zero()) const;

  virtual void
  weightedApply(ConstHostView X,
                HostView Y,
                ConstHostView D,
                int blockIndex,
                Teuchos::ETransp mode = Teuchos::NO_TRANS,
                SC alpha              = Teuchos::ScalarTraits<SC>::one(),
                SC beta               = Teuchos::ScalarTraits<SC>::zero()) const;

  // The underlying container implements the splitting <tt>A = D + R</tt>. Only
  // it can have efficient access to D and R, as these are constructed in the
  // symbolic and numeric phases.
  //
  // This is the first performance-portable implementation of a block
  // relaxation, and it is supported currently only by BlockTriDiContainer.
  virtual void applyInverseJacobi(const mv_type& /* X */, mv_type& /* Y */,
                                  SC dampingFactor,
                                  bool /* zeroStartingSolution = false */,
                                  int /* numSweeps = 1 */) const;

  void applyMV(const mv_type& X, mv_type& Y) const;

  void weightedApplyMV(const mv_type& X,
                       mv_type& Y,
                       vector_type& W) const;

  virtual void clearBlocks();

  virtual std::ostream& print(std::ostream& os) const = 0;

  static std::string getName();

 protected:
  // Do Gauss-Seidel on only block i (this is used by DoGaussSeidel and DoSGS)
  void DoGSBlock(ConstHostView X, HostView Y, HostView Y2, HostView Resid,
                 SC dampingFactor, LO i) const;

  virtual void
  solveBlock(ConstHostSubviewLocal X,
             HostSubviewLocal Y,
             int blockIndex,
             Teuchos::ETransp mode,
             const LSC alpha,
             const LSC beta) const;

  mutable HostViewLocal X_local_;  // length: blockRows_.size()
  mutable HostViewLocal Y_local_;  // length: blockRows_.size()

  mutable HostViewLocal weightedApplyScratch_;

  mutable std::vector<HostSubviewLocal> X_localBlocks_;

  mutable std::vector<HostSubviewLocal> Y_localBlocks_;
};

namespace Details {
template <typename Scalar, typename LocalOrdinal, typename GlobalOrdinal, typename Node>
struct StridedRowView {
  using SC = Scalar;
  using LO = LocalOrdinal;

  using block_crs_matrix_type = Tpetra::BlockCrsMatrix<SC, LO, GlobalOrdinal, Node>;

  using h_inds_type = typename block_crs_matrix_type::local_inds_host_view_type;
  using h_vals_type = typename block_crs_matrix_type::values_host_view_type;
  StridedRowView(h_vals_type vals_, h_inds_type inds_, int blockSize_, size_t nnz_);

  //    StridedRowView(const SC* vals_, const LO* inds_, int blockSize_, size_t nnz_);

  StridedRowView(Teuchos::Array<SC>& vals_, Teuchos::Array<LO>& inds_);

  SC val(size_t i) const;
  LO ind(size_t i) const;

  size_t size() const;

 private:
  h_vals_type vals;
  h_inds_type inds;
  int blockSize;
  size_t nnz;
  // These arrays are only used if the inputMatrix_ doesn't support row views.
  Teuchos::Array<SC> valsCopy;
  Teuchos::Array<LO> indsCopy;
};
}  // namespace Details

}  // namespace Ifpack2

template <class MatrixType>
std::ostream& operator<<(std::ostream& os, const Ifpack2::Container<MatrixType>& obj);

namespace Teuchos {

template <class MatrixType>
class TEUCHOSCORE_LIB_DLL_EXPORT TypeNameTraits< ::Ifpack2::Container<MatrixType> > {
 public:
  static std::string name() {
    return std::string("Ifpack2::Container<") +
           TypeNameTraits<MatrixType>::name() + ">";
  }

  static std::string concreteName(const ::Ifpack2::Container<MatrixType>&) {
    return name();
  }
};

}  // namespace Teuchos

#endif  // IFPACK2_CONTAINER_HPP