KokkosSparse_pcg.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 <Kokkos_Timer.hpp>

#include <Kokkos_Atomic.hpp>
#include <Kokkos_MemoryTraits.hpp>

#include <iostream>
#include "KokkosKernels_Handle.hpp"
#include <KokkosSparse_spmv.hpp>
#include <KokkosBlas.hpp>
#include <KokkosSparse_gauss_seidel.hpp>
#include <KokkosSparse_sor_sequential_impl.hpp>
//----------------------------------------------------------------------------
//----------------------------------------------------------------------------


//#define KK_TICTOCPRINT
namespace KokkosKernels {
namespace Experimental{
namespace Example {


struct CGSolveResult {
  size_t  iteration ;
  double  iter_time ;
  double  matvec_time ;
  double  norm_res ;
  double precond_time;
  double precond_init_time;
};

template< typename KernelHandle_t,
          typename crsMatrix_t,
          typename y_vector_t,
          typename x_vector_t
          >
void block_pcgsolve(
               KernelHandle_t &kh
            ,  const crsMatrix_t &point_crsMat
            ,  const crsMatrix_t &_block_crsMat, int block_size
            ,  const y_vector_t &y_vector
            ,  x_vector_t x_vector
            ,  const size_t  maximum_iteration = 200
            ,  const double  tolerance = std::numeric_limits<double>::epsilon()
            ,  CGSolveResult * result = 0
            ,  bool use_sgs = true)
{
  using namespace KokkosSparse;
  using namespace KokkosSparse::Experimental;
  typedef typename KernelHandle_t::HandleExecSpace Space;

  const size_t count_total = point_crsMat.numRows();

  size_t  iteration = 0 ;
  double  iter_time = 0 ;
  double  matvec_time = 0 ;
  double  norm_res = 0 ;
  double precond_time = 0;
  double precond_init_time = 0;

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

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

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

  // r = b - A * x ;
  // p  = x      
  Kokkos::deep_copy( p , x_vector );

  // Ap = A * p 
  KokkosSparse::spmv("N", 1, point_crsMat, pAll, 0, Ap);

  // r  = Ap
  Kokkos::deep_copy( r , Ap );

  // r = b - r
  KokkosBlas::axpby(1.0, y_vector, -1.0, r);

  // p  = r
  Kokkos::deep_copy( p , r );
;
  double old_rdot = KokkosBlas::dot( r , r );
  norm_res  = sqrt( old_rdot );

  int apply_count = 1;
  y_vector_t z;

  double precond_old_rdot = 1;
  //Kokkos::deep_copy( p , z );

  bool owner_handle = false;

  KernelHandle_t block_kh;
  block_kh.create_gs_handle();
  block_kh.get_point_gs_handle()->set_block_size(block_size);
    //block_kh.set_shmem_size(8032);
  if (use_sgs){
    if (kh.get_gs_handle() == NULL){
      owner_handle = true;
      kh.create_gs_handle();
    }

    timer.reset();

    //gauss_seidel_numeric
    //  (&kh, count_total, count_total, point_crsMat.graph.row_map, point_crsMat.graph.entries, point_crsMat.values);

    //Space().fence();
    //timer.reset();

    //block_kh.set_verbose(true);
    block_gauss_seidel_numeric
          (&block_kh, _block_crsMat.numRows(), _block_crsMat.numCols(), block_size, _block_crsMat.graph.row_map, _block_crsMat.graph.entries, _block_crsMat.values);

    precond_init_time += timer.seconds();

    z = y_vector_t( "pcg::z" , count_total );
    Space().fence();
    timer.reset();
    symmetric_block_gauss_seidel_apply
            (&block_kh, _block_crsMat.numRows(), _block_crsMat.numCols(),block_size,  _block_crsMat.graph.row_map, _block_crsMat.graph.entries, _block_crsMat.values,
            		z, r, true, true, 1.0, apply_count);

    //symmetric_gauss_seidel_apply
    //    (&kh, count_total, count_total, point_crsMat.graph.row_map, point_crsMat.graph.entries, point_crsMat.values, z, r, true, true, apply_count);
    Space().fence();
    precond_time += timer.seconds();
    precond_old_rdot = KokkosBlas::dot( r , z );
    Kokkos::deep_copy( p , z );
  }

  iteration = 0 ;

#ifdef KK_TICTOCPRINT

  std::cout << "norm_res:" << norm_res << " old_rdot:" << old_rdot<<  std::endl;

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


    timer.reset();
    //Ap = A * p
    KokkosSparse::spmv("N", 1, point_crsMat, pAll, 0, Ap);


    Space().fence();
    matvec_time += timer.seconds();

    //const double pAp_dot = Kokkos::Example::all_reduce( dot( count_owned , p , Ap ) , import.comm );
    //const double pAp_dot = dot<y_vector_t,y_vector_t, Space>( count_total , p , Ap ) ;

    // pAp_dot = dot(Ap , p);
    const double pAp_dot = KokkosBlas::dot( p , Ap ) ;


    double alpha  = 0;
    if (use_sgs){
      alpha = precond_old_rdot / pAp_dot ;
    }
    else {
      alpha = old_rdot / pAp_dot ;
    }

    // x +=  alpha * p ;
    KokkosBlas::axpby(alpha, p, 1.0, x_vector);

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

    const double r_dot = KokkosBlas::dot( r , r );

    const double beta_original  = r_dot / old_rdot ;
    double precond_r_dot = 1;
    double precond_beta = 1;
    if (use_sgs){
      Space().fence();
      timer.reset();
      symmetric_block_gauss_seidel_apply
                  (&block_kh, _block_crsMat.numRows(), _block_crsMat.numCols(),block_size, _block_crsMat.graph.row_map, _block_crsMat.graph.entries, _block_crsMat.values,
                  		z, r, true, true, 1.0, apply_count);

      //symmetric_gauss_seidel_apply(
      //    &kh,
      //    count_total, count_total,
      //    point_crsMat.graph.row_map,
      //    point_crsMat.graph.entries,
      //    point_crsMat.values, z, r, true,
      //    apply_count);

      Space().fence();
      precond_time += timer.seconds();
      precond_r_dot = KokkosBlas::dot(r , z );
      precond_beta  = precond_r_dot / precond_old_rdot ;
    }

    double beta  = 1;
    if (!use_sgs){
      beta = beta_original;
      // p = r + beta * p ;
      KokkosBlas::axpby(1.0, r, beta, p);
    }
    else {
      beta = precond_beta;
      KokkosBlas::axpby(1.0, z, beta, p);
    }

#ifdef KK_TICTOCPRINT
    std::cout << "\tbeta_original:" << beta_original <<  std::endl;
    if (use_sgs)
    std::cout << "\tprecond_beta:" << precond_beta <<  std::endl;

#endif

    norm_res = sqrt( old_rdot = r_dot );
    precond_old_rdot = precond_r_dot;

#ifdef KK_TICTOCPRINT
    std::cout << "\tnorm_res:" << norm_res << " old_rdot:" << old_rdot<<  std::endl;
#endif
    ++iteration ;
  }

