BelosKokkosDenseAdapter.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//                 Belos: Block Linear Solvers Package
//
// Copyright 2004-2016 NTESS and the Belos contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER
//
#ifndef BELOS_KOKKOS_DENSE_MAT_TRAITS_HPP
#define BELOS_KOKKOS_DENSE_MAT_TRAITS_HPP

 
#include "Teuchos_Assert.hpp"
#include "Teuchos_RCP.hpp"
#include "Teuchos_ScalarTraits.hpp"
#include "Teuchos_BLAS.hpp"
#include "Teuchos_LAPACK.hpp"
#include "BelosDenseMatTraits.hpp"
 
#include "Kokkos_DualView.hpp"
#include "Kokkos_Random.hpp"
#include "KokkosKernels_ArithTraits.hpp"
#include "KokkosBlas1_scal.hpp"
#include "KokkosBlas1_axpby.hpp"
 
#include <vector>
 
namespace Belos {

  template<typename V>
  void kokkos_transpose(const V& dst, const V& src)
  {
    Kokkos::parallel_for(Kokkos::MDRangePolicy<typename V::execution_space, Kokkos::Rank<2>>({0, 0}, {dst.extent(0), dst.extent(1)}),
    KOKKOS_LAMBDA(int i, int j)
    {
      dst(i, j) = KokkosKernels::ArithTraits<typename V::non_const_value_type>::conj( src(j, i) );
    });
  }

  template<class Scalar, class DM>
  class KokkosDenseSolver : public DenseSolver<Scalar, DM>
  {
  public:
    typedef typename KokkosKernels::ArithTraits<Scalar>::val_type IST; //Impl Scalar Type, as used in Tpetra
    typedef DenseMatTraits<Scalar, DM> DMT;



    KokkosDenseSolver() {}

    virtual ~KokkosDenseSolver() {}



    using DenseSolver<Scalar, DM>::setMatrix;

    using DenseSolver<Scalar, DM>::setVectors;



    using DenseSolver<Scalar, DM>::setSPD;


    using DenseSolver<Scalar, DM>::factorWithEquilibration;


    using DenseSolver<Scalar, DM>::solveWithTransposeFlag;




    int factor()
    {
      int INFO = 0;
 
      // Only factor matrix if it is new, otherwise factors have been computed
      if (newMatrix_)
      {
        Teuchos::LAPACK<int,Scalar> lapack;
 
        int M = DMT::GetNumRows(*A_);
        int N = DMT::GetNumCols(*A_);
        int Min_MN = TEUCHOS_MIN(M,N);
        int LDA = DMT::GetStride(*A_);
 
        IPIV_.resize( Min_MN );
 
        DMT::SyncDeviceToHost(*A_);
        Scalar * Aptr = DMT::GetRawHostPtr(*A_);
 
        if (equilibrate_)
        {
          // Compute equilibration scalings
          R_.resize(M);
          if (!spd_)
            C_.resize(N);
 
          MagnitudeType ROWCND, COLCND, AMAX;
          if (spd_)
            lapack.POEQU (M, Aptr, LDA, &R_[0], &ROWCND, &AMAX, &INFO);
          else
            lapack.GEEQU (M, N, Aptr, LDA, &R_[0], &C_[0], &ROWCND, &COLCND, &AMAX, &INFO);
 
          if (INFO)
            return INFO;
 
          // Apply equilibration to matrix
          if (spd_) {
            Scalar * ptr = 0;
            for (int j=0; j<N; j++) {
              ptr = Aptr + j*LDA;
              Scalar s1 = R_[j];
              for (int i=0; i<=j; i++) {
                *ptr = *ptr*s1*R_[i];
                ptr++;
              }
            }
          }
          else {
            Scalar * ptr = 0;
            for (int j=0; j<N; j++) {
              ptr = Aptr + j*LDA;
              Scalar s1 = C_[j];
              for (int i=0; i<M; i++) {
                *ptr = *ptr*s1*R_[i];
                ptr++;
              }
            }
          }
        }
 
        // Compute LU factor
        if (spd_) {
          lapack.POTRF('U', M, Aptr, LDA, &INFO);
        }
        else {
          lapack.GETRF(M, N, Aptr, LDA, &IPIV_[0], &INFO);
        }
 
        DMT::SyncHostToDevice(*A_);
      }
 
      return INFO;
    }


