HypreKernels.hpp

Go to the documentation of this file.

/*
 * ------------------------------------------------------------------------------------------------------------
 * SPDX-License-Identifier: LGPL-2.1-only
 *
 * Copyright (c) 2016-2024 Lawrence Livermore National Security LLC
 * Copyright (c) 2018-2024 TotalEnergies
 * Copyright (c) 2018-2024 The Board of Trustees of the Leland Stanford Junior University
 * Copyright (c) 2023-2024 Chevron
 * Copyright (c) 2019-     GEOS/GEOSX Contributors
 * All rights reserved
 *
 * See top level LICENSE, COPYRIGHT, CONTRIBUTORS, NOTICE, and ACKNOWLEDGEMENTS files for details.
 * ------------------------------------------------------------------------------------------------------------
 */
 
#ifndef GEOS_LINEARALGEBRA_INTERFACES_HYPREKERNELS_HPP_
#define GEOS_LINEARALGEBRA_INTERFACES_HYPREKERNELS_HPP_
 
#include "codingUtilities/Utilities.hpp"
#include "codingUtilities/traits.hpp"
#include "common/DataTypes.hpp"
#include "common/GEOS_RAJA_Interface.hpp"
#include "linearAlgebra/common/common.hpp"
#include "linearAlgebra/interfaces/hypre/HypreUtils.hpp"
 
#include <_hypre_parcsr_mv.h>
 
namespace geos
{
namespace hypre
{
 
 
template< bool CONST >
struct CSRData
{
  HYPRE_Int const * rowptr;
  HYPRE_Int const * colind;
  add_const_if_t< HYPRE_Real, CONST > * values;
  HYPRE_Int nrow;
  HYPRE_Int ncol;
  HYPRE_Int nnz;
 
