Xpetra_TpetraHalfPrecisionOperator.hpp

Go to the documentation of this file.

// @HEADER
//
// ***********************************************************************
//
//        MueLu: A package for multigrid based preconditioning
//                  Copyright 2012 Sandia Corporation
//
// Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact
//                    Jonathan Hu       (jhu@sandia.gov)
//                    Andrey Prokopenko (aprokop@sandia.gov)
//                    Ray Tuminaro      (rstumin@sandia.gov)
//
// ***********************************************************************
//
// @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