Xpetra_TpetraHalfPrecisionOperator.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//        MueLu: A package for multigrid based preconditioning
//
// Copyright 2012 NTESS and the MueLu contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef XPETRA_TPETRAHALFPRECISIONOPEARTOR_HPP
#define XPETRA_TPETRAHALFPRECISIONOPEARTOR_HPP

#include "Xpetra_ConfigDefs.hpp"

#include <Teuchos_ScalarTraits.hpp>

#include <Tpetra_CrsMatrix.hpp>
#include <Xpetra_TpetraOperator.hpp>
#include <Xpetra_MultiVector.hpp>
#include <Xpetra_CrsMatrixWrap.hpp>
#include <Xpetra_MultiVectorFactory.hpp>

namespace Xpetra {

template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
RCP<Xpetra::Matrix<typename Teuchos::ScalarTraits<Scalar>::halfPrecision, LocalOrdinal, GlobalOrdinal, Node>>
convertToHalfPrecision(RCP<Xpetra::Matrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>>& A) {
#if defined(HAVE_TPETRA_INST_DOUBLE) && defined(HAVE_TPETRA_INST_FLOAT)
  using HalfScalar = typename Teuchos::ScalarTraits<Scalar>::halfPrecision;
  typedef typename Xpetra::CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node> XpCrs;

  RCP<XpCrs> xpCrs = Teuchos::rcp_dynamic_cast<Xpetra::CrsMatrixWrap<Scalar, LocalOrdinal, GlobalOrdinal, Node>>(A, true)->getCrsMatrix();
  auto tpCrs       = Teuchos::rcp_dynamic_cast<Xpetra::TpetraCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>>(xpCrs, true)->getTpetra_CrsMatrix();
  auto newTpCrs    = tpCrs->template convert<HalfScalar>();
  auto newXpCrs    = Teuchos::rcp(new Xpetra::TpetraCrsMatrix<HalfScalar, LocalOrdinal, GlobalOrdinal, Node>(newTpCrs));
  auto newA        = Teuchos::rcp(new Xpetra::CrsMatrixWrap<HalfScalar, LocalOrdinal, GlobalOrdinal, Node>(Teuchos::rcp_dynamic_cast<Xpetra::CrsMatrix<HalfScalar, LocalOrdinal, GlobalOrdinal, Node>>(newXpCrs)));

  return newA;
#else
  TEUCHOS_TEST_FOR_EXCEPTION(true, Xpetra::Exceptions::RuntimeError,
                             "Xpetra::convertToHalfPrecision only available for Tpetra with SC=double and SC=float enabled");
  TEUCHOS_UNREACHABLE_RETURN(false);
#endif
}

template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
RCP<Xpetra::MultiVector<typename Teuchos::ScalarTraits<Scalar>::halfPrecision, LocalOrdinal, GlobalOrdinal, Node>>
convertToHalfPrecision(RCP<Xpetra::MultiVector<Scalar, LocalOrdinal, GlobalOrdinal, Node>>& X) {
#if defined(HAVE_TPETRA_INST_DOUBLE) && defined(HAVE_TPETRA_INST_FLOAT)
  using HalfScalar = typename Teuchos::ScalarTraits<Scalar>::halfPrecision;

  auto tpX    = toTpetra(*X);
  auto newTpX = tpX.template convert<HalfScalar>();
  auto newX   = rcp(new Xpetra::TpetraMultiVector<HalfScalar, LocalOrdinal, GlobalOrdinal, Node>(newTpX));

  return newX;
#else
  TEUCHOS_TEST_FOR_EXCEPTION(true, Xpetra::Exceptions::RuntimeError,
                             "Xpetra::convertToHalfPrecision only available for Tpetra with SC=double and SC=float enabled");
  TEUCHOS_UNREACHABLE_RETURN(false);
#endif
}

template <class Scalar,
          class LocalOrdinal,
          class GlobalOrdinal,
          class Node = Tpetra::KokkosClassic::DefaultNode::DefaultNodeType>
class TpetraHalfPrecisionOperator : public Xpetra::Operator<Scalar, LocalOrdinal, GlobalOrdinal, Node> {
 public:
  typedef typename Teuchos::ScalarTraits<Scalar>::halfPrecision HalfScalar;



  TpetraHalfPrecisionOperator(const RCP<Xpetra::Operator<HalfScalar, LocalOrdinal, GlobalOrdinal, Node>>& op)
    : Op_(op) {
    Allocate(1);
  }

  void Allocate(int numVecs) {
    X_ = Xpetra::MultiVectorFactory<HalfScalar, LocalOrdinal, GlobalOrdinal, Node>::Build(Op_->getDomainMap(), numVecs);
    Y_ = Xpetra::MultiVectorFactory<HalfScalar, LocalOrdinal, GlobalOrdinal, Node>::Build(Op_->getRangeMap(), numVecs);
  }

  virtual ~TpetraHalfPrecisionOperator() {}


  const Teuchos::RCP<const Xpetra::Map<LocalOrdinal, GlobalOrdinal, Node>> getDomainMap() const {
    return Op_->getDomainMap();
  }

  const Teuchos::RCP<const Xpetra::Map<LocalOrdinal, GlobalOrdinal, Node>> getRangeMap() const {
    return Op_->getRangeMap();
  }


  void apply(const Xpetra::MultiVector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& X,
             Xpetra::MultiVector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Y,
             Teuchos::ETransp mode = Teuchos::NO_TRANS,
             Scalar alpha          = Teuchos::ScalarTraits<Scalar>::one(),
             Scalar beta           = Teuchos::ScalarTraits<Scalar>::one()) const {
    typedef Xpetra::TpetraMultiVector<HalfScalar, LocalOrdinal, GlobalOrdinal, Node> tMVHalf;
    Tpetra::deep_copy(*Teuchos::rcp_dynamic_cast<tMVHalf>(X_)->getTpetra_MultiVector(),
                      toTpetra(X));
    Op_->apply(*X_, *Y_, mode, Teuchos::as<HalfScalar>(alpha), Teuchos::as<HalfScalar>(beta));
    Tpetra::deep_copy(toTpetra(Y),
                      *Teuchos::rcp_dynamic_cast<tMVHalf>(Y_)->getTpetra_MultiVector());
  }

  bool hasTransposeApply() const { return false; }

  void residual(const Xpetra::MultiVector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& X,
                const Xpetra::MultiVector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& B,
                Xpetra::MultiVector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& R) const {
    using STS = Teuchos::ScalarTraits<Scalar>;
    R.update(STS::one(), B, STS::zero());
    this->apply(X, R, Teuchos::NO_TRANS, -STS::one(), STS::one());
  }



  RCP<Xpetra::Operator<HalfScalar, LocalOrdinal, GlobalOrdinal, Node>> GetHalfPrecisionOperator() const { return Op_; }

  void SetHalfPrecisionOperator(const RCP<Xpetra::Operator<HalfScalar, LocalOrdinal, GlobalOrdinal, Node>>& op) { Op_ = op; };


 private:
  RCP<Xpetra::Operator<HalfScalar, LocalOrdinal, GlobalOrdinal, Node>> Op_;
  RCP<Xpetra::MultiVector<HalfScalar, LocalOrdinal, GlobalOrdinal, Node>> X_, Y_;
};

}  // namespace Xpetra

#endif  // XPETRA_TPETRAHALFPRECISIONOPEARTOR_HPP