CSRLocalMatrixData.h

Go to the documentation of this file.

#ifndef included_AMP_CSRLocalMatrixData_h
#define included_AMP_CSRLocalMatrixData_h
 
#include "AMP/matrices/MatrixParametersBase.h"
#include "AMP/matrices/data/MatrixData.h"
#include "AMP/utils/Memory.h"
#include "AMP/utils/Utilities.h"
 
#include <algorithm>
#include <functional>
#include <map>
#include <numeric>
#include <tuple>
#include <unordered_map>
 
namespace AMP::IO {
class RestartManager;
}
 
namespace AMP::Discretization {
class DOFManager;
}
 
namespace AMP::LinearAlgebra {
 
// Forward declare CSRMatrix{Data,Communicator} to make them friends
template<typename C>
class CSRMatrixData;
template<typename C>
class CSRMatrixCommunicator;
template<typename C>
class CSRMatrixSpGEMMDefault;
template<typename C>
class CSRMatrixSpGEMMDevice;
 
template<typename Config>
class CSRLocalMatrixData : public AMP::enable_shared_from_this<CSRLocalMatrixData<Config>>
{
public:
    template<typename C>
    friend class CSRMatrixData;
    template<typename C>
    friend class CSRLocalMatrixData;
    template<typename C>
    friend class CSRMatrixCommunicator;
    template<typename C>
    friend class CSRMatrixSpGEMMDefault;
    template<typename C>
    friend class CSRMatrixSpGEMMDevice;
 
    using mask_t         = unsigned char;
    using gidx_t         = typename Config::gidx_t;
    using lidx_t         = typename Config::lidx_t;
    using scalar_t       = typename Config::scalar_t;
    using allocator_type = typename Config::allocator_type;
    static_assert( std::is_same_v<typename allocator_type::value_type, void> );
 
    explicit CSRLocalMatrixData( std::shared_ptr<MatrixParametersBase> params,
                                 AMP::Utilities::MemoryType memory_location,
                                 typename Config::gidx_t first_row,
                                 typename Config::gidx_t last_row,
                                 typename Config::gidx_t first_col,
                                 typename Config::gidx_t last_col,
                                 bool is_diag,
                                 bool is_symbolic = false,
                                 uint64_t hash    = 0 );
 
    virtual ~CSRLocalMatrixData();
 
    std::string type() const;
 
    std::uint16_t mode() const { return static_cast<std::uint16_t>( Config::mode ); }
 
    std::tuple<lidx_t *, gidx_t *, lidx_t *, scalar_t *> getDataFields()
    {
        return std::make_tuple(
            d_row_starts.get(), d_cols.get(), d_cols_loc.get(), d_coeffs.get() );
    }
 
    std::tuple<const lidx_t *, const gidx_t *, const lidx_t *, const scalar_t *>
    getDataFields() const
    {
        return std::make_tuple(
            d_row_starts.get(), d_cols.get(), d_cols_loc.get(), d_coeffs.get() );
    }
 
    void swapDataFields( CSRLocalMatrixData<Config> &other );
 
    lidx_t *getRowStarts() { return d_row_starts.get(); }
 
    auto getMemoryLocation() const { return d_memory_location; }
 
    bool isDiag() const { return d_is_diag; }
 
    bool isEmpty() const { return d_is_empty; }
 
    lidx_t numberOfNonZeros() const { return d_nnz; }
 
    lidx_t numLocalRows() const { return d_last_row - d_first_row; }
 
    lidx_t numLocalColumns() const { return d_last_col - d_first_col; }
 
    lidx_t numUniqueColumns() const
    {
        return d_is_diag ? ( d_last_col - d_first_col ) : d_ncols_unq;
    }
 
    gidx_t beginRow() const { return d_first_row; }
 
    gidx_t endRow() const { return d_last_row; }
 
    gidx_t beginCol() const { return d_first_col; }
 
    gidx_t endCol() const { return d_last_col; }
 
    void globalToLocalColumns();
 
    gidx_t *getColumnMap() const
    {
        if ( d_is_diag ) {
            return nullptr;
        }
        return d_cols_unq.get();
    }
 
    size_t *getColumnMapSizeT() const;
 
    scalar_t *getGhostCache() const;
 
    template<typename idx_t>
    void getColumnMap( std::vector<idx_t> &colMap ) const
    {
        // Don't do anything if offd and empty
        if ( !d_is_diag && d_is_empty ) {
            return;
        }
 
        colMap.resize( d_is_diag ? ( d_last_col - d_first_col ) : d_ncols_unq, 0 );
 
        if ( d_is_diag ) {
            std::iota( colMap.begin(), colMap.end(), d_first_col );
        } else {
            AMP::Utilities::copy<gidx_t, idx_t>( d_ncols_unq, d_cols_unq.get(), colMap.data() );
        }
    }
 
 
    void setNNZ( lidx_t tot_nnz );
 
