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.");

  using ISC = typename Kokkos::ArithTraits<SC>::val_type;

  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;
  using LISC = typename Kokkos::ArithTraits<LSC>::val_type;

  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