CGSolve.hpp

/*
//@HEADER
// ************************************************************************
//
//                        Kokkos v. 3.0
//       Copyright (2020) National Technology & Engineering
//               Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// 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 NTESS "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 NTESS 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 Siva Rajamanickam (srajama@sandia.gov)
//
// ************************************************************************
//@HEADER
*/

#ifndef KOKKOS_EXAMPLE_CG_SOLVE
#define KOKKOS_EXAMPLE_CG_SOLVE

#include <cmath>
#include <limits>
#include <Kokkos_Core.hpp>
#include <KokkosBlas.hpp>
#include <KokkosSparse_spmv.hpp>
#include <Kokkos_Timer.hpp>

#include <WrapMPI.hpp>

//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

namespace Kokkos {
namespace Example {

template< class ImportType , class SparseMatrixType , class VectorType , class TagType = void >
struct CGSolve ;


template< class ImportType , class SparseMatrixType , class VectorType >
struct CGSolve< ImportType , SparseMatrixType , VectorType ,
  typename std::enable_if<(
    Kokkos::Impl::is_view< VectorType >::value &&
    VectorType::rank == 1
  )>::type >
{
  typedef typename VectorType::value_type scalar_type ;
  typedef typename VectorType::execution_space execution_space;

  size_t iteration ;
  double iter_time ;
  double matvec_time ;
  double norm_res ;

  CGSolve( const ImportType       & import ,
           const SparseMatrixType & A ,
           const VectorType       & b ,
           const VectorType       & x ,
           const size_t             maximum_iteration = 200 ,
           const double             tolerance = std::numeric_limits<double>::epsilon() )
    : iteration(0)
    , iter_time(0)
    , matvec_time(0)
    , norm_res(0)
  {
    const size_t count_owned = import.count_owned ;
    const size_t count_total = import.count_owned + import.count_receive;

    // Need input vector to matvec to be owned + received
    VectorType pAll ( "cg::p" , count_total );

    VectorType p = Kokkos::subview( pAll , std::pair<size_t,size_t>(0,count_owned) );
    VectorType r ( "cg::r" , count_owned );
    VectorType Ap( "cg::Ap", count_owned );

    /* r = b - A * x ; */

    /* p  = x       */  Kokkos::deep_copy( p , x );
    /* import p     */  import( pAll );
    /* Ap = A * p   */  KokkosSparse::spmv( "N" , 1.0 , A , pAll , 0.0 , Ap);
    /* b - Ap => r  */  KokkosBlas::update( 1.0 , b , -1.0 , Ap , 0.0 , r);
    /* p  = r       */  Kokkos::deep_copy( p , r );

    double old_rdot = Kokkos::Example::all_reduce( KokkosBlas::dot( r , r ) , import.comm );

    norm_res  = sqrt( old_rdot );
    iteration = 0 ;

    Kokkos::Impl::Timer wall_clock ;
    Kokkos::Impl::Timer timer;

    while ( tolerance < norm_res && iteration < maximum_iteration ) {

      /* pAp_dot = dot( p , Ap = A * p ) */

      timer.reset();
      /* import p    */  import( pAll );
      /* Ap = A * p  */  KokkosSparse::spmv( "N", 1.0, A , pAll, 0.0, Ap);
      execution_space().fence();
      matvec_time += timer.seconds();

      const double pAp_dot = Kokkos::Example::all_reduce( KokkosBlas::dot( p , Ap ) , import.comm );
      const double alpha   = old_rdot / pAp_dot ;

      /* x +=  alpha * p ;  */ KokkosBlas::axpby( alpha, p  , 1.0 , x );
      /* r += -alpha * Ap ; */ KokkosBlas::axpby(-alpha, Ap , 1.0 , r );

      const double r_dot = Kokkos::Example::all_reduce( KokkosBlas::dot( r , r ) , import.comm );
      const double beta  = r_dot / old_rdot ;

      /* p = r + beta * p ; */ KokkosBlas::axpby( 1.0 , r , beta , p );

      norm_res = std::sqrt( old_rdot = r_dot );

      ++iteration ;
    }

    execution_space().fence();
    iter_time = wall_clock.seconds();
  }
};

} // namespace Example
} // namespace Kokkos

//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

#endif /* #ifndef KOKKOS_EXAMPLE_CG_SOLVE */