Tpetra_TsqrAdaptor.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_TSQRADAPTOR_HPP
#define TPETRA_TSQRADAPTOR_HPP


#include "Tpetra_ConfigDefs.hpp"

#ifdef HAVE_TPETRA_TSQR
#  include "Tsqr_NodeTsqrFactory.hpp" // create intranode TSQR object
#  include "Tsqr.hpp" // full (internode + intranode) TSQR
#  include "Tsqr_DistTsqr.hpp" // internode TSQR
// Subclass of TSQR::MessengerBase, implemented using Teuchos
// communicator template helper functions
#  include "Tsqr_TeuchosMessenger.hpp"
#  include "Tpetra_MultiVector.hpp"
#  include "Teuchos_ParameterListAcceptorDefaultBase.hpp"
#  include <stdexcept>

namespace Tpetra {

  template<class MV>
  class TsqrAdaptor : public Teuchos::ParameterListAcceptorDefaultBase {
  public:
    using scalar_type = typename MV::scalar_type;
    using ordinal_type = typename MV::local_ordinal_type;
    using dense_matrix_type =
      Teuchos::SerialDenseMatrix<ordinal_type, scalar_type>;
    using magnitude_type =
      typename Teuchos::ScalarTraits<scalar_type>::magnitudeType;

  private:
    using node_tsqr_factory_type =
      TSQR::NodeTsqrFactory<scalar_type, ordinal_type,
                            typename MV::device_type>;
    using node_tsqr_type = TSQR::NodeTsqr<ordinal_type, scalar_type>;
    using dist_tsqr_type = TSQR::DistTsqr<ordinal_type, scalar_type>;
    using tsqr_type = TSQR::Tsqr<ordinal_type, scalar_type>;

    TSQR::MatView<ordinal_type, scalar_type>
    get_mat_view(MV& X)
    {
      TEUCHOS_ASSERT( ! tsqr_.is_null() );
      // FIXME (mfh 18 Oct 2010, 22 Dec 2019) Check Teuchos::Comm<int>
      // object in Q to make sure it is the same communicator as the
      // one we are using in our dist_tsqr_type implementation.

      const ordinal_type lclNumRows(X.getLocalLength());
      const ordinal_type numCols(X.getNumVectors());
      scalar_type* X_ptr = nullptr;
      // LAPACK and BLAS functions require "LDA" >= 1, even if the
      // corresponding matrix dimension is zero.
      ordinal_type X_stride = 1;
      if(tsqr_->wants_device_memory()) {
        auto X_view = X.getLocalViewDevice(Access::ReadWrite);
        X_ptr = reinterpret_cast<scalar_type*>(X_view.data());
        X_stride = static_cast<ordinal_type>(X_view.stride(1));
        if(X_stride == 0) {
          X_stride = ordinal_type(1); // see note above
        }
      }
      else {
        auto X_view = X.getLocalViewHost(Access::ReadWrite);
        X_ptr = reinterpret_cast<scalar_type*>(X_view.data());
        X_stride = static_cast<ordinal_type>(X_view.stride(1));
        if(X_stride == 0) {
          X_stride = ordinal_type(1); // see note above
        }
      }
      using mat_view_type = TSQR::MatView<ordinal_type, scalar_type>;
      return mat_view_type(lclNumRows, numCols, X_ptr, X_stride);
    }

  public:
    TsqrAdaptor(const Teuchos::RCP<Teuchos::ParameterList>& plist) :
      nodeTsqr_(node_tsqr_factory_type::getNodeTsqr()),
      distTsqr_(new dist_tsqr_type),
      tsqr_(new tsqr_type(nodeTsqr_, distTsqr_))
    {
      setParameterList(plist);
    }

    TsqrAdaptor() :
      nodeTsqr_(node_tsqr_factory_type::getNodeTsqr()),
      distTsqr_(new dist_tsqr_type),
      tsqr_(new tsqr_type(nodeTsqr_, distTsqr_))
    {
      setParameterList(Teuchos::null);
    }

