Tpetra_CrsMatrixMultiplyOp.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_CRSMATRIXMULTIPLYOP_HPP
#define TPETRA_CRSMATRIXMULTIPLYOP_HPP


#include "Tpetra_CrsMatrixMultiplyOp_fwd.hpp"
#include "Tpetra_CrsMatrix.hpp"
#include "Tpetra_Util.hpp"
#include "Tpetra_Details_Behavior.hpp"
#include "Tpetra_Details_Profiling.hpp"
#include "Tpetra_LocalCrsMatrixOperator.hpp"

namespace Tpetra {

  template <class Scalar,
            class MatScalar,
            class LocalOrdinal,
            class GlobalOrdinal,
            class Node>
  class CrsMatrixMultiplyOp :
    public Operator<Scalar, LocalOrdinal, GlobalOrdinal, Node>
  {
  public:
    using crs_matrix_type =
      CrsMatrix<MatScalar, LocalOrdinal, GlobalOrdinal, Node>;
    using map_type = Map<LocalOrdinal, GlobalOrdinal, Node>;

  private:
    using local_matrix_device_type =
      typename crs_matrix_type::local_matrix_device_type;

  public:


    CrsMatrixMultiplyOp (const Teuchos::RCP<const crs_matrix_type>& A) :
      matrix_ (A),
      localMultiply_ (std::make_shared<local_matrix_device_type> (
                                       A->getLocalMatrixDevice ()))
    {}

    ~CrsMatrixMultiplyOp () override = default;



    void
    apply (const MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node>& X,
           MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node>& Y,
           Teuchos::ETransp mode = Teuchos::NO_TRANS,
           Scalar alpha = Teuchos::ScalarTraits<Scalar>::one (),
           Scalar beta = Teuchos::ScalarTraits<Scalar>::zero ()) const override
    {
      TEUCHOS_TEST_FOR_EXCEPTION
        (! matrix_->isFillComplete (), std::runtime_error,
         Teuchos::typeName (*this) << "::apply(): underlying matrix is not fill-complete.");
      TEUCHOS_TEST_FOR_EXCEPTION
        (X.getNumVectors () != Y.getNumVectors (), std::runtime_error,
         Teuchos::typeName (*this) << "::apply(X,Y): X and Y must have the same number of vectors.");
      TEUCHOS_TEST_FOR_EXCEPTION
        (Teuchos::ScalarTraits<Scalar>::isComplex && mode == Teuchos::TRANS, std::logic_error,
         Teuchos::typeName (*this) << "::apply() does not currently support transposed multiplications for complex scalar types.");
      if (mode == Teuchos::NO_TRANS) {
        applyNonTranspose (X, Y, alpha, beta);
      }
      else {
        applyTranspose (X, Y, mode, alpha, beta);
      }
    }

    bool hasTransposeApply() const override {
      return true;
    }

    Teuchos::RCP<const map_type> getDomainMap () const override {
      return matrix_->getDomainMap ();
    }

    Teuchos::RCP<const map_type> getRangeMap () const override {
      return matrix_->getRangeMap ();
    }


  protected:
    typedef MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node> MV;

    const Teuchos::RCP<const crs_matrix_type> matrix_;

    LocalCrsMatrixOperator<Scalar, MatScalar, typename
                           crs_matrix_type::device_type> localMultiply_;

    mutable Teuchos::RCP<MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node> > importMV_;

    mutable Teuchos::RCP<MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node> > exportMV_;

    void
    applyTranspose (const MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node>& X_in,
                    MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node> &Y_in,
                    Teuchos::ETransp mode,
                    Scalar alpha,
                    Scalar beta) const
    {
      using Teuchos::null;
      using Teuchos::RCP;
      using Teuchos::rcp;
      using export_type = Export<LocalOrdinal, GlobalOrdinal, Node>;
      using import_type = Import<LocalOrdinal, GlobalOrdinal, Node>;
      using STS = Teuchos::ScalarTraits<Scalar>;
      typedef typename MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node>:: dual_view_type::t_dev nonconst_view_type;

      const size_t numVectors = X_in.getNumVectors();
      // because of Views, it is difficult to determine if X and Y point to the same data.
      // however, if they reference the exact same object, we will do the user the favor of copying X into new storage (with a warning)
      // we ony need to do this if we have trivial importers; otherwise, we don't actually apply the operator from X into Y
      RCP<const import_type> importer = matrix_->getGraph()->getImporter();
      RCP<const export_type> exporter = matrix_->getGraph()->getExporter();

      // some parameters for below
      const bool Y_is_replicated = ! Y_in.isDistributed ();
      const bool Y_is_overwritten = (beta == STS::zero ());
      const int myRank = matrix_->getComm ()->getRank ();
      if (Y_is_replicated && myRank != 0) {
        beta = STS::zero ();
      }

      // access X indirectly, in case we need to create temporary storage
      RCP<const MV> X;
      // currently, cannot multiply from multivector of non-constant stride
      if (! X_in.isConstantStride () && importer == null) {
        // generate a strided copy of X_in
        X = Teuchos::rcp (new MV (X_in, Teuchos::Copy));
      }
      else {
        // just temporary, so this non-owning RCP is okay
        X = Teuchos::rcpFromRef (X_in);
      }

      // set up import/export temporary multivectors
      if (importer != null) {
        if (importMV_ != null && importMV_->getNumVectors () != numVectors) {
          importMV_ = null;
        }
        if (importMV_ == null) {
          importMV_ = rcp (new MV (matrix_->getColMap (), numVectors));
        }
      }
      if (exporter != null) {
        if (exportMV_ != null && exportMV_->getNumVectors () != numVectors) {
          exportMV_ = null;
        }
        if (exportMV_ == null) {
          exportMV_ = rcp (new MV (matrix_->getRowMap (), numVectors));
        }
      }

      // If we have a non-trivial exporter, we must import elements that are permuted or are on other processors
      if (exporter != null) {
        exportMV_->doImport(X_in,*exporter,INSERT);
        X = exportMV_; // multiply out of exportMV_
      }

      auto X_lcl = X->getLocalViewDevice(Access::ReadOnly);

      // If we have a non-trivial importer, we must export elements that are permuted or belong to other processors
      // We will compute solution into the to-be-exported MV; get a view
      if (importer != null) {
        // Beta is zero here, so we clobber Y_lcl
        auto Y_lcl = importMV_->getLocalViewDevice(Access::OverwriteAll);

        localMultiply_.apply (X_lcl, Y_lcl, mode, alpha, STS::zero ());
        if (Y_is_overwritten) {
          Y_in.putScalar (STS::zero ());
        }
        else {
          Y_in.scale (beta);
        }
        Y_in.doExport (*importMV_, *importer, ADD_ASSIGN);
      }
      // otherwise, multiply into Y
      else {
        // can't multiply in-situ; can't multiply into non-strided multivector
        if (! Y_in.isConstantStride () || X.getRawPtr () == &Y_in) {
          // generate a strided copy of Y
          MV Y (Y_in, Teuchos::Copy);
          nonconst_view_type Y_lcl;
          if(Y_is_overwritten) Y_lcl = Y.getLocalViewDevice(Access::OverwriteAll);
          else                 Y_lcl = Y.getLocalViewDevice(Access::ReadWrite); 

          localMultiply_.apply (X_lcl, Y_lcl, mode, alpha, beta);
          Tpetra::deep_copy (Y_in, Y);
        }
        else {
          nonconst_view_type Y_lcl;
          if(Y_is_overwritten) Y_lcl = Y_in.getLocalViewDevice(Access::OverwriteAll);
          else                 Y_lcl = Y_in.getLocalViewDevice(Access::ReadWrite); 

          localMultiply_.apply (X_lcl, Y_lcl, mode, alpha, beta);
        }
      }
      // Handle case of rangemap being a local replicated map: in this case, sum contributions from each processor
      if (Y_is_replicated) {
        Y_in.reduce();
      }
    }

    void
    applyNonTranspose (const MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node>& X_in,
                       MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node>& Y_in,
                       Scalar alpha,
                       Scalar beta) const
    {
      using Tpetra::Details::ProfilingRegion;
      using Teuchos::NO_TRANS;
      using Teuchos::RCP;
      using Teuchos::rcp;
      using Teuchos::rcp_const_cast;
      using Teuchos::rcpFromRef;
      typedef Export<LocalOrdinal,GlobalOrdinal,Node> export_type;
      typedef Import<LocalOrdinal,GlobalOrdinal,Node> import_type;
      typedef Teuchos::ScalarTraits<Scalar> STS;
      typedef typename MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node>:: dual_view_type::t_dev nonconst_view_type;

      if (alpha == STS::zero ()) {
        if (beta == STS::zero ()) {
          Y_in.putScalar (STS::zero ());
        } else if (beta != STS::one ()) {
          Y_in.scale (beta);
        }
        return;
      }

      // It's possible that X is a view of Y or vice versa.  We don't
      // allow this (apply() requires that X and Y not alias one
      // another), but it's helpful to detect and work around this
      // case.  We don't try to to detect the more subtle cases (e.g.,
      // one is a subview of the other, but their initial pointers
      // differ).  We only need to do this if this matrix's Import is
      // trivial; otherwise, we don't actually apply the operator from
      // X into Y.

      RCP<const import_type> importer = matrix_->getGraph()->getImporter();
      RCP<const export_type> exporter = matrix_->getGraph()->getExporter();

      // If beta == 0, then the output MV will be overwritten; none of
      // its entries should be read.  (Sparse BLAS semantics say that we
      // must ignore any Inf or NaN entries in Y_in, if beta is zero.)
      // This matters if we need to do an Export operation; see below.
      const bool Y_is_overwritten = (beta == STS::zero());

      // We treat the case of a replicated MV output specially.
      const bool Y_is_replicated =
        (! Y_in.isDistributed () && matrix_->getComm ()->getSize () != 1);

      // This is part of the special case for replicated MV output.
      // We'll let each process do its thing, but do an all-reduce at
      // the end to sum up the results.  Setting beta=0 on all
      // processes but Proc 0 makes the math work out for the
      // all-reduce.  (This assumes that the replicated data is
      // correctly replicated, so that the data are the same on all
      // processes.)
      if (Y_is_replicated && matrix_->getComm ()->getRank () > 0) {
        beta = STS::zero();
      }

      // Temporary MV for Import operation.  After the block of code
      // below, this will be an (Imported if necessary) column Map MV
      // ready to give to localMultiply_.apply(...).
      RCP<const MV> X_colMap;
      if (importer.is_null ()) {
        if (! X_in.isConstantStride ()) {
          // Not all sparse mat-vec kernels can handle an input MV with
          // nonconstant stride correctly, so we have to copy it in that
          // case into a constant stride MV.  To make a constant stride
          // copy of X_in, we force creation of the column (== domain)
          // Map MV (if it hasn't already been created, else fetch the
          // cached copy).  This avoids creating a new MV each time.
          RCP<MV> X_colMapNonConst = getColumnMapMultiVector (X_in, true);
          Tpetra::deep_copy (*X_colMapNonConst, X_in);
          X_colMap = rcp_const_cast<const MV> (X_colMapNonConst);
        }
        else {
          // The domain and column Maps are the same, so do the local
          // multiply using the domain Map input MV X_in.
          X_colMap = rcpFromRef (X_in);
        }
      }
      else { // need to Import source (multi)vector
        ProfilingRegion regionImport ("Tpetra::CrsMatrixMultiplyOp::apply: Import");

        // We're doing an Import anyway, which will copy the relevant
        // elements of the domain Map MV X_in into a separate column Map
        // MV.  Thus, we don't have to worry whether X_in is constant
        // stride.
        RCP<MV> X_colMapNonConst = getColumnMapMultiVector (X_in);

        // Import from the domain Map MV to the column Map MV.
        X_colMapNonConst->doImport (X_in, *importer, INSERT);
        X_colMap = rcp_const_cast<const MV> (X_colMapNonConst);
      }

      // Temporary MV for doExport (if needed), or for copying a
      // nonconstant stride output MV into a constant stride MV.  This
      // is null if we don't need the temporary MV, that is, if the
      // Export is trivial (null).
      RCP<MV> Y_rowMap = getRowMapMultiVector (Y_in);
      
      auto X_lcl = X_colMap->getLocalViewDevice(Access::ReadOnly);

      // If we have a nontrivial Export object, we must perform an
      // Export.  In that case, the local multiply result will go into
      // the row Map multivector.  We don't have to make a
      // constant-stride version of Y_in in this case, because we had to
      // make a constant stride Y_rowMap MV and do an Export anyway.
      if (! exporter.is_null ()) {
        auto Y_lcl = Y_rowMap->getLocalViewDevice(Access::OverwriteAll);

        localMultiply_.apply (X_lcl, Y_lcl, NO_TRANS,
                              alpha, STS::zero ());
        {
          ProfilingRegion regionExport ("Tpetra::CrsMatrixMultiplyOp::apply: Export");

          // If we're overwriting the output MV Y_in completely (beta
          // == 0), then make sure that it is filled with zeros before
          // we do the Export.  Otherwise, the ADD combine mode will
          // use data in Y_in, which is supposed to be zero.
          if (Y_is_overwritten) {
            Y_in.putScalar (STS::zero());
          }
          else {
            // Scale the output MV by beta, so that the Export sums in
            // the mat-vec contribution: Y_in = beta*Y_in + alpha*A*X_in.
            Y_in.scale (beta);
          }
          // Do the Export operation.
          Y_in.doExport (*Y_rowMap, *exporter, ADD_ASSIGN);
        }
      }
      else { // Don't do an Export: row Map and range Map are the same.
        //
        // If Y_in does not have constant stride, or if the column Map
        // MV aliases Y_in, then we can't let the kernel write directly
        // to Y_in.  Instead, we have to use the cached row (== range)
        // Map MV as temporary storage.
        //
        // FIXME (mfh 05 Jun 2014, mfh 07 Dec 2018) This test for
        // aliasing only tests if the user passed in the same
        // MultiVector for both X and Y.  It won't detect whether one
        // MultiVector views the other.  We should also check the
        // MultiVectors' raw data pointers.
        if (! Y_in.isConstantStride () || X_colMap.getRawPtr () == &Y_in) {
          // Force creating the MV if it hasn't been created already.
          // This will reuse a previously created cached MV.
          Y_rowMap = getRowMapMultiVector (Y_in, true);

          // If beta == 0, we don't need to copy Y_in into Y_rowMap,
          // since we're overwriting it anyway.
          if (beta != STS::zero ()) {
            Tpetra::deep_copy (*Y_rowMap, Y_in);
          }
          nonconst_view_type Y_lcl;
          if(Y_is_overwritten) Y_lcl = Y_rowMap->getLocalViewDevice(Access::OverwriteAll);
          else                 Y_lcl = Y_rowMap->getLocalViewDevice(Access::ReadWrite); 

          localMultiply_.apply (X_lcl, Y_lcl, NO_TRANS, alpha, beta);
          Tpetra::deep_copy (Y_in, *Y_rowMap);
        }
        else {
          nonconst_view_type Y_lcl;
          if(Y_is_overwritten) Y_lcl = Y_in.getLocalViewDevice(Access::OverwriteAll);
          else                 Y_lcl = Y_in.getLocalViewDevice(Access::ReadWrite); 

          localMultiply_.apply (X_lcl, Y_lcl, NO_TRANS, alpha, beta);
        }
      }

      // If the range Map is a locally replicated Map, sum up
      // contributions from each process.  We set beta = 0 on all
      // processes but Proc 0 initially, so this will handle the scaling
      // factor beta correctly.
      if (Y_is_replicated) {
        ProfilingRegion regionReduce ("Tpetra::CrsMatrixMultiplyOp::apply: Reduce Y");
        Y_in.reduce ();
      }
    }

  private:
    Teuchos::RCP<MV>
    getColumnMapMultiVector (const MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node>& X_domainMap,
                             const bool force = false) const
    {
      using Teuchos::null;
      using Teuchos::RCP;
      using Teuchos::rcp;
      typedef Import<LocalOrdinal,GlobalOrdinal,Node> import_type;

      const size_t numVecs = X_domainMap.getNumVectors ();
      RCP<const import_type> importer = matrix_->getGraph ()->getImporter ();
      RCP<const map_type> colMap = matrix_->getColMap ();

      RCP<MV> X_colMap; // null by default

      // If the Import object is trivial (null), then we don't need a
      // separate column Map multivector.  Just return null in that
      // case.  The caller is responsible for knowing not to use the
      // returned null pointer.
      //
      // If the Import is nontrivial, then we do need a separate
      // column Map multivector for the Import operation.  Check in
      // that case if we have to (re)create the column Map
      // multivector.
      if (! importer.is_null () || force) {
        if (importMV_.is_null () || importMV_->getNumVectors () != numVecs) {
          X_colMap = rcp (new MV (colMap, numVecs));

          // Cache the newly created multivector for later reuse.
          importMV_ = X_colMap;
        }
        else { // Yay, we can reuse the cached multivector!
          X_colMap = importMV_;
          // mfh 09 Jan 2013: We don't have to fill with zeros first,
          // because the Import uses INSERT combine mode, which overwrites
          // existing entries.
          //
          //X_colMap->putScalar (STS::zero ());
        }
      }
      return X_colMap;
    }

    Teuchos::RCP<MV>
    getRowMapMultiVector (const MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node>& Y_rangeMap,
                          const bool force = false) const
    {
      using Teuchos::null;
      using Teuchos::RCP;
      using Teuchos::rcp;
      typedef Export<LocalOrdinal,GlobalOrdinal,Node> export_type;

      const size_t numVecs = Y_rangeMap.getNumVectors ();
      RCP<const export_type> exporter = matrix_->getGraph ()->getExporter ();
      RCP<const map_type> rowMap = matrix_->getRowMap ();

      RCP<MV> Y_rowMap; // null by default

      // If the Export object is trivial (null), then we don't need a
      // separate row Map multivector.  Just return null in that case.
      // The caller is responsible for knowing not to use the returned
      // null pointer.
      //
      // If the Export is nontrivial, then we do need a separate row
      // Map multivector for the Export operation.  Check in that case
      // if we have to (re)create the row Map multivector.
      if (! exporter.is_null () || force) {
        if (exportMV_.is_null () || exportMV_->getNumVectors () != numVecs) {
          Y_rowMap = rcp (new MV (rowMap, numVecs));
          exportMV_ = Y_rowMap; // Cache the newly created MV for later reuse
        }
        else { // Yay, we can reuse the cached multivector!
          Y_rowMap = exportMV_;
        }
      }
      return Y_rowMap;
    }
  };

  template <class OpScalar,
            class MatScalar,
            class LocalOrdinal,
            class GlobalOrdinal,
            class Node>
  Teuchos::RCP<
    CrsMatrixMultiplyOp<OpScalar, MatScalar, LocalOrdinal, GlobalOrdinal, Node> >
  createCrsMatrixMultiplyOp (const Teuchos::RCP<
    const CrsMatrix<MatScalar, LocalOrdinal, GlobalOrdinal, Node> >& A)
  {
    typedef CrsMatrixMultiplyOp<OpScalar, MatScalar, LocalOrdinal,
      GlobalOrdinal, Node> op_type;
    return Teuchos::rcp (new op_type (A));
  }

} // end of namespace Tpetra

#endif // TPETRA_CRSMATRIXMULTIPLYOP_HPP