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