    int solve()
    {
      bool transpose = (TRANS_ != Teuchos::NO_TRANS) ? true : false;
 
      DMT::SyncDeviceToHost(*X_);
 
      // LAPACK overwrites RHS vector with solution vector, so copy if necessary
      if (B_ != X_)
        DMT::Assign(*X_, *B_); // Copy B to X if needed
 
      // Since B_ = X_, perform operations on X_.
 
      int M = DMT::GetNumRows(*X_);
      int NRHS = DMT::GetNumCols(*X_);
      int LDX = DMT::GetStride(*X_);
      Scalar * X = DMT::GetRawHostPtr(*X_);
 
      if (equilibrate_)
      {
        // Apply equilibration scalings to RHS vector
        MagnitudeType * R_tmp = (transpose && !spd_) ? &C_[0] : &R_[0];
 
        Scalar * ptr = 0;
        for (int j=0; j<NRHS; j++) {
          ptr = X + j*LDX;
          for (int i=0; i<M; i++) {
            *ptr = *ptr*R_tmp[i];
            ptr++;
          }
        }
      }
 
      int INFO = 0;
 
      int LDA = DMT::GetStride(*A_);
      Scalar * Aptr = DMT::GetRawHostPtr(*A_);
      Teuchos::LAPACK<int,Scalar> lapack;
 
      if (spd_) {
        lapack.POTRS('U', M, NRHS, Aptr, LDA, X, LDX, &INFO);
      }
      else {
        lapack.GETRS(Teuchos::ETranspChar[TRANS_], M, NRHS,
                     Aptr, LDA, &IPIV_[0], X, LDX, &INFO);
      }
 
      if (equilibrate_)
      {
        // Apply equilibration scalings to X vector
        MagnitudeType * C_tmp = (spd_ || transpose) ? &R_[0] : &C_[0];
 
        Scalar * ptr = 0;
        for (int j=0; j<NRHS; j++) {
          ptr = X + j*LDX;
            for (int i=0; i<M; i++) {
            *ptr = *ptr*C_tmp[i];
            ptr++;
          }
        }
      }
 
      // Synchronize solution vector to the device
      DMT::SyncHostToDevice(*X_);
 
      return INFO;
    }

 
  private:
 
    typedef typename Teuchos::ScalarTraits<Scalar>::magnitudeType MagnitudeType;
 
    std::vector<int> IPIV_;
    std::vector<MagnitudeType> R_, C_;
 
    using DenseSolver<Scalar, DM>::newMatrix_;
    using DenseSolver<Scalar, DM>::equilibrate_;
    using DenseSolver<Scalar, DM>::TRANS_;
    using DenseSolver<Scalar, DM>::spd_;
 
    using DenseSolver<Scalar, DM>::A_;
    using DenseSolver<Scalar, DM>::X_;
    using DenseSolver<Scalar, DM>::B_;
 
  };

  //
  template<class Scalar, class... Properties>
  class DenseMatTraits<Scalar, Kokkos::DualView<typename KokkosKernels::ArithTraits<Scalar>::val_type **,Properties...>>{
 
  public:
    typedef typename KokkosKernels::ArithTraits<Scalar>::val_type IST;
    using DM = Kokkos::DualView<IST**,Properties...>;


    static Teuchos::RCP<DM> Create() {
      return Teuchos::rcp(new DM("BelosDenseCreate",0,0));
    }

    static Teuchos::RCP<DM>
           Create( const int numRows, const int numCols, bool initZero = true) {
      if(initZero){
        return Teuchos::rcp(new DM("BelosDenseCreate2",numRows,numCols));
      }
      else {
        return Teuchos::rcp(new DM(Kokkos::view_alloc(Kokkos::WithoutInitializing,"BelosDenseCreate2"),numRows,numCols));
      }
    }

    static Teuchos::RCP<DM>
           CreateCopy(const DM & dm, bool transpose=false)
    {
      Teuchos::RCP<DM> tmpCopyRCP = Teuchos::null;
 
      if (transpose) {
        // want tmpCopyRCP to end up as dm^H in the end. Prefer doing transpose on device
        tmpCopyRCP = Teuchos::rcp(new DM
                                 (Kokkos::view_alloc(Kokkos::WithoutInitializing,"BelosDenseCreateCopy"),dm.extent_int(1),dm.extent_int(0)));
        if(tmpCopyRCP->need_sync_device()) {
          // tmpCopyRCP is only up to date on the host
          kokkos_transpose(tmpCopyRCP->view_host(), dm.view_host());
          tmpCopyRCP->clear_sync_state();
          tmpCopyRCP->modify_host();
        }
        else {
          kokkos_transpose(tmpCopyRCP->view_device(), dm.view_device());
          tmpCopyRCP->clear_sync_state();
          tmpCopyRCP->modify_device();
        }
      }
      else {
        tmpCopyRCP = Teuchos::rcp(new DM
                                 (Kokkos::view_alloc(Kokkos::WithoutInitializing,"BelosDenseCreateCopy"),dm.extent_int(0),dm.extent_int(1)));
        Kokkos::deep_copy(*tmpCopyRCP, dm);
      }
      return tmpCopyRCP;
    }

