Ifpack2_LocalSparseTriangularSolver_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_LOCALSPARSETRIANGULARSOLVER_DECL_HPP
#define IFPACK2_LOCALSPARSETRIANGULARSOLVER_DECL_HPP
 
#include "Ifpack2_Preconditioner.hpp"
#include "Ifpack2_Details_CanChangeMatrix.hpp"
#include "Tpetra_CrsMatrix_fwd.hpp"
#include "Teuchos_FancyOStream.hpp"
#include <type_traits>
 
#include "KokkosSparse_sptrsv.hpp"
 
namespace Ifpack2 {
 
template <class MatrixType>
class LocalSparseTriangularSolver : virtual public Ifpack2::Preconditioner<typename MatrixType::scalar_type,
                                                                           typename MatrixType::local_ordinal_type,
                                                                           typename MatrixType::global_ordinal_type,
                                                                           typename MatrixType::node_type>,
                                    virtual public Ifpack2::Details::CanChangeMatrix<Tpetra::RowMatrix<typename MatrixType::scalar_type,
                                                                                                       typename MatrixType::local_ordinal_type,
                                                                                                       typename MatrixType::global_ordinal_type,
                                                                                                       typename MatrixType::node_type> > {
 public:
  typedef typename MatrixType::scalar_type scalar_type;
  typedef typename MatrixType::impl_scalar_type impl_scalar_type;
  typedef typename MatrixType::local_ordinal_type local_ordinal_type;
  typedef typename MatrixType::global_ordinal_type global_ordinal_type;
  typedef typename MatrixType::node_type node_type;
 
  typedef typename MatrixType::mag_type magnitude_type;
  typedef Tpetra::Map<local_ordinal_type, global_ordinal_type, node_type> map_type;
  typedef Tpetra::RowMatrix<scalar_type, local_ordinal_type,
                            global_ordinal_type, node_type>
      row_matrix_type;
  typedef Tpetra::CrsMatrix<scalar_type, local_ordinal_type,
                            global_ordinal_type, node_type>
      crs_matrix_type;
 
  static_assert(std::is_same<MatrixType, row_matrix_type>::value,
                "Ifpack2::LocalSparseTriangularSolver: The template parameter "
                "MatrixType must be a Tpetra::RowMatrix specialization.  "
                "Please don't use Tpetra::CrsMatrix (a subclass of "
                "Tpetra::RowMatrix) here anymore.  The constructor can take "
                "either a RowMatrix or a CrsMatrix just fine.");
 
  // Use the local matrix types
  using local_matrix_device_type       = typename crs_matrix_type::local_matrix_device_type;
  using local_matrix_graph_device_type = typename local_matrix_device_type::StaticCrsGraphType;
  using lno_row_view_t                 = typename local_matrix_graph_device_type::row_map_type;
  using lno_nonzero_view_t             = typename local_matrix_graph_device_type::entries_type;
  using scalar_nonzero_view_t          = typename local_matrix_device_type::values_type;
  using TemporaryMemorySpace           = typename local_matrix_graph_device_type::device_type::memory_space;
  using PersistentMemorySpace          = typename local_matrix_graph_device_type::device_type::memory_space;
  using HandleExecSpace                = typename local_matrix_graph_device_type::device_type::execution_space;
  using k_handle                       = typename KokkosKernels::Experimental::KokkosKernelsHandle<typename lno_row_view_t::const_value_type, typename lno_nonzero_view_t::const_value_type, typename scalar_nonzero_view_t::value_type, HandleExecSpace, TemporaryMemorySpace, PersistentMemorySpace>;
 
  LocalSparseTriangularSolver(const Teuchos::RCP<const row_matrix_type>& A);
 
  LocalSparseTriangularSolver(const Teuchos::RCP<const row_matrix_type>& A,
                              const Teuchos::RCP<Teuchos::FancyOStream>& out);
 
  LocalSparseTriangularSolver();
 
  LocalSparseTriangularSolver(const bool /* unused */, const Teuchos::RCP<Teuchos::FancyOStream>& out);
 
  virtual ~LocalSparseTriangularSolver();
 
  void setParameters(const Teuchos::ParameterList& params);
 
  void initialize();
 
  inline bool isInitialized() const {
    return isInitialized_;
  }
 
  void compute();
 
  inline bool isComputed() const {
    return isComputed_;
  }
 
 
 
  void
  apply(const Tpetra::MultiVector<scalar_type, local_ordinal_type, global_ordinal_type, node_type>& X,
        Tpetra::MultiVector<scalar_type, local_ordinal_type, global_ordinal_type, node_type>& Y,
        Teuchos::ETransp mode = Teuchos::NO_TRANS,
        scalar_type alpha     = Teuchos::ScalarTraits<scalar_type>::one(),
        scalar_type beta      = Teuchos::ScalarTraits<scalar_type>::zero()) const;
 
  Teuchos::RCP<const map_type> getDomainMap() const;
 
  Teuchos::RCP<const map_type> getRangeMap() const;
 
  void
  applyMat(const Tpetra::MultiVector<scalar_type, local_ordinal_type,
                                     global_ordinal_type, node_type>& X,
           Tpetra::MultiVector<scalar_type, local_ordinal_type,
                               global_ordinal_type, node_type>& Y,
           Teuchos::ETransp mode = Teuchos::NO_TRANS) const;
 
  Teuchos::RCP<const Teuchos::Comm<int> > getComm() const;
 
  Teuchos::RCP<const row_matrix_type> getMatrix() const {
    return A_;
  }
 
  double getComputeFlops() const;
 
  double getApplyFlops() const;
 
  int getNumInitialize() const;
 
  int getNumCompute() const;
 
  int getNumApply() const;
 
  double getInitializeTime() const;
 
  double getComputeTime() const;
 
  double getApplyTime() const;
 
 
 
  std::string description() const;
 
  void
  describe(Teuchos::FancyOStream& out,
           const Teuchos::EVerbosityLevel verbLevel =
               Teuchos::Describable::verbLevel_default) const;
 
  virtual void setMatrix(const Teuchos::RCP<const row_matrix_type>& A);
 
  void setStreamInfo(const bool& isKokkosKernelsStream, const int& num_streams, const std::vector<HandleExecSpace>& exec_space_instances);
 
  void setMatrices(const std::vector<Teuchos::RCP<crs_matrix_type> >& A_crs_v);
 
 
 private:
  Teuchos::RCP<const row_matrix_type> A_;
  Teuchos::RCP<Teuchos::FancyOStream> out_;
  Teuchos::RCP<const crs_matrix_type> A_crs_;
  std::vector<Teuchos::RCP<crs_matrix_type> > A_crs_v_;
 
  typedef Tpetra::MultiVector<scalar_type, local_ordinal_type, global_ordinal_type, node_type> MV;
  mutable Teuchos::RCP<MV> X_colMap_;
  mutable Teuchos::RCP<MV> Y_rowMap_;
 
  bool isInitialized_;
  bool isComputed_;
  bool isInternallyChanged_;
  bool reverseStorage_;
 
  mutable int numInitialize_;
  mutable int numCompute_;
  mutable int numApply_;
 
  double initializeTime_;
  double computeTime_;
  double applyTime_;
 
  class HtsImpl;
  Teuchos::RCP<HtsImpl> htsImpl_;
 
  bool isKokkosKernelsSptrsv_;
  Teuchos::RCP<k_handle> kh_;
  std::vector<Teuchos::RCP<k_handle> > kh_v_;
  int num_streams_;
  bool isKokkosKernelsStream_;
  bool kh_v_nonnull_;
  std::vector<HandleExecSpace> exec_space_instances_;
 
  std::string uplo_;
  std::string diag_;
 
  void
  localApply(const MV& X,
             MV& Y,
             const Teuchos::ETransp mode,
             const scalar_type& alpha,
             const scalar_type& beta) const;
 
  void
  localTriangularSolve(const MV& Y,
                       MV& X,
                       const Teuchos::ETransp mode) const;
 
  void initializeState();
 
  KokkosSparse::Experimental::SPTRSVAlgorithm kokkosKernelsAlgorithm() const;
};
 
}  // namespace Ifpack2
 
#endif  // IFPACK2_LOCALSPARSETRIANGULARSOLVER_DECL_HPP