CSRMatrixOperationsKokkos.h

Go to the documentation of this file.

#ifndef included_CSRMatrixOperationsKokkos_H_
#define included_CSRMatrixOperationsKokkos_H_
 
#include "AMP/AMP_TPLs.h"
#include "AMP/matrices/CSRConfig.h"
#include "AMP/matrices/data/CSRMatrixData.h"
#include "AMP/matrices/data/MatrixData.h"
#include "AMP/matrices/operations/MatrixOperations.h"
#include "AMP/matrices/operations/default/CSRMatrixOperationsDefault.h"
#include "AMP/matrices/operations/kokkos/CSRLocalMatrixOperationsKokkos.h"
#include "AMP/utils/Memory.h"
#include "AMP/vectors/Vector.h"
 
#ifdef AMP_USE_DEVICE
    #include "AMP/matrices/operations/device/CSRMatrixOperationsDevice.h"
#endif
 
#include <type_traits>
 
#ifdef AMP_USE_KOKKOS
 
    #include "Kokkos_Core.hpp"
 
namespace AMP::LinearAlgebra {
 
template<typename Config,
    #ifdef AMP_USE_DEVICE
         class ExecSpace = typename std::conditional<alloc_info<Config::allocator>::mem_loc ==
                                                         AMP::Utilities::MemoryType::host,
                                                     Kokkos::DefaultHostExecutionSpace,
                                                     Kokkos::DefaultExecutionSpace>::type,
         class ViewSpace = typename std::conditional<
             alloc_info<Config::allocator>::mem_loc == AMP::Utilities::MemoryType::host,
             Kokkos::HostSpace,
             typename std::conditional<
                 alloc_info<Config::allocator>::mem_loc == AMP::Utilities::MemoryType::managed,
                 Kokkos::SharedSpace,
                 typename Kokkos::DefaultExecutionSpace::memory_space>::type>::type
    #else
         class ExecSpace = Kokkos::DefaultHostExecutionSpace,
         class ViewSpace = Kokkos::HostSpace
    #endif
         >
class CSRMatrixOperationsKokkos : public MatrixOperations
{
public:
    static_assert( std::is_same_v<typename Config::allocator_type::value_type, void> );
 
    using config_type       = Config;
    using allocator_type    = typename Config::allocator_type;
    using matrixdata_t      = CSRMatrixData<Config>;
    using localmatrixdata_t = typename matrixdata_t::localmatrixdata_t;
 
    using localops_t = CSRLocalMatrixOperationsKokkos<Config, ExecSpace, ViewSpace>;
 
    using gidx_t   = typename Config::gidx_t;
    using lidx_t   = typename Config::lidx_t;
    using scalar_t = typename Config::scalar_t;
 
    CSRMatrixOperationsKokkos()
        : d_exec_space(),
          d_localops_diag( std::make_shared<localops_t>( d_exec_space ) ),
          d_localops_offd( std::make_shared<localops_t>( d_exec_space ) )
    {
    }
 
    void mult( std::shared_ptr<const Vector> x,
               MatrixData const &A,
               std::shared_ptr<Vector> y ) override;
 
    void multTranspose( std::shared_ptr<const Vector> in,
                        MatrixData const &A,
                        std::shared_ptr<Vector> out ) override;
 
    void scale( AMP::Scalar alpha, MatrixData &A ) override;
 
    void scale( AMP::Scalar alpha, std::shared_ptr<const Vector> D, MatrixData &A ) override;
 
    void scaleInv( AMP::Scalar alpha, std::shared_ptr<const Vector> D, MatrixData &A ) override;
 
    void matMatMult( std::shared_ptr<MatrixData> A,
                     std::shared_ptr<MatrixData> B,
                     std::shared_ptr<MatrixData> C ) override;
 
    void axpy( AMP::Scalar alpha, const MatrixData &X, MatrixData &Y ) override;
 
    void setScalar( AMP::Scalar alpha, MatrixData &A ) override;
 
    void zero( MatrixData &A ) override;
 
    void setDiagonal( std::shared_ptr<const Vector> in, MatrixData &A ) override;
 
    void setIdentity( MatrixData &A ) override;
 
    void extractDiagonal( MatrixData const &A, std::shared_ptr<Vector> buf ) override;
 
    void getRowSums( MatrixData const &A, std::shared_ptr<Vector> buf ) override;
 
    void getRowSumsAbsolute( MatrixData const &A,
                             std::shared_ptr<Vector> buf,
                             const bool remove_zeros = false ) override;
 
    AMP::Scalar LinfNorm( const MatrixData &X ) const override;
 
    void copy( const MatrixData &X, MatrixData &Y ) override;
 
    void copyCast( const MatrixData &X, MatrixData &Y ) override;
 
    template<typename ConfigIn>
    static void copyCast( CSRMatrixData<ConfigIn> *X, CSRMatrixData<Config> *Y );
 
    std::string type() const override { return "CSRMatrixOperationsKokkos"; }
 
    void writeRestart( int64_t fid ) const override;
 
    CSRMatrixOperationsKokkos( int64_t, AMP::IO::RestartManager * )
        : d_exec_space(),
          d_localops_diag( std::make_shared<localops_t>( d_exec_space ) ),
          d_localops_offd( std::make_shared<localops_t>( d_exec_space ) )
    {
    }
 
protected:
    ExecSpace d_exec_space;
    std::shared_ptr<localops_t> d_localops_diag;
    std::shared_ptr<localops_t> d_localops_offd;
 
    // This currently forwards SpGEMM operations to either default ops or device ops
    // so internal versions of each are held
    CSRMatrixOperationsDefault<Config> d_matrixOpsDefault;
    #ifdef AMP_USE_DEVICE
    CSRMatrixOperationsDevice<Config> d_matrixOpsDevice;
    #endif
};
 
} // namespace AMP::LinearAlgebra
 
#endif
#endif