  Space().fence();
  iter_time = wall_clock.seconds();

  if ( 0 != result ) {
    result->iteration   = iteration ;
    result->iter_time   = iter_time ;
    result->matvec_time = matvec_time ;
    result->norm_res    = norm_res ;
    result->precond_time = precond_time;
    result->precond_init_time = precond_init_time;
  }

  if (use_sgs & owner_handle ){

    kh.destroy_gs_handle();
  }
}

template< typename KernelHandle_t,
          typename crsMatrix_t,
          typename y_vector_t,
          typename x_vector_t
          >
void pcgsolve(
               KernelHandle_t &kh
            ,  const crsMatrix_t &crsMat
            ,  const y_vector_t &y_vector
            ,  x_vector_t x_vector
            ,  const size_t  maximum_iteration = 200
            ,  const double tolerance = std::numeric_limits<double>::epsilon()
            ,  CGSolveResult * result = 0
            ,  bool use_sgs = true
            ,  int /*clusterSize*/ = 1
            ,  bool use_sequential_sgs = false)
{
  using namespace KokkosSparse;
  using namespace KokkosSparse::Experimental;
  using size_type = typename KernelHandle_t::size_type;
  using nnz_lno_t = typename KernelHandle_t::nnz_lno_t;
  using Space = typename KernelHandle_t::HandleExecSpace;
  static_assert(std::is_same<double, typename KernelHandle_t::nnz_scalar_t>::value,
      "The PCG performance test only works with scalar = double.");

  const nnz_lno_t count_total = crsMat.numRows();

  size_t  iteration = 0 ;
  double  iter_time = 0 ;
  double  matvec_time = 0 ;
  double  norm_res = 0 ;
  double precond_time = 0;
  double precond_init_time = 0;

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

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

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

  /* r = b - A * x ; */
  /* p  = x       */  Kokkos::deep_copy( p , x_vector );

  /* Ap = A * p   */  KokkosSparse::spmv("N", 1, crsMat, pAll, 0, Ap);

  /* r  = Ap       */  Kokkos::deep_copy( r , Ap );

  /* r = b - r   */  KokkosBlas::axpby(1.0, y_vector, -1.0, r);

  /* p  = r       */  Kokkos::deep_copy( p , r );

  double old_rdot = KokkosBlas::dot( r , r );
  norm_res  = sqrt( old_rdot );

  int apply_count = 1;
  y_vector_t z;

  double precond_old_rdot = 1;
  //Kokkos::deep_copy( p , z );

  bool use_par_sgs = use_sgs && !use_sequential_sgs;

  auto ptrHost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), crsMat.graph.row_map);
  auto indHost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), crsMat.graph.entries);
  auto valHost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), crsMat.values);
  Kokkos::View<double*, Kokkos::HostSpace> diagHost;
  if(use_sequential_sgs)
  {
    diagHost = Kokkos::View<double*, Kokkos::HostSpace>("Diag for Seq SOR", count_total);
    for(int i = 0; i < count_total; i++)
    {
      for(size_type j = ptrHost(i); j < ptrHost(i + 1); j++)
      {
        if(indHost(j) == i)
          diagHost(i) = 1.0 / valHost(j);
      }
    }
  }
  auto xHost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), x_vector);
  auto yHost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), y_vector);
  
  if(use_sgs) {
    timer.reset();
    z = y_vector_t( "pcg::z" , count_total );
    if (use_par_sgs) {
      gauss_seidel_numeric
        (&kh, count_total, count_total, crsMat.graph.row_map, crsMat.graph.entries, crsMat.values);

      Space().fence();

      precond_init_time += timer.seconds();
      Space().fence();
      timer.reset();

      symmetric_gauss_seidel_apply
          (&kh, count_total, count_total, crsMat.graph.row_map, crsMat.graph.entries, crsMat.values, z, r, true, true, 1.0, apply_count);

      Space().fence();
    }
    else if(use_sequential_sgs) {
      //z = LHS (aka x), r RHS (aka y or b)
      Kokkos::deep_copy(z, 0.0);
      auto zhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), z);
      auto rhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), r);
      //as with par_sgs, init unknown to 0
      timer.reset();
      for(int sweep = 0; sweep < apply_count; sweep++)
      {
        KokkosSparse::Impl::Sequential::gaussSeidel<nnz_lno_t, size_type, double, double, double>
          (count_total, // rows = cols of the matrix
           1,           // number of vectors in X and B
           ptrHost.data(), indHost.data(), valHost.data(),
           rhost.data(), count_total, //raw ptr to B vector, and B column stride (for when multiple RHS gets added to MTSGS)
           zhost.data(), count_total, //raw ptr to X vector, and X column stride
           diagHost.data(),
           1.0,
           "F");
        KokkosSparse::Impl::Sequential::gaussSeidel<nnz_lno_t, size_type, double, double, double>
          (count_total, 1,
           ptrHost.data(), indHost.data(), valHost.data(), 
           rhost.data(), count_total,
           zhost.data(), count_total,
           diagHost.data(),
           1.0,
           "B");
      }
      //result is in z (but r doesn't change)
      Kokkos::deep_copy(z, zhost);
      Kokkos::deep_copy(r, rhost);
    }
    precond_time += timer.seconds();
    precond_old_rdot = KokkosBlas::dot(r , z);
    Kokkos::deep_copy(p, z);
  }

  iteration = 0 ;

