CRSMatrixView.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2021, Lawrence Livermore National Security, LLC and LvArray contributors.
 * All rights reserved.
 * See the LICENSE file for details.
 * SPDX-License-Identifier: (BSD-3-Clause)
 */

#pragma once

#include "SparsityPatternView.hpp"
#include "arrayManipulation.hpp"
#include "ArraySlice.hpp"
#include "umpireInterface.hpp"

namespace LvArray
{

namespace internal
{

template< typename POLICY, typename T >
LVARRAY_HOST_DEVICE inline
void atomicAdd( POLICY, T * const acc, T const & val )
{ RAJA::atomicAdd< POLICY >( acc, val ); }

DISABLE_HD_WARNING
template< typename T >
LVARRAY_HOST_DEVICE constexpr inline
void atomicAdd( RAJA::seq_atomic, T * acc, T const & val )
{ *acc += val; }

} // namespace internal

template< typename T,
          typename COL_TYPE,
          typename INDEX_TYPE,
          template< typename > class BUFFER_TYPE
          >
class CRSMatrixView : protected SparsityPatternView< COL_TYPE, INDEX_TYPE, BUFFER_TYPE >
{

  using ParentClass = SparsityPatternView< COL_TYPE, INDEX_TYPE, BUFFER_TYPE >;

  using typename ParentClass::INDEX_TYPE_NC;

  using typename ParentClass::SIZE_TYPE;

public:
  static_assert( !std::is_const< T >::value ||
                 (std::is_const< COL_TYPE >::value && std::is_const< INDEX_TYPE >::value),
                 "When T is const COL_TYPE and INDEX_TYPE must also be const." );

  using EntryType = T;
  using typename ParentClass::ColType;
  using typename ParentClass::IndexType;


  CRSMatrixView() = default;

  CRSMatrixView( CRSMatrixView const & ) = default;

  inline
  CRSMatrixView( CRSMatrixView && ) = default;

  LVARRAY_HOST_DEVICE constexpr inline
  CRSMatrixView( INDEX_TYPE const nRows,
                 INDEX_TYPE const nCols,
                 BUFFER_TYPE< INDEX_TYPE > const & offsets,
                 BUFFER_TYPE< SIZE_TYPE > const & nnz,
                 BUFFER_TYPE< COL_TYPE > const & columns,
                 BUFFER_TYPE< T > const & entries ):
    ParentClass( nRows, nCols, offsets, nnz, columns ),
    m_entries( entries )
  {}

  inline
  CRSMatrixView & operator=( CRSMatrixView const & ) = default;

  inline
  CRSMatrixView & operator=( CRSMatrixView && ) = default;


  LVARRAY_HOST_DEVICE constexpr inline
  CRSMatrixView< T, COL_TYPE, INDEX_TYPE const, BUFFER_TYPE >
  toView() const
  {
    return CRSMatrixView< T, COL_TYPE, INDEX_TYPE const, BUFFER_TYPE >( numRows(),
                                                                        numColumns(),
                                                                        this->m_offsets,
                                                                        this->m_sizes,
                                                                        this->m_values,
                                                                        this->m_entries );
  }

  LVARRAY_HOST_DEVICE constexpr inline
  CRSMatrixView< T, COL_TYPE const, INDEX_TYPE const, BUFFER_TYPE >
  toViewConstSizes() const
  {
    return CRSMatrixView< T, COL_TYPE const, INDEX_TYPE const, BUFFER_TYPE >( numRows(),
                                                                              numColumns(),
                                                                              this->m_offsets,
                                                                              this->m_sizes,
                                                                              this->m_values,
                                                                              this->m_entries );
  }

  LVARRAY_HOST_DEVICE constexpr inline
  CRSMatrixView< T const, COL_TYPE const, INDEX_TYPE const, BUFFER_TYPE >
  toViewConst() const
  {
    return CRSMatrixView< T const, COL_TYPE const, INDEX_TYPE const, BUFFER_TYPE >( numRows(),
                                                                                    numColumns(),
                                                                                    this->m_offsets,
                                                                                    this->m_sizes,
                                                                                    this->m_values,
                                                                                    this->m_entries );
  }