    static Scalar* GetRawHostPtr(DM & dm ) {
      dm.sync_host();
      dm.modify_host();
      return reinterpret_cast<Scalar*>(dm.view_host().data());
      //TODO: Is there any way that the user could hold on to this pointer...
      // and everything works fine the first time they pass to LAPACK.
      // But then... they call MvTimesMatAddMv which syncs to device.
      // But then they keep the same pointer and pass to LAPACK again.
      // Then they call MvTimesMatAddMv... but since they didn't call this
      // function again, we miss the sync... See thread with Heidi on this.
    }

    static Scalar const * GetConstRawHostPtr(const DM & dm ) {
      // CAG: This is a bit naughty.
      const_cast<DM*>(&dm)->sync_host();
      return reinterpret_cast<Scalar const *>(dm.view_host().data());
    }

    //        Row and column indexing is zero-based.
    static Teuchos::RCP<DM>
    Subview( DM & source, int numRows, int numCols, int startRow=0, int startCol=0){
      return Teuchos::rcp(new DM(source,
          Kokkos::pair<int,int>(startRow,startRow+numRows), Kokkos::pair<int,int>(startCol,startCol+numCols)));
    }
 
    static Teuchos::RCP<const DM>
    SubviewConst( const DM& source, int numRows, int numCols, int startRow=0, int startCol=0){
      return Teuchos::rcp(new DM(source,
          Kokkos::pair<int,int>(startRow,startRow+numRows), Kokkos::pair<int,int>(startCol,startCol+numCols)));
    }

    static Teuchos::RCP<DM>
    SubviewCopy( const DM& source, int numRows, int numCols, int startRow=0, int startCol=0){
      //Maaybe we could get away with just a host copy here??
      //Hmmm... but it says we return a dual view.
      //Maybe it should return a dual view with only host data copied in. Require sync to work on device.
      //This is related to the functionality of the Assign function.
      auto tmpViewRCP = Teuchos::rcp(new DM
                                  (Kokkos::view_alloc(Kokkos::WithoutInitializing,"BelosDenseSubViewCopy"),numRows,numCols));
      // I am keeping this where it copies the whole view on host and device because:
      // a) I feel like we might be inviting some weird bugs later if we don't.
      // But TODO Clarify to developer that this function needs to work on both host and device.
      // b) It's not a host-device copy or vice versa. Its a copy from device to same device and from host to host.
      // So shouldn't add much extra overhead.
      Kokkos::deep_copy(*tmpViewRCP, Kokkos::subview(source,
          Kokkos::pair<int,int>(startRow,startRow+numRows), Kokkos::pair<int,int>(startCol,startCol+numCols)));
      return tmpViewRCP;
    }



    static int GetNumRows( const DM& dm ) {
      return dm.extent_int(0);
    }

    static int GetNumCols( const DM& dm ) {
      return dm.extent_int(1);
    }

    static int GetStride( const DM& dm ) {
      // Note: We force LayoutLeft, which is column major, so the stride_0 is always 1.
      // (This is the distance between two elts in same col, different rows.)
      // Lapack wants stride_1, the distance from one col to the next col if we stay
      // in the same row.
      int strides[8]; // There are 8 possible strides and all will be returned
      dm.stride(strides);
      return strides[1];
      //return dm.stride_1();  //This shortcut doesn't work for dualView.
    }


 
    /* \brief Reshaping method for changing the size of \c dm to have \c numRows rows and \c numCols columns.
     *        All values will be initialized to zero if the final argument is true.
     *        If the final argument is fale, the previous entries in
     *        the matrix will be maintained. For new entries that did not exist in the previous matrix, values will
     *        contain noise from memory.
    */
    static void Reshape( DM& dm, const int numRows, const int numCols, bool initZero = false) {
      if(initZero){
        dm.realloc(numRows,numCols);
        Kokkos::deep_copy(dm.view_device(), 0.0);
        dm.modify_device();
      } else{
        dm.resize(numRows,numCols); //keeps values in old array.
      }
    }



    static Scalar & Value( DM& dm, const int i, const int j )
    {
      // Mark as modified on host, since we don't know if it will be.
      dm.sync_host();
      dm.modify_host();
      return reinterpret_cast<Scalar&>((dm.view_host())(i,j));
    }

    static const Scalar & ValueConst( const DM& dm, const int i, const int j ) {
      // CAG: This is a bit naughty.
      const_cast<DM*>(&dm)->sync_host();
      return reinterpret_cast<Scalar const &>((dm.view_host())(i,j));
    }