#ifdef KK_TICTOCPRINT

  std::cout << "norm_res:" << norm_res << " old_rdot:" << old_rdot<<  std::endl;

#endif
  while (tolerance < norm_res && iteration < maximum_iteration ) {
    std::cout << "Running CG iteration " << iteration << ", current resnorm = " << norm_res << '\n';


    timer.reset();
    /* Ap = A * p   */  KokkosSparse::spmv("N", 1, crsMat, pAll, 0, Ap);


    Space().fence();
    matvec_time += timer.seconds();

    //const double pAp_dot = Kokkos::Example::all_reduce( dot( count_owned , p , Ap ) , import.comm );
    //const double pAp_dot = dot<y_vector_t,y_vector_t, Space>( count_total , p , Ap ) ;

    /* pAp_dot = dot(Ap , p ) */ const double pAp_dot = KokkosBlas::dot( p , Ap ) ;

    double alpha  = 0;
    if (use_sgs){
      alpha = precond_old_rdot / pAp_dot ;
    }
    else {
      alpha = old_rdot / pAp_dot ;
    }

    /* x +=  alpha * p ;  */  KokkosBlas::axpby(alpha, p, 1.0, x_vector);

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

    const double r_dot = KokkosBlas::dot( r , r );

    const double beta_original  = r_dot / old_rdot ;
    double precond_r_dot = 1;
    double precond_beta = 1;
    if(use_sgs)
    {
      Space().fence();
      timer.reset();
      if (use_par_sgs)
      {
        symmetric_gauss_seidel_apply(
            &kh,
            count_total, count_total,
            crsMat.graph.row_map,
            crsMat.graph.entries,
            crsMat.values, z, r, true, true,
            1.0, apply_count);
      }
      else if(use_sequential_sgs)
      {
        //z = LHS (aka x), r RHS (aka y or b)
        Kokkos::deep_copy(z, 0.0);
        auto zhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), z);
        auto rhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), r);
        //as with the par_sgs version, init unknown (here, zhost) to 0
        for(int sweep = 0; sweep < apply_count; sweep++)
        {
          KokkosSparse::Impl::Sequential::gaussSeidel<nnz_lno_t, size_type, double, double, double>
            (count_total, 1,
             ptrHost.data(), indHost.data(), valHost.data(), 
             rhost.data(), count_total,
             zhost.data(), count_total,
             diagHost.data(),
             1.0,
             "F");
          KokkosSparse::Impl::Sequential::gaussSeidel<nnz_lno_t , size_type, double, double, double>
            (count_total, 1,
             ptrHost.data(), indHost.data(), valHost.data(), 
             rhost.data(), count_total,
             zhost.data(), count_total,
             diagHost.data(),
             1.0,
             "B");
        }
        Kokkos::deep_copy(z, zhost);
        Kokkos::deep_copy(r, rhost);
      }
      precond_time += timer.seconds();
      precond_r_dot = KokkosBlas::dot(r , z );
      precond_beta  = precond_r_dot / precond_old_rdot ;
    }
    double beta = 1;
    if (!use_sgs) {
      beta = beta_original;
      /* p = r + beta * p ; */  KokkosBlas::axpby(1.0, r, beta, p);
    }
    else {
      beta = precond_beta;
      KokkosBlas::axpby(1.0, z, beta, p);
    }

#ifdef KK_TICTOCPRINT
    std::cout << "\tbeta_original:" << beta_original <<  std::endl;
    if (use_sgs)
    std::cout << "\tprecond_beta:" << precond_beta <<  std::endl;

#endif

    norm_res = sqrt( old_rdot = r_dot );
    precond_old_rdot = precond_r_dot;

#ifdef KK_TICTOCPRINT
    std::cout << "\tnorm_res:" << norm_res << " old_rdot:" << old_rdot<<  std::endl;
#endif
    ++iteration ;
  }

  Space().fence();
  iter_time = wall_clock.seconds();

  if ( 0 != result ) {
    result->iteration   = iteration ;
    result->iter_time   = iter_time ;
    result->matvec_time = matvec_time ;
    result->norm_res    = norm_res ;
    result->precond_time = precond_time;
    result->precond_init_time = precond_init_time;
  }
}

} // namespace Example
} // namespace Kokkos
}
//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

#endif /* #ifndef KOKKOS_EXAMPLE_CG_SOLVE */