  LVARRAY_HOST_DEVICE constexpr inline
  SparsityPatternView< COL_TYPE const, INDEX_TYPE const, BUFFER_TYPE >
  toSparsityPatternView() const &
  {
    return SparsityPatternView< COL_TYPE const, INDEX_TYPE const, BUFFER_TYPE >( numRows(),
                                                                                 numColumns(),
                                                                                 this->m_offsets,
                                                                                 this->m_sizes,
                                                                                 this->m_values );
  }



  using ParentClass::numRows;
  using ParentClass::numColumns;
  using ParentClass::numNonZeros;
  using ParentClass::nonZeroCapacity;
  using ParentClass::empty;



  LVARRAY_HOST_DEVICE inline
  ArraySlice< T, 1, 0, INDEX_TYPE_NC > getEntries( INDEX_TYPE const row ) const
  {
    ARRAYOFARRAYS_CHECK_BOUNDS( row );
    return ArraySlice< T, 1, 0, INDEX_TYPE_NC >( m_entries.data() + this->m_offsets[ row ],
                                                 &this->m_sizes[ row ],
                                                 nullptr );
  }

  LVARRAY_HOST_DEVICE inline
  T const * getEntries() const
  {
    return m_entries.data();
  }

  using ParentClass::getColumns;
  using ParentClass::getOffsets;
  using ParentClass::getSizes;



  LVARRAY_HOST_DEVICE inline
  bool insertNonZero( INDEX_TYPE const row, COL_TYPE const col, T const & entry ) const
  { return ParentClass::insertIntoSetImpl( row, col, CallBacks( *this, row, &entry ) ); }

  LVARRAY_HOST_DEVICE inline
  INDEX_TYPE_NC insertNonZeros( INDEX_TYPE const row,
                                COL_TYPE const * const LVARRAY_RESTRICT cols,
                                T const * const LVARRAY_RESTRICT entriesToInsert,
                                INDEX_TYPE const ncols ) const
  { return ParentClass::insertIntoSetImpl( row, cols, cols + ncols, CallBacks( *this, row, entriesToInsert ) ); }

  LVARRAY_HOST_DEVICE inline
  bool removeNonZero( INDEX_TYPE const row, COL_TYPE const col ) const
  { return ParentClass::removeFromSetImpl( row, col, CallBacks( *this, row, nullptr )); }

  DISABLE_HD_WARNING
  LVARRAY_HOST_DEVICE inline
  INDEX_TYPE_NC removeNonZeros( INDEX_TYPE const row,
                                COL_TYPE const * const LVARRAY_RESTRICT cols,
                                INDEX_TYPE const ncols ) const
  {
    T * const entries = getEntries( row );
    INDEX_TYPE const rowNNZ = numNonZeros( row );
    INDEX_TYPE const nRemoved = ParentClass::removeFromSetImpl( row, cols, cols + ncols, CallBacks( *this, row, nullptr ) );

    arrayManipulation::destroy( entries + rowNNZ - nRemoved, nRemoved );

    return nRemoved;
  }



  template< typename POLICY >
  inline void setValues( T const & value ) const
  {
    CRSMatrixView< T, COL_TYPE const, INDEX_TYPE const, BUFFER_TYPE > const view = toViewConstSizes();
    RAJA::forall< POLICY >( RAJA::TypedRangeSegment< INDEX_TYPE >( 0, numRows() ),
                            [view, value] LVARRAY_HOST_DEVICE ( INDEX_TYPE const row )
      {
        INDEX_TYPE const nnz = view.numNonZeros( row );
        ArraySlice< T, 1, 0, INDEX_TYPE_NC > const entries = view.getEntries( row );

        for( INDEX_TYPE_NC i = 0; i < nnz; ++i )
        { entries[ i ] = value; }
      } );
  }

