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

// clang-format off
#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"

#if KOKKOSKERNELS_VERSION >= 40299
#include "Tpetra_Details_MatrixApplyHelper.hpp"
#endif

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
}

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::HostMirror 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>;

#if KOKKOSKERNELS_VERSION >= 40299
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;
    }

#endif

  using local_matrix_host_type =
    typename local_matrix_device_type::HostMirror;



  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;

  // clang-format on
  using dist_object_type::
      copyAndPermute; 
                      // clang-format off

  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;

  // clang-format on
  using dist_object_type::packAndPrepare; 
                                          // clang-format off

  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;

  // clang-format on
  using dist_object_type::unpackAndCombine; 
                                            // clang-format off

  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::HostMirror;
  graph_row_offset_host_type ptrHost_;

  using graph_column_indices_host_type =   typename crs_graph_type::local_graph_device_type::entries_type::HostMirror;
  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 ();

  template<class Device>
  struct is_cuda {
#if defined(KOKKOS_ENABLE_CUDA)
    // CudaHostPinnedSpace::execution_space ==
    // HostSpace::execution_space.  That's OK; it's host memory, that
    // just happens to be Cuda accessible.  But what if somebody gives
    // us Device<Cuda, CudaHostPinnedSpace>?  It looks like they mean
    // to run on device then, so we should sync to device.
    static constexpr bool value =
      std::is_same<typename Device::execution_space, Kokkos::Cuda>::value;
    // Gonna badly fake this here for other execspaces
#elif defined(KOKKOS_ENABLE_HIP)
    static constexpr bool value =
      std::is_same<typename Device::execution_space, Kokkos::HIP>::value;
#elif defined(KOKKOS_ENABLE_SYCL)
    static constexpr bool value =
      std::is_same<typename Device::execution_space, Kokkos::Experimental::SYCL>::value;
#else
    static constexpr bool value = false;
#endif
  };

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