    //
    //  \note The only Belos function that results in a need to sync to
    //  host is MvTransMv. You MUST call SyncDeviceToHost before calling
    //  any other DenseMatTraits functions after a call to MvTransMv.
    //  All DenseMatTraits functions assume the necessary data is on host
    //  and perform computations only on the host.
    //
    static void SyncDeviceToHost(DM & dm) {
      if(dm.need_sync_host()){
        if(dm.view_host().span_is_contiguous() && dm.view_device().span_is_contiguous()){
        dm.sync_host();}
        else{
          DM compat_view("compat view",dm.extent_int(0),dm.extent_int(1));
          Kokkos::deep_copy(compat_view,dm);
          compat_view.sync_host();
          Kokkos::deep_copy(dm,compat_view);
          dm.clear_sync_state();
        }
      }
    }
 
    static void SyncHostToDevice(DM & dm) {
      if(dm.need_sync_device()){
        if(dm.view_host().span_is_contiguous() && dm.view_device().span_is_contiguous()){
          dm.sync_device();
        }
        else{
          DM compat_view("compat view",dm.extent_int(0),dm.extent_int(1));
          Kokkos::deep_copy(compat_view,dm);
          compat_view.sync_device();
          Kokkos::deep_copy(dm,compat_view);
          dm.clear_sync_state();
        }
      }
    }



    static void Add( DM& thisDM, const DM& sourceDM) {
      thisDM.sync_device();
      // CAG: This is a bit naughty.
      const_cast<DM*>(&sourceDM)->sync_device();
      KokkosBlas::axpy(1.0,sourceDM.view_device(), thisDM.view_device()); //axpy(alpha,x,y), y = y + alpha*x
      thisDM.modify_device();
    }

    static void PutScalar( DM& dm, Scalar value = Teuchos::ScalarTraits<Scalar>::zero()){
      dm.clear_sync_state();
      Kokkos::deep_copy( dm.view_device(), value);
      dm.modify_device();
    }

    static void Scale( DM& dm, Scalar value) {
      dm.sync_device();
      KokkosBlas::scal( dm.view_device(), value, dm.view_device());
      dm.modify_device();
    }

    static void Randomize( DM& dm) {
      int rand_seed = std::rand();
      Kokkos::Random_XorShift64_Pool<> pool(rand_seed);
      dm.clear_sync_state();
      Kokkos::fill_random(dm.view_device(), pool, -1,1);
      dm.modify_device();
    }

    static void Assign( DM& dest, const DM& source) {
      Kokkos::deep_copy(dest,source);
    }

    static typename Teuchos::ScalarTraits<Scalar>::magnitudeType NormFrobenius(const DM& dm) {
      using KAT = KokkosKernels::ArithTraits<IST>;
      using mag_t = typename KAT::mag_type;
      // CAG: This is a bit naughty.
      const_cast<DM*>(&dm)->sync_device();
      mag_t frobNorm;
      Kokkos::parallel_reduce(Kokkos::MDRangePolicy<Kokkos::Rank<2>>({0, 0}, {dm.extent(0), dm.extent(1)}),
          KOKKOS_LAMBDA(size_t i, size_t j, mag_t& lfrobNorm)
          {
          mag_t absVal = KAT::abs((dm.view_device())(i, j));
          lfrobNorm += absVal * absVal;
          }, frobNorm);
      return Kokkos::sqrt(frobNorm);
    }

    static typename Teuchos::ScalarTraits<Scalar>::magnitudeType NormOne(const DM& dm) {
      using KAT = KokkosKernels::ArithTraits<IST>;
      using mag_t = typename KAT::mag_type;
      // CAG: This is a bit naughty.
      const_cast<DM*>(&dm)->sync_device();
      mag_t max_sum = 0;
 
      Kokkos::parallel_reduce(dm.extent(1), KOKKOS_LAMBDA(const int j, mag_t& norm) {
        mag_t sum = 0;
        for(int i = 0; i < dm.extent_int(0); i++){  //rows
          sum += KAT::abs((dm.view_device())(i,j));
        }
        norm = Kokkos::max(norm, sum);
      }, Kokkos::Max<mag_t>(max_sum));
      return KAT::abs(max_sum);
    }



    static Teuchos::RCP<DenseSolver<Scalar, DM>>
      createDenseSolver() {
 
      Teuchos::RCP<DenseSolver<Scalar, DM>> newSolver
        = Teuchos::rcp( new KokkosDenseSolver<Scalar, DM>() );
      return newSolver;
    }

 
  };
 
} // namespace Belos
 
#endif // end file BELOS_KOKKOS_DENSE_MAT_TRAITS_HPP