  inline void zero() const
  {
  #if !defined( LVARRAY_USE_UMPIRE )
    LVARRAY_ERROR_IF_NE_MSG( m_entries.getPreviousSpace(), MemorySpace::host, "Without Umpire only host memory is supported." );
  #endif

    if( m_entries.capacity() > 0 )
    {
      ParentClass::move( m_entries.getPreviousSpace(), false );
      m_entries.move( m_entries.getPreviousSpace(), true );

      size_t const numBytes = m_entries.capacity() * sizeof( T );

      umpireInterface::memset( m_entries.data(), 0, numBytes );

    }
  }

  template< typename AtomicPolicy >
  LVARRAY_HOST_DEVICE inline
  void addToRow( INDEX_TYPE const row,
                 COL_TYPE const * const LVARRAY_RESTRICT cols,
                 T const * const LVARRAY_RESTRICT vals,
                 INDEX_TYPE const nCols ) const
  {
    INDEX_TYPE const nnz = numNonZeros( row );
    if( nCols < nnz / 4 && nnz > 64 )
    {
      addToRowBinarySearch< AtomicPolicy >( row, cols, vals, nCols );
    }
    else
    {
      addToRowLinearSearch< AtomicPolicy >( row, cols, vals, nCols );
    }
  }

  template< typename AtomicPolicy >
  LVARRAY_HOST_DEVICE inline
  void addToRowBinarySearch( INDEX_TYPE const row,
                             COL_TYPE const * const LVARRAY_RESTRICT cols,
                             T const * const LVARRAY_RESTRICT vals,
                             INDEX_TYPE const nCols ) const
  {
    LVARRAY_ASSERT( sortedArrayManipulation::isSortedUnique( cols, cols + nCols ) );

    INDEX_TYPE const nnz = numNonZeros( row );
    COL_TYPE const * const columns = getColumns( row );
    T * const entries = getEntries( row );

    INDEX_TYPE_NC curPos = 0;
    for( INDEX_TYPE_NC i = 0; i < nCols; ++i )
    {
      INDEX_TYPE const pos = sortedArrayManipulation::find( columns + curPos, nnz - curPos, cols[ i ] ) + curPos;
      LVARRAY_ASSERT_GT( nnz, pos );
      LVARRAY_ASSERT_EQ( columns[ pos ], cols[ i ] );

      internal::atomicAdd( AtomicPolicy{}, entries + pos, vals[ i ] );
      curPos = pos + 1;
    }
  }

  template< typename AtomicPolicy >
  LVARRAY_HOST_DEVICE inline
  void addToRowLinearSearch( INDEX_TYPE const row,
                             COL_TYPE const * const LVARRAY_RESTRICT cols,
                             T const * const LVARRAY_RESTRICT vals,
                             INDEX_TYPE const nCols ) const
  {
    LVARRAY_ASSERT( sortedArrayManipulation::isSortedUnique( cols, cols + nCols ) );

    INDEX_TYPE const nnz = numNonZeros( row );
    COL_TYPE const * const columns = getColumns( row );
    T * const entries = getEntries( row );

    INDEX_TYPE_NC curPos = 0;
    for( INDEX_TYPE_NC i = 0; i < nCols; ++i )
    {
      for( INDEX_TYPE_NC j = curPos; j < nnz; ++j )
      {
        if( columns[ j ] == cols[ i ] )
        {
          curPos = j;
          break;
        }
      }
      LVARRAY_ASSERT_EQ( columns[ curPos ], cols[ i ] );
      internal::atomicAdd( AtomicPolicy{}, entries + curPos, vals[ i ] );
      ++curPos;
    }
  }