    Teuchos::RCP<const Teuchos::ParameterList>
    getValidParameters() const
    {
      if(defaultParams_.is_null()) {
        auto params = Teuchos::parameterList("TSQR implementation");
        params->set("NodeTsqr", *(nodeTsqr_->getValidParameters()));
        params->set("DistTsqr", *(distTsqr_->getValidParameters()));
        defaultParams_ = params;
      }
      return defaultParams_;
    }

    void
    setParameterList(const Teuchos::RCP<Teuchos::ParameterList>& plist)
    {
      auto params = plist.is_null() ?
        Teuchos::parameterList(*getValidParameters()) : plist;
      using Teuchos::sublist;
      nodeTsqr_->setParameterList(sublist(params, "NodeTsqr"));
      distTsqr_->setParameterList(sublist(params, "DistTsqr"));

      this->setMyParamList(params);
    }

    void
    factorExplicit(MV& A,
                   MV& Q,
                   dense_matrix_type& R,
                   const bool forceNonnegativeDiagonal=false)
    {
      TEUCHOS_TEST_FOR_EXCEPTION
        (! A.isConstantStride(), std::invalid_argument, "TsqrAdaptor::"
         "factorExplicit: Input MultiVector A must have constant stride.");
      TEUCHOS_TEST_FOR_EXCEPTION
        (! Q.isConstantStride(), std::invalid_argument, "TsqrAdaptor::"
         "factorExplicit: Input MultiVector Q must have constant stride.");
      prepareTsqr(Q); // Finish initializing TSQR.
      TEUCHOS_ASSERT( ! tsqr_.is_null() );

      auto A_view = get_mat_view(A);
      auto Q_view = get_mat_view(Q);
      constexpr bool contiguousCacheBlocks = false;
      tsqr_->factorExplicitRaw(A_view.extent(0),
                               A_view.extent(1),
                               A_view.data(), A_view.stride(1),
                               Q_view.data(), Q_view.stride(1),
                               R.values(), R.stride(),
                               contiguousCacheBlocks,
                               forceNonnegativeDiagonal);
    }

    int
    revealRank(MV& Q,
               dense_matrix_type& R,
               const magnitude_type& tol)
    {
      TEUCHOS_TEST_FOR_EXCEPTION
        (! Q.isConstantStride(), std::invalid_argument, "TsqrAdaptor::"
         "revealRank: Input MultiVector Q must have constant stride.");
      prepareTsqr(Q); // Finish initializing TSQR.

      auto Q_view = get_mat_view(Q);
      constexpr bool contiguousCacheBlocks = false;
      return tsqr_->revealRankRaw(Q_view.extent(0),
                                  Q_view.extent(1),
                                  Q_view.data(), Q_view.stride(1),
                                  R.values(), R.stride(),
                                  tol, contiguousCacheBlocks);
    }

  private:
    Teuchos::RCP<node_tsqr_type> nodeTsqr_;

    Teuchos::RCP<dist_tsqr_type> distTsqr_;

    Teuchos::RCP<tsqr_type> tsqr_;

    mutable Teuchos::RCP<const Teuchos::ParameterList> defaultParams_;

    bool ready_ = false;

    void
    prepareTsqr(const MV& mv)
    {
      if(! ready_) {
        prepareDistTsqr(mv);
        ready_ = true;
      }
    }

    void
    prepareDistTsqr(const MV& mv)
    {
      using Teuchos::RCP;
      using Teuchos::rcp_implicit_cast;
      using mess_type = TSQR::TeuchosMessenger<scalar_type>;
      using base_mess_type = TSQR::MessengerBase<scalar_type>;

      auto comm = mv.getMap()->getComm();
      RCP<mess_type> mess(new mess_type(comm));
      auto messBase = rcp_implicit_cast<base_mess_type>(mess);
      distTsqr_->init(messBase);
    }
  };

} // namespace Tpetra

#endif // HAVE_TPETRA_TSQR

#endif // TPETRA_TSQRADAPTOR_HPP