CSRLocalMatrixOperationsKokkos.h

Go to the documentation of this file.

#ifndef included_CSRLocalMatrixOperationsKokkos_H_
#define included_CSRLocalMatrixOperationsKokkos_H_
 
#include "AMP/matrices/operations/MatrixOperations.h"
 
#include <type_traits>
 
#ifdef AMP_USE_KOKKOS
 
    #include "Kokkos_Core.hpp"
 
namespace AMP::LinearAlgebra {
 
template<typename Config, class ExecSpace, class ViewSpace>
class CSRLocalMatrixOperationsKokkos
{
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 localmatrixdata_t = CSRLocalMatrixData<Config>;
 
    using gidx_t   = typename Config::gidx_t;
    using lidx_t   = typename Config::lidx_t;
    using scalar_t = typename Config::scalar_t;
 
    CSRLocalMatrixOperationsKokkos( const ExecSpace &exec_space ) : d_exec_space( exec_space ) {}
 
    void mult( const scalar_t *in,
               const scalar_t alpha,
               std::shared_ptr<localmatrixdata_t> A,
               const scalar_t beta,
               scalar_t *out );
 
    void multTranspose( const scalar_t *in, std::shared_ptr<localmatrixdata_t> A, scalar_t *out );
 
    void scale( scalar_t alpha, std::shared_ptr<localmatrixdata_t> A );
 
    void scale( scalar_t alpha, const scalar_t *D, std::shared_ptr<localmatrixdata_t> A );
 
    void scaleInv( scalar_t alpha, const scalar_t *D, std::shared_ptr<localmatrixdata_t> A );
 
    void matMatMult( std::shared_ptr<localmatrixdata_t>,
                     std::shared_ptr<localmatrixdata_t>,
                     std::shared_ptr<localmatrixdata_t> );
 
    void axpy( scalar_t alpha,
               std::shared_ptr<localmatrixdata_t> X,
               std::shared_ptr<localmatrixdata_t> Y );
 
    void setScalar( scalar_t alpha, std::shared_ptr<localmatrixdata_t> A );
 
    void zero( std::shared_ptr<localmatrixdata_t> A );
 
    void setDiagonal( const scalar_t *in, std::shared_ptr<localmatrixdata_t> A );
 
    void extractDiagonal( std::shared_ptr<localmatrixdata_t> A, scalar_t *buf );
 
    void
    getRowSums( std::shared_ptr<localmatrixdata_t> A, scalar_t *buf, const bool zero_first ) const;
 
    void getRowSumsAbsolute( std::shared_ptr<localmatrixdata_t> A,
                             scalar_t *buf,
                             const bool zero_first ) const;
 
    void setIdentity( std::shared_ptr<localmatrixdata_t> A );
 
    void copy( std::shared_ptr<const localmatrixdata_t> X, std::shared_ptr<localmatrixdata_t> Y );
 
    template<typename ConfigIn>
    static void copyCast( std::shared_ptr<CSRLocalMatrixData<ConfigIn>> X,
                          std::shared_ptr<localmatrixdata_t> Y );
 
protected:
    ExecSpace d_exec_space;
};
 
} // namespace AMP::LinearAlgebra
 
#endif
 
#endif