  template< typename AtomicPolicy >
  LVARRAY_HOST_DEVICE inline
  void addToRowBinarySearchUnsorted( INDEX_TYPE const row,
                                     COL_TYPE const * const LVARRAY_RESTRICT cols,
                                     T const * const LVARRAY_RESTRICT vals,
                                     INDEX_TYPE const nCols ) const
  {
    INDEX_TYPE const nnz = numNonZeros( row );
    COL_TYPE const * const columns = getColumns( row );
    T * const entries = getEntries( row );

    for( INDEX_TYPE_NC i = 0; i < nCols; ++i )
    {
      INDEX_TYPE const pos = sortedArrayManipulation::find( columns, nnz, cols[ i ] );
      LVARRAY_ASSERT_GT( nnz, pos );
      LVARRAY_ASSERT_EQ( columns[ pos ], cols[ i ] );

      internal::atomicAdd( AtomicPolicy{}, entries + pos, vals[ i ] );
    }
  }



  void move( MemorySpace const space, bool const touch=true ) const
  {
    ParentClass::move( space, touch );
    m_entries.move( space, touch );
  }


protected:

  CRSMatrixView( bool ):
    ParentClass( true ),
    m_entries( true )
  {}

  template< typename U >
  void setName( std::string const & name )
  {
    ParentClass::template setName< U >( name );
    m_entries.template setName< U >( name + "/entries" );
  }

  BUFFER_TYPE< T > m_entries;

private:

  class CallBacks
  {
public:

    LVARRAY_HOST_DEVICE inline
    CallBacks( CRSMatrixView const & crsMV,
               INDEX_TYPE const row, T const * const entriesToInsert ):
      m_crsMV( crsMV ),
      m_row( row ),
      m_rowNNZ( crsMV.numNonZeros( row ) ),
      m_rowCapacity( crsMV.nonZeroCapacity( row ) ),
      m_entries( crsMV.getEntries( row ) ),
      m_entriesToInsert( entriesToInsert )
    {}

    LVARRAY_HOST_DEVICE inline
    COL_TYPE * incrementSize( COL_TYPE * const curPtr, INDEX_TYPE const nToAdd ) const
    {
      LVARRAY_UNUSED_VARIABLE( curPtr );
#ifdef LVARRAY_BOUNDS_CHECK
      LVARRAY_ERROR_IF_GT_MSG( m_rowNNZ + nToAdd, m_rowCapacity, "CRSMatrixView cannot do reallocation." );
#else
      LVARRAY_DEBUG_VAR( nToAdd );
#endif
      return m_crsMV.getSetValues( m_row );
    }

    LVARRAY_HOST_DEVICE inline
    void insert( INDEX_TYPE const insertPos ) const
    { arrayManipulation::emplace( m_entries, m_rowNNZ, insertPos, m_entriesToInsert[0] ); }

    DISABLE_HD_WARNING
    LVARRAY_HOST_DEVICE inline
    void set( INDEX_TYPE const pos, INDEX_TYPE const colPos ) const
    { new (&m_entries[ pos ]) T( m_entriesToInsert[ colPos ] ); }

    DISABLE_HD_WARNING
    LVARRAY_HOST_DEVICE inline
    void insert( INDEX_TYPE const nLeftToInsert,
                 INDEX_TYPE const colPos,
                 INDEX_TYPE const pos,
                 INDEX_TYPE const prevPos ) const
    {
      arrayManipulation::shiftUp( m_entries, prevPos, pos, nLeftToInsert );
      new (&m_entries[pos + nLeftToInsert - 1]) T( m_entriesToInsert[colPos] );
    }

    LVARRAY_HOST_DEVICE inline
    void remove( INDEX_TYPE_NC removePos ) const
    { arrayManipulation::erase( m_entries, m_rowNNZ, removePos ); }

    LVARRAY_HOST_DEVICE inline
    void remove( INDEX_TYPE const nRemoved,
                 INDEX_TYPE const curPos,
                 INDEX_TYPE const nextPos ) const
    { arrayManipulation::shiftDown( m_entries, nextPos, curPos + 1, nRemoved ); }

private:
    CRSMatrixView const & m_crsMV;

    INDEX_TYPE const m_row;

    INDEX_TYPE const m_rowNNZ;

    INDEX_TYPE const m_rowCapacity;

    T * const m_entries;

    T const * const m_entriesToInsert;
  };

};

} /* namespace LvArray */