  explicit CSRData( add_const_if_t< hypre_CSRMatrix, CONST > * const mat )
    : rowptr( hypre_CSRMatrixI( mat ) ),
    colind( hypre_CSRMatrixJ( mat ) ),
    values( hypre_CSRMatrixData( mat ) ),
    nrow( hypre_CSRMatrixNumRows( mat ) ),
    ncol( hypre_CSRMatrixNumCols( mat ) ),
    nnz( hypre_CSRMatrixNumNonzeros( mat ) )
  {}
};
 
HYPRE_BigInt const * getOffdColumnMap( hypre_ParCSRMatrix const * const mat );
 
void scaleMatrixValues( hypre_CSRMatrix * const mat,
                        real64 const factor );
 
void scaleMatrixRows( hypre_CSRMatrix * const mat,
                      hypre_Vector const * const vec );
 
void clampMatrixEntries( hypre_CSRMatrix * const mat,
                         real64 const lo,
                         real64 const hi,
                         bool const skip_diag );
 
real64 computeMaxNorm( hypre_CSRMatrix const * const mat );
 
real64 computeMaxNorm( hypre_CSRMatrix const * const mat,
                       arrayView1d< globalIndex const > const & rowIndices,
                       globalIndex const firstLocalRow );
 
template< typename F, typename R >
void rescaleMatrixRows( hypre_ParCSRMatrix * const mat,
                        arrayView1d< globalIndex const > const & rowIndices,
                        F transform,
                        R reduce )
{
  CSRData< false > diag{ hypre_ParCSRMatrixDiag( mat ) };
  CSRData< false > offd{ hypre_ParCSRMatrixOffd( mat ) };
  HYPRE_BigInt const firstLocalRow = hypre_ParCSRMatrixFirstRowIndex( mat );
 
  forAll< execPolicy >( rowIndices.size(), [diag, offd, transform, reduce, rowIndices, firstLocalRow] GEOS_HYPRE_DEVICE ( localIndex const i )
  {
    HYPRE_Int const localRow = LvArray::integerConversion< HYPRE_Int >( rowIndices[i] - firstLocalRow );
    GEOS_ASSERT( 0 <= localRow && localRow < diag.nrow );
 
    HYPRE_Real scale = 0.0;
    for( HYPRE_Int k = diag.rowptr[localRow]; k < diag.rowptr[localRow + 1]; ++k )
    {
      scale = reduce( scale, transform( diag.values[k] ) );
    }
    if( offd.ncol > 0 )
    {
      for( HYPRE_Int k = offd.rowptr[localRow]; k < offd.rowptr[localRow + 1]; ++k )
      {
        scale = reduce( scale, transform( offd.values[k] ) );
      }
    }
 
    GEOS_ASSERT_MSG( !isZero( scale ), "Zero row sum in row " << rowIndices[i] );
    scale = 1.0 / scale;
    for( HYPRE_Int k = diag.rowptr[localRow]; k < diag.rowptr[localRow + 1]; ++k )
    {
      diag.values[k] *= scale;
    }
    if( offd.ncol > 0 )
    {
      for( HYPRE_Int k = offd.rowptr[localRow]; k < offd.rowptr[localRow + 1]; ++k )
      {
        offd.values[k] *= scale;
      }
    }
  } );
}
 
template< typename F, typename R >
void computeRowsSums( hypre_ParCSRMatrix const * const mat,
                      hypre_ParVector * const vec,
                      F transform,
                      R reduce )
{
  CSRData< true > const diag{ hypre_ParCSRMatrixDiag( mat ) };
  CSRData< true > const offd{ hypre_ParCSRMatrixOffd( mat ) };
  HYPRE_Real * const values = hypre_VectorData( hypre_ParVectorLocalVector( vec ) );
 
  forAll< execPolicy >( diag.nrow, [diag, offd, transform, reduce, values] GEOS_HYPRE_DEVICE ( HYPRE_Int const localRow )
  {
    HYPRE_Real sum = 0.0;
    for( HYPRE_Int k = diag.rowptr[localRow]; k < diag.rowptr[localRow + 1]; ++k )
    {
      sum = reduce( sum, transform( diag.values[k] ) );
    }
    if( offd.ncol )
    {
      for( HYPRE_Int k = offd.rowptr[localRow]; k < offd.rowptr[localRow + 1]; ++k )
      {
        sum = reduce( sum, transform( offd.values[k] ) );
      }
    }
    values[localRow] = sum;
  } );
}
 
namespace internal
{
 
template< typename MAP >
void GEOS_HYPRE_DEVICE
makeSortedPermutation( HYPRE_Int const * const indices,
                       HYPRE_Int const size,
                       HYPRE_Int * const perm,
                       MAP map )
{
  for( HYPRE_Int i = 0; i < size; ++i )
  {
    perm[i] = i; // std::iota
  }
  auto const comp = [indices, map] GEOS_HYPRE_DEVICE ( HYPRE_Int i, HYPRE_Int j )
  {
    return map( indices[i] ) < map( indices[j] );
  };
  LvArray::sortedArrayManipulation::makeSorted( perm, perm + size, comp );
}
 
} // namespace internal
 
template< typename KERNEL >
void addMatrixEntries( hypre_ParCSRMatrix const * const src,
                       hypre_ParCSRMatrix * const dst,
                       real64 const scale )
{
  GEOS_LAI_ASSERT( src != nullptr );
  GEOS_LAI_ASSERT( dst != nullptr );
  KERNEL::launch( hypre_ParCSRMatrixDiag( src ),
                  [] GEOS_HYPRE_DEVICE ( auto x ){ return x; },
                  hypre_ParCSRMatrixDiag( dst ),
                  [] GEOS_HYPRE_DEVICE ( auto x ){ return x; },
                  scale );
  if( hypre_CSRMatrixNumCols( hypre_ParCSRMatrixOffd( dst ) ) > 0 )
  {
    HYPRE_BigInt const * const src_colmap = hypre::getOffdColumnMap( src );
    HYPRE_BigInt const * const dst_colmap = hypre::getOffdColumnMap( dst );
    KERNEL::launch( hypre_ParCSRMatrixOffd( src ),
                    [src_colmap] GEOS_HYPRE_DEVICE ( auto i ){ return src_colmap[i]; },
                    hypre_ParCSRMatrixOffd( dst ),
                    [dst_colmap] GEOS_HYPRE_DEVICE ( auto i ){ return dst_colmap[i]; },
                    scale );
  }
}
 
struct AddEntriesRestrictedKernel
{
  template< typename SRC_COLMAP, typename DST_COLMAP >
  static void
  launch( hypre_CSRMatrix const * const src_mat,
          SRC_COLMAP const src_colmap,
          hypre_CSRMatrix * const dst_mat,
          DST_COLMAP const dst_colmap,
          real64 const scale )
  {
    GEOS_LAI_ASSERT( src_mat != nullptr );
    GEOS_LAI_ASSERT( dst_mat != nullptr );
 
    CSRData< true > src{ src_mat };
    CSRData< false > dst{ dst_mat };
    GEOS_LAI_ASSERT_EQ( src.nrow, dst.nrow );
 
    if( src.ncol == 0 || isZero( scale ) )
    {
      return;
    }
 
    // Allocate contiguous memory to store sorted column permutations of each row
    array1d< HYPRE_Int > const src_permutation_arr( src.nnz );
    array1d< HYPRE_Int > const dst_permutation_arr( dst.nnz );
 
    arrayView1d< HYPRE_Int > const src_permutation = src_permutation_arr.toView();
    arrayView1d< HYPRE_Int > const dst_permutation = dst_permutation_arr.toView();
    // Each thread adds one row of src into dst
    forAll< hypre::execPolicy >( dst.nrow,
                                 [src,
                                  src_colmap,
                                  dst,
                                  dst_colmap,
                                  scale,
                                  src_permutation,
                                  dst_permutation ] GEOS_HYPRE_DEVICE ( HYPRE_Int const localRow )
    {
      HYPRE_Int const src_offset = src.rowptr[localRow];
      HYPRE_Int const src_length = src.rowptr[localRow + 1] - src_offset;
      HYPRE_Int const * const src_indices = src.colind + src_offset;
      HYPRE_Real const * const src_values = src.values + src_offset;
      HYPRE_Int * const src_perm = src_permutation.data() + src_offset;
 
      HYPRE_Int const dst_offset = dst.rowptr[localRow];
      HYPRE_Int const dst_length = dst.rowptr[localRow + 1] - dst_offset;
      HYPRE_Int const * const dst_indices = dst.colind + dst_offset;
      HYPRE_Real * const dst_values = dst.values + dst_offset;
      HYPRE_Int * const dst_perm = dst_permutation.data() + dst_offset;
 
      // Since hypre does not store columns in sorted order, create a sorted "view" of src and dst rows
      // TODO: it would be nice to cache the permutation arrays somewhere to avoid recomputing
      internal::makeSortedPermutation( src_indices, src_length, src_perm, src_colmap );
      internal::makeSortedPermutation( dst_indices, dst_length, dst_perm, dst_colmap );
 
      // Add entries looping through them in sorted column order, skipping src entries not in dst
      for( HYPRE_Int i = 0, j = 0; i < dst_length && j < src_length; ++i )
      {
        while( j < src_length && src_colmap( src_indices[src_perm[j]] ) < dst_colmap( dst_indices[dst_perm[i]] ) )
        {
          ++j;
        }
        if( j < src_length && src_colmap( src_indices[src_perm[j]] ) == dst_colmap( dst_indices[dst_perm[i]] ) )
        {
          dst_values[dst_perm[i]] += scale * src_values[src_perm[j++]];
        }
      }
    } );
  }
};
 
struct AddEntriesSamePatternKernel
{
  template< typename SRC_COLMAP, typename DST_COLMAP >
  static void
  launch( hypre_CSRMatrix const * const src_mat,
          SRC_COLMAP const src_colmap,
          hypre_CSRMatrix * const dst_mat,
          DST_COLMAP const dst_colmap,
          real64 const scale )
  {
    GEOS_LAI_ASSERT( src_mat != nullptr );
    GEOS_LAI_ASSERT( dst_mat != nullptr );
 
    CSRData< true > src{ src_mat };
    CSRData< false > dst{ dst_mat };
    GEOS_LAI_ASSERT_EQ( src.nrow, dst.nrow );
 
    if( src.ncol == 0 || isZero( scale ) )
    {
      return;
    }
 
    // Each thread adds one row of src into dst
    forAll< hypre::execPolicy >( dst.nrow,
                                 [src, src_colmap, dst, dst_colmap, scale] GEOS_HYPRE_DEVICE ( HYPRE_Int const localRow )
    {
      HYPRE_Int const src_offset = src.rowptr[localRow];
      HYPRE_Int const src_length = src.rowptr[localRow + 1] - src_offset;
      HYPRE_Int const * const src_indices = src.colind + src_offset;
      HYPRE_Real const * const src_values = src.values + src_offset;
 
      HYPRE_Int const dst_offset = dst.rowptr[localRow];
      HYPRE_Int const dst_length = dst.rowptr[localRow + 1] - dst_offset;
      HYPRE_Int const * const dst_indices = dst.colind + dst_offset;
      HYPRE_Real * const dst_values = dst.values + dst_offset;
 
      GEOS_ASSERT_EQ( src_offset, dst_offset );
      GEOS_ASSERT_EQ( src_length, dst_length );
      GEOS_DEBUG_VAR( src_length, dst_length, src_indices, dst_indices, src_colmap, dst_colmap );
 
      // NOTE: this assumes that entries are in the exact same order, to avoid creating a sorted view
      for( HYPRE_Int i = 0; i < dst_length; ++i )
      {
        GEOS_ASSERT_EQ( src_colmap( src_indices[i] ), dst_colmap( dst_indices[i] ) );
        dst_values[i] += scale * src_values[i];
      }
    } );
  }
};
 
 
} // namespace hypre
} // namespace geos
 
#endif //GEOS_LINEARALGEBRA_INTERFACES_HYPREKERNELS_HPP_