    void setNNZ( const lidx_t *nnz );
 
    void setNNZ( bool do_accum );
 
    void removeRange( const scalar_t bnd_lo, const scalar_t bnd_up );
 
    void getColPtrs( std::vector<gidx_t *> &col_ptrs );
 
    void printStats( bool verbose, bool show_zeros ) const;
 
    void printAll( bool force = false ) const;
 
    std::shared_ptr<CSRLocalMatrixData> maskMatrixData( const mask_t *mask,
                                                        const bool is_symbolic ) const;
 
    static std::shared_ptr<CSRLocalMatrixData>
    ConcatHorizontal( std::shared_ptr<MatrixParametersBase> params,
                      std::map<int, std::shared_ptr<CSRLocalMatrixData>> blocks );
 
    static std::shared_ptr<CSRLocalMatrixData>
    ConcatVertical( std::shared_ptr<MatrixParametersBase> params,
                    std::map<int, std::shared_ptr<CSRLocalMatrixData>> blocks,
                    const gidx_t first_col,
                    const gidx_t last_col,
                    const bool is_diag );
 
    template<typename U>
    static std::shared_ptr<U[]> sharedArrayBuilder( const size_t N )
    {
        using alloc_t = typename std::allocator_traits<allocator_type>::template rebind_alloc<U>;
        alloc_t alloc;
        return std::shared_ptr<typename alloc_t::value_type[]>(
            alloc.allocate( N ), [N, &alloc]( auto p ) -> void { alloc.deallocate( p, N ); } );
    }
 
    static std::shared_ptr<lidx_t[]> makeLidxArray( const size_t N )
    {
        return sharedArrayBuilder<lidx_t>( N );
    }
 
    static std::shared_ptr<gidx_t[]> makeGidxArray( const size_t N )
    {
        return sharedArrayBuilder<gidx_t>( N );
    }
 
    static std::shared_ptr<scalar_t[]> makeScalarArray( const size_t N )
    {
        return sharedArrayBuilder<scalar_t>( N );
    }
 
public: // Non virtual functions
    uint64_t getID() const { return d_hash; }
 
public: // Write/read restart data
    void registerChildObjects( AMP::IO::RestartManager *manager ) const;
 
    void writeRestart( int64_t fid ) const;
 
    CSRLocalMatrixData( int64_t fid, AMP::IO::RestartManager *manager );
 
protected:
    void sortColumns();
 
    std::shared_ptr<CSRLocalMatrixData> cloneMatrixData();
 
    template<typename ConfigOut>
    std::shared_ptr<CSRLocalMatrixData<ConfigOut>> migrate() const;
 
    std::shared_ptr<CSRLocalMatrixData>
    transpose( std::shared_ptr<MatrixParametersBase> params ) const;
 
    void getRowByGlobalID( const size_t local_row,
                           std::vector<size_t> &cols,
                           std::vector<double> &values ) const;
 
    void getValuesByGlobalID( const size_t local_row,
                              const size_t num_cols,
                              const size_t *cols,
                              scalar_t *values ) const;
 
    void addValuesByGlobalID( const size_t local_row,
                              const size_t num_cols,
                              const size_t *cols,
                              const scalar_t *vals );
 
    void setValuesByGlobalID( const size_t local_row,
                              const size_t num_cols,
                              const size_t *cols,
                              const scalar_t *vals );
 
    size_t numberColumnIDs( size_t local_row ) const;
 
    std::vector<size_t> getColumnIDs( const size_t local_row ) const;
 
    // Data members passed from outer CSRMatrixData object
    AMP::Utilities::MemoryType d_memory_location;
    gidx_t d_first_row;
    gidx_t d_last_row;
    gidx_t d_first_col;
    gidx_t d_last_col;
 
    bool d_is_diag;
    bool d_is_empty = true;
    bool d_is_symbolic = false;
 
    std::shared_ptr<lidx_t[]> d_row_starts;
    std::shared_ptr<gidx_t[]> d_cols;
    std::shared_ptr<gidx_t[]> d_cols_unq;
    std::shared_ptr<lidx_t[]> d_cols_loc;
    std::shared_ptr<scalar_t[]> d_coeffs;
 
    lidx_t d_num_rows;
    lidx_t d_nnz = 0;
    lidx_t d_ncols_unq = 0;
    uint64_t d_hash = 0;
 
    mutable std::shared_ptr<size_t[]> d_cols_unq_size_t;
 
    mutable std::shared_ptr<scalar_t[]> d_ghost_cache;
};
 
} // namespace AMP::LinearAlgebra
 
#endif