Tpetra_BlockCrsMatrix_decl.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//          Tpetra: Templated Linear Algebra Services Package
//
// Copyright 2008 NTESS and the Tpetra contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef TPETRA_BLOCKCRSMATRIX_DECL_HPP
#define TPETRA_BLOCKCRSMATRIX_DECL_HPP


#include "Tpetra_CrsGraph.hpp"
#include "Tpetra_RowMatrix.hpp"
#include "Tpetra_BlockMultiVector_decl.hpp"
#include "Tpetra_CrsMatrix_decl.hpp"

#include "KokkosSparse_BsrMatrix.hpp"

#include "Tpetra_Details_MatrixApplyHelper.hpp"

namespace Tpetra {

template <class BlockCrsMatrixType>
Teuchos::RCP<BlockCrsMatrixType>
importAndFillCompleteBlockCrsMatrix(const Teuchos::RCP<const BlockCrsMatrixType>& sourceMatrix,
                                    const Import<typename BlockCrsMatrixType::local_ordinal_type,
                                                 typename BlockCrsMatrixType::global_ordinal_type,
                                                 typename BlockCrsMatrixType::node_type>& importer);
template <class BlockCrsMatrixType>
Teuchos::RCP<BlockCrsMatrixType>
exportAndFillCompleteBlockCrsMatrix(const Teuchos::RCP<const BlockCrsMatrixType>& sourceMatrix,
                                    const Export<typename BlockCrsMatrixType::local_ordinal_type,
                                                 typename BlockCrsMatrixType::global_ordinal_type,
                                                 typename BlockCrsMatrixType::node_type>& exporter);


namespace Impl {
#if defined(TPETRA_ENABLE_BLOCKCRS_LITTLEBLOCK_LAYOUTLEFT)
using BlockCrsMatrixLittleBlockArrayLayout = Kokkos::LayoutLeft;
#else
using BlockCrsMatrixLittleBlockArrayLayout = Kokkos::LayoutRight;
#endif
}  // namespace Impl

template <class Scalar,
          class LO,
          class GO,
          class Node>
class BlockCrsMatrix : virtual public ::Tpetra::RowMatrix<Scalar, LO, GO, Node>,
                       virtual public ::Tpetra::DistObject<char, LO, GO, Node> {
 private:
  using dist_object_type = ::Tpetra::DistObject<char, LO, GO, Node>;
  using BMV              = BlockMultiVector<Scalar, LO, GO, Node>;
  using STS              = Teuchos::ScalarTraits<Scalar>;

 protected:
  typedef char packet_type;

 public:


  using scalar_type = Scalar;

  using impl_scalar_type = typename BMV::impl_scalar_type;

  typedef LO local_ordinal_type;
  typedef GO global_ordinal_type;
  typedef Node node_type;

  typedef typename Node::device_type device_type;
  typedef typename device_type::execution_space execution_space;
  typedef typename device_type::memory_space memory_space;

  typedef ::Tpetra::Map<LO, GO, node_type> map_type;
  typedef ::Tpetra::MultiVector<Scalar, LO, GO, node_type> mv_type;
  typedef ::Tpetra::CrsGraph<LO, GO, node_type> crs_graph_type;

  typedef Kokkos::View<impl_scalar_type**,
                       Impl::BlockCrsMatrixLittleBlockArrayLayout,
                       device_type,
                       Kokkos::MemoryTraits<Kokkos::Unmanaged> >
      little_block_type;
  typedef typename little_block_type::host_mirror_type little_block_host_type;

  typedef Kokkos::View<const impl_scalar_type**,
                       Impl::BlockCrsMatrixLittleBlockArrayLayout,
                       device_type,
                       Kokkos::MemoryTraits<Kokkos::Unmanaged> >
      const_little_block_type;
  typedef typename BMV::little_vec_type little_vec_type;
  typedef typename BMV::little_host_vec_type little_host_vec_type;

  typedef typename BMV::const_little_vec_type const_little_vec_type;
  typedef typename BMV::const_little_host_vec_type const_host_little_vec_type;

  using row_matrix_type = RowMatrix<Scalar, LO, GO, node_type>;
  using local_inds_device_view_type =
      typename row_matrix_type::local_inds_device_view_type;
  using local_inds_host_view_type =
      typename row_matrix_type::local_inds_host_view_type;
  using nonconst_local_inds_host_view_type =
      typename row_matrix_type::nonconst_local_inds_host_view_type;

  using global_inds_device_view_type =
      typename row_matrix_type::global_inds_device_view_type;
  using global_inds_host_view_type =
      typename row_matrix_type::global_inds_host_view_type;
  using nonconst_global_inds_host_view_type =
      typename row_matrix_type::nonconst_global_inds_host_view_type;

  using values_device_view_type =
      typename row_matrix_type::values_device_view_type;
  using values_host_view_type =
      typename row_matrix_type::values_host_view_type;
  using nonconst_values_host_view_type =
      typename row_matrix_type::nonconst_values_host_view_type;

  using local_graph_device_type = typename crs_graph_type::local_graph_device_type;

  using local_matrix_device_type =
      KokkosSparse::Experimental::BsrMatrix<impl_scalar_type,
                                            local_ordinal_type,
                                            device_type,
                                            void,
                                            typename local_graph_device_type::size_type>;

 private:
  // TODO: When KokkosKernels 4.4 is released, local_matrix_device_type can be permanently modified to use the default_size_type
  // of KK. This is always a type that is enabled by KK's ETI (preferring int if both or neither int and size_t are enabled).
  using local_matrix_int_rowptrs_device_type =
      KokkosSparse::Experimental::BsrMatrix<impl_scalar_type,
                                            local_ordinal_type,
                                            device_type,
                                            void,
                                            int>;

  using ApplyHelper = Details::MatrixApplyHelper<
      local_matrix_device_type,
      local_matrix_int_rowptrs_device_type,
      typename mv_type::device_view_type>;

  mutable std::shared_ptr<ApplyHelper> applyHelper;

 public:
  std::shared_ptr<ApplyHelper> getApplyHelper() const {
    if (!applyHelper) {
      auto A_lcl  = getLocalMatrixDevice();
      applyHelper = std::make_shared<ApplyHelper>(A_lcl.nnz(), A_lcl.graph.row_map);
    }
    return applyHelper;
  }

#if KOKKOS_VERSION >= 40799
  using local_matrix_host_type =
      typename local_matrix_device_type::host_mirror_type;
#else
  using local_matrix_host_type =
      typename local_matrix_device_type::HostMirror;
#endif



  BlockCrsMatrix();

  BlockCrsMatrix(const crs_graph_type& graph, const LO blockSize);

  BlockCrsMatrix(const crs_graph_type& graph,
                 const typename local_matrix_device_type::values_type& values,
                 const LO blockSize);

  BlockCrsMatrix(const crs_graph_type& graph,
                 const map_type& domainPointMap,
                 const map_type& rangePointMap,
                 const LO blockSize);

  virtual ~BlockCrsMatrix() {}



  Teuchos::RCP<const map_type> getDomainMap() const override;

  Teuchos::RCP<const map_type> getRangeMap() const override;

  Teuchos::RCP<const map_type> getRowMap() const override;

  Teuchos::RCP<const map_type> getColMap() const override;

  global_size_t getGlobalNumRows() const override;

  size_t getLocalNumRows() const override;

  size_t getLocalMaxNumRowEntries() const override;

  void
  apply(const mv_type& X,
        mv_type& Y,
        Teuchos::ETransp mode = Teuchos::NO_TRANS,
        Scalar alpha          = Teuchos::ScalarTraits<Scalar>::one(),
        Scalar beta           = Teuchos::ScalarTraits<Scalar>::zero()) const override;

  bool hasTransposeApply() const override {
    // FIXME (mfh 04 May 2014) Transpose and conjugate transpose modes
    // are not implemented yet.  Fill in applyBlockTrans() to fix this.
    return false;
  }

  void setAllToScalar(const Scalar& alpha);



  std::string description() const override;

  void
  describe(Teuchos::FancyOStream& out,
           const Teuchos::EVerbosityLevel verbLevel) const override;



  virtual LO getBlockSize() const override { return blockSize_; }

  virtual Teuchos::RCP<const ::Tpetra::RowGraph<LO, GO, Node> > getGraph() const override;

  const crs_graph_type& getCrsGraph() const { return graph_; }

  void
  applyBlock(const BlockMultiVector<Scalar, LO, GO, Node>& X,
             BlockMultiVector<Scalar, LO, GO, Node>& Y,
             Teuchos::ETransp mode = Teuchos::NO_TRANS,
             const Scalar alpha    = Teuchos::ScalarTraits<Scalar>::one(),
             const Scalar beta     = Teuchos::ScalarTraits<Scalar>::zero());

  void
  importAndFillComplete(Teuchos::RCP<BlockCrsMatrix<Scalar, LO, GO, Node> >& destMatrix,
                        const Import<LO, GO, Node>& importer) const;

  void
  exportAndFillComplete(Teuchos::RCP<BlockCrsMatrix<Scalar, LO, GO, Node> >& destMatrix,
                        const Export<LO, GO, Node>& exporter) const;

  LO replaceLocalValues(const LO localRowInd,
                        const LO colInds[],
                        const Scalar vals[],
                        const LO numColInds) const;

  LO sumIntoLocalValues(const LO localRowInd,
                        const LO colInds[],
                        const Scalar vals[],
                        const LO numColInds) const;

  void
  getLocalRowView(LO LocalRow,
                  local_inds_host_view_type& indices,
                  values_host_view_type& values) const override;

  void
  getLocalRowViewNonConst(LO LocalRow,
                          local_inds_host_view_type& indices,
                          nonconst_values_host_view_type& values) const;

  virtual void
  getLocalRowCopy(LO LocalRow,
                  nonconst_local_inds_host_view_type& Indices,
                  nonconst_values_host_view_type& Values,
                  size_t& NumEntries) const override;
  little_block_type
  getLocalBlockDeviceNonConst(const LO localRowInd, const LO localColInd) const;

  little_block_host_type
  getLocalBlockHostNonConst(const LO localRowInd, const LO localColInd) const;

  LO getLocalRowOffsets(const LO localRowInd,
                        ptrdiff_t offsets[],
                        const LO colInds[],
                        const LO numColInds) const;

  LO replaceLocalValuesByOffsets(const LO localRowInd,
                                 const ptrdiff_t offsets[],
                                 const Scalar vals[],
                                 const LO numOffsets) const;

  LO absMaxLocalValuesByOffsets(const LO localRowInd,
                                const ptrdiff_t offsets[],
                                const Scalar vals[],
                                const LO numOffsets) const;

  LO sumIntoLocalValuesByOffsets(const LO localRowInd,
                                 const ptrdiff_t offsets[],
                                 const Scalar vals[],
                                 const LO numOffsets) const;

  size_t getNumEntriesInLocalRow(const LO localRowInd) const override;

  local_matrix_device_type getLocalMatrixDevice() const;

  bool localError() const {
    return *localError_;
  }

  std::string errorMessages() const {
    return (*errs_).is_null() ? std::string("") : (*errs_)->str();
  }

  void
  getLocalDiagOffsets(const Kokkos::View<size_t*, device_type,
                                         Kokkos::MemoryUnmanaged>& offsets) const;

  void
  getLocalDiagCopy(const Kokkos::View<impl_scalar_type***, device_type,
                                      Kokkos::MemoryUnmanaged>& diag,
                   const Kokkos::View<const size_t*, device_type,
                                      Kokkos::MemoryUnmanaged>& offsets) const;

 protected:
  LO absMaxLocalValues(const LO localRowInd,
                       const LO colInds[],
                       const Scalar vals[],
                       const LO numColInds) const;



  using buffer_device_type = typename DistObject<Scalar, LO, GO,
                                                 Node>::buffer_device_type;

  virtual bool checkSizes(const ::Tpetra::SrcDistObject& source) override;

  using dist_object_type::
      copyAndPermute;  

  virtual void
  copyAndPermute(const SrcDistObject& sourceObj,
                 const size_t numSameIDs,
                 const Kokkos::DualView<const local_ordinal_type*,
                                        buffer_device_type>& permuteToLIDs,
                 const Kokkos::DualView<const local_ordinal_type*,
                                        buffer_device_type>& permuteFromLIDs,
                 const CombineMode CM) override;

  using dist_object_type::packAndPrepare;  

  virtual void
  packAndPrepare(const SrcDistObject& sourceObj,
                 const Kokkos::DualView<const local_ordinal_type*,
                                        buffer_device_type>& exportLIDs,
                 Kokkos::DualView<packet_type*,
                                  buffer_device_type>& exports,
                 Kokkos::DualView<size_t*,
                                  buffer_device_type>
                     numPacketsPerLID,
                 size_t& constantNumPackets) override;

  using dist_object_type::unpackAndCombine;  

  virtual void
  unpackAndCombine(const Kokkos::DualView<const local_ordinal_type*,
                                          buffer_device_type>& importLIDs,
                   Kokkos::DualView<packet_type*,
                                    buffer_device_type>
                       imports,
                   Kokkos::DualView<size_t*,
                                    buffer_device_type>
                       numPacketsPerLID,
                   const size_t constantNumPackets,
                   const CombineMode combineMode) override;

 private:
  crs_graph_type graph_;
  Teuchos::RCP<crs_graph_type> graphRCP_;
  map_type rowMeshMap_;
  map_type domainPointMap_;
  map_type rangePointMap_;
  LO blockSize_;

  using graph_row_offset_host_type = typename crs_graph_type::local_graph_device_type::row_map_type::host_mirror_type;
  graph_row_offset_host_type ptrHost_;

  using graph_column_indices_host_type = typename crs_graph_type::local_graph_device_type::entries_type::host_mirror_type;
  graph_column_indices_host_type indHost_;

  using impl_scalar_type_dualview         = Kokkos::DualView<impl_scalar_type*, device_type>;
  using impl_scalar_type_wrapped_dualview = Details::WrappedDualView<impl_scalar_type_dualview>;
  mutable impl_scalar_type_wrapped_dualview val_;

  Teuchos::RCP<Teuchos::RCP<BMV> > X_colMap_;
  Teuchos::RCP<Teuchos::RCP<BMV> > Y_rowMap_;

  Teuchos::RCP<Teuchos::RCP<typename crs_graph_type::import_type> > pointImporter_;

  LO offsetPerBlock_;

  Teuchos::RCP<bool> localError_;

  Teuchos::RCP<Teuchos::RCP<std::ostringstream> > errs_;

  std::ostream& markLocalErrorAndGetStream();

  // //! Clear the local error state and stream.
  // void clearLocalErrorStateAndStream ();

 public:
  typename impl_scalar_type_dualview::t_host::const_type
  getValuesHost() const;

  typename impl_scalar_type_dualview::t_dev::const_type
  getValuesDevice() const;

  typename impl_scalar_type_dualview::t_host
  getValuesHostNonConst() const;

  typename impl_scalar_type_dualview::t_dev
  getValuesDeviceNonConst() const;

  typename impl_scalar_type_dualview::t_host::const_type
  getValuesHost(const LO& lclRow) const;

  typename impl_scalar_type_dualview::t_dev::const_type
  getValuesDevice(const LO& lclRow) const;

  typename impl_scalar_type_dualview::t_host
  getValuesHostNonConst(const LO& lclRow);

  typename impl_scalar_type_dualview::t_dev
  getValuesDeviceNonConst(const LO& lclRow);

 private:
  void
  applyBlockTrans(const BlockMultiVector<Scalar, LO, GO, Node>& X,
                  BlockMultiVector<Scalar, LO, GO, Node>& Y,
                  const Teuchos::ETransp mode,
                  const Scalar alpha,
                  const Scalar beta);

  void
  applyBlockNoTrans(const BlockMultiVector<Scalar, LO, GO, Node>& X,
                    BlockMultiVector<Scalar, LO, GO, Node>& Y,
                    const Scalar alpha,
                    const Scalar beta);

  void
  localApplyBlockNoTrans(const BlockMultiVector<Scalar, LO, GO, Node>& X,
                         BlockMultiVector<Scalar, LO, GO, Node>& Y,
                         const Scalar alpha,
                         const Scalar beta);

  LO findRelOffsetOfColumnIndex(const LO localRowIndex,
                                const LO colIndexToFind,
                                const LO hint = 0) const;

  LO offsetPerBlock() const;

  const_little_block_type
  getConstLocalBlockFromInput(const impl_scalar_type* val, const size_t pointOffset) const;

  little_block_type
  getNonConstLocalBlockFromInput(impl_scalar_type* val, const size_t pointOffset) const;

  little_block_host_type
  getNonConstLocalBlockFromInputHost(impl_scalar_type* val, const size_t pointOffset) const;

 public:
  virtual Teuchos::RCP<const Teuchos::Comm<int> > getComm() const override;

  virtual global_size_t getGlobalNumCols() const override;

  virtual size_t getLocalNumCols() const override;

  virtual GO getIndexBase() const override;

  virtual global_size_t getGlobalNumEntries() const override;

  virtual size_t getLocalNumEntries() const override;
  virtual size_t getNumEntriesInGlobalRow(GO globalRow) const override;

  virtual size_t getGlobalMaxNumRowEntries() const override;

  virtual bool hasColMap() const override;

  virtual bool isLocallyIndexed() const override;

  virtual bool isGloballyIndexed() const override;

  virtual bool isFillComplete() const override;

  virtual bool supportsRowViews() const override;



  virtual void
  getGlobalRowCopy(GO GlobalRow,
                   nonconst_global_inds_host_view_type& Indices,
                   nonconst_values_host_view_type& Values,
                   size_t& NumEntries) const override;
  virtual void
  getGlobalRowView(GO GlobalRow,
                   global_inds_host_view_type& indices,
                   values_host_view_type& values) const override;

  virtual void getLocalDiagCopy(::Tpetra::Vector<Scalar, LO, GO, Node>& diag) const override;



  virtual void leftScale(const ::Tpetra::Vector<Scalar, LO, GO, Node>& x) override;

  virtual void rightScale(const ::Tpetra::Vector<Scalar, LO, GO, Node>& x) override;

  virtual typename ::Tpetra::RowMatrix<Scalar, LO, GO, Node>::mag_type
  getFrobeniusNorm() const override;

  // Friend declaration for nonmember function.
  template <class BlockCrsMatrixType>
  friend Teuchos::RCP<BlockCrsMatrixType>
  Tpetra::importAndFillCompleteBlockCrsMatrix(const Teuchos::RCP<const BlockCrsMatrixType>& sourceMatrix,
                                              const Import<typename BlockCrsMatrixType::local_ordinal_type,
                                                           typename BlockCrsMatrixType::global_ordinal_type,
                                                           typename BlockCrsMatrixType::node_type>& importer);
  // Friend declaration for nonmember function.
  template <class BlockCrsMatrixType>
  friend Teuchos::RCP<BlockCrsMatrixType>
  Tpetra::exportAndFillCompleteBlockCrsMatrix(const Teuchos::RCP<const BlockCrsMatrixType>& sourceMatrix,
                                              const Export<typename BlockCrsMatrixType::local_ordinal_type,
                                                           typename BlockCrsMatrixType::global_ordinal_type,
                                                           typename BlockCrsMatrixType::node_type>& exporter);
};

template <class BlockCrsMatrixType>
Teuchos::RCP<BlockCrsMatrixType>
importAndFillCompleteBlockCrsMatrix(const Teuchos::RCP<const BlockCrsMatrixType>& sourceMatrix,
                                    const Import<typename BlockCrsMatrixType::local_ordinal_type,
                                                 typename BlockCrsMatrixType::global_ordinal_type,
                                                 typename BlockCrsMatrixType::node_type>& importer) {
  Teuchos::RCP<BlockCrsMatrixType> destMatrix;
  sourceMatrix->importAndFillComplete(destMatrix, importer);
  return destMatrix;
}

template <class BlockCrsMatrixType>
Teuchos::RCP<BlockCrsMatrixType>
exportAndFillCompleteBlockCrsMatrix(const Teuchos::RCP<const BlockCrsMatrixType>& sourceMatrix,
                                    const Export<typename BlockCrsMatrixType::local_ordinal_type,
                                                 typename BlockCrsMatrixType::global_ordinal_type,
                                                 typename BlockCrsMatrixType::node_type>& exporter) {
  Teuchos::RCP<BlockCrsMatrixType> destMatrix;
  sourceMatrix->exportAndFillComplete(destMatrix, exporter);
  return destMatrix;
}

}  // namespace Tpetra

#endif  // TPETRA_BLOCKCRSMATRIX_DECL_HPP