mpi_matrix.h

Go to the documentation of this file.

#pragma once
 
#include "mpi_gather_kron.h"
#include "sparseblockmat.h"
#include "memory.h"
#include "timer.h"
 
 
 
namespace dg {
namespace blas2 {
//forward declare blas2 symv functions
template< class MatrixType, class ContainerType1, class ContainerType2>
void symv( MatrixType&& M,
                  const ContainerType1& x,
                  ContainerType2& y);
template< class FunctorType, class MatrixType, class ContainerType1, class ContainerType2>
void stencil( FunctorType f, MatrixType&& M,
                  const ContainerType1& x,
                  ContainerType2& y);
template< class MatrixType, class ContainerType1, class ContainerType2>
void symv( get_value_type<ContainerType1> alpha,
                  MatrixType&& M,
                  const ContainerType1& x,
                  get_value_type<ContainerType1> beta,
                  ContainerType2& y);
namespace detail
{
template< class Stencil, class ContainerType, class ...ContainerTypes>
void doParallelFor( SharedVectorTag, Stencil f, unsigned N, ContainerType&& x, ContainerTypes&&... xs);
}
}//namespace blas2
 
 
template<template<class > class Vector, class LocalMatrixInner, class LocalMatrixOuter = LocalMatrixInner>
struct MPISparseBlockMat
{
    MPISparseBlockMat() = default;
    MPISparseBlockMat( const LocalMatrixInner& inside,
                const LocalMatrixOuter& outside,
                const MPIKroneckerGather<Vector>& mpi_gather // only gather foreign messages
                )
    : m_i(inside), m_o(outside), m_g(mpi_gather)
    {
    }
    template< template<class> class V, class LI, class LO>
    friend class MPISparseBlockMat; // enable copy
 
    template< template<class> class V, class LI, class LO>
    MPISparseBlockMat( const MPISparseBlockMat<V,LI,LO>& src)
    : m_i(src.m_i), m_o(src.m_o), m_g(src.m_g)
    {}
 
    const LocalMatrixInner& inner_matrix() const{return m_i;}
    const LocalMatrixOuter& outer_matrix() const{return m_o;}
    LocalMatrixInner& inner_matrix() {return m_i;}
    LocalMatrixOuter& outer_matrix() {return m_o;}
 
    MPI_Comm communicator() const { return m_g.communicator();}
 
    template<class ContainerType1, class ContainerType2>
    void symv( dg::get_value_type<ContainerType1> alpha, const ContainerType1&
        x, dg::get_value_type<ContainerType1> beta, ContainerType2& y) const
    {
        // ContainerType is MPI_Vector here...
        //the blas2 functions should make enough static assertions on tpyes
        if( !m_g.isCommunicating()) //no communication needed
        {
            dg::blas2::symv( alpha, m_i, x.data(), beta, y.data());
            return;
 
        }
        int result;
        MPI_Comm_compare( x.communicator(), y.communicator(), &result);
        assert( result == MPI_CONGRUENT || result == MPI_IDENT);
        using value_type = dg::get_value_type<ContainerType1>;
 
        m_buffer_ptrs.template set<const value_type*>( m_o.num_cols);
        auto& buffer_ptrs = m_buffer_ptrs.template get<const value_type*>();
        // 1 initiate communication
        m_g.global_gather_init( x.data());
        // 2 compute inner points
        dg::blas2::symv( alpha, m_i, x.data(), beta, y.data());
        // 3 wait for communication to finish
        m_g.global_gather_wait( x.data(), buffer_ptrs);
        // Local computation may be unnecessary
        if( buffer_ptrs.size() > 0)
        {
            // 4 compute and add outer points
            const value_type** b_ptrs = thrust::raw_pointer_cast( buffer_ptrs.data());
                  value_type*  y_ptr  = thrust::raw_pointer_cast( y.data().data());
            m_o.symv( SharedVectorTag(), dg::get_execution_policy<ContainerType1>(),
                alpha, b_ptrs, value_type(1.), y_ptr);
        }
    }
    template<class ContainerType1, class ContainerType2>
    void symv(const ContainerType1& x, ContainerType2& y) const
    {
        symv( 1, x, 0, y);
    }
    private:
    LocalMatrixInner m_i;
    LocalMatrixOuter m_o;
    MPIKroneckerGather<Vector> m_g;
    mutable detail::AnyVector<Vector>  m_buffer_ptrs;
};
 
 
template<class real_type, class ConversionPolicyRows, class ConversionPolicyCols>
auto make_mpi_sparseblockmat( // not actually communicating
    const EllSparseBlockMat<real_type, thrust::host_vector>& src,
    const ConversionPolicyRows& g_rows, // must be 1d
    const ConversionPolicyCols& g_cols
)
{
    // Indices in src are BLOCK INDICES!
    // g_rows and g_cols can have different n!!
    // Our approach here is:
    // 1. some preliminary MPI tests
    // 2. grab the correct rows of src and mark rows that are communicating
    // we need to grab the entire row to ensure reproducibility (which is guaranteed by order of computations)
    // 3. Convert inner rows to local and move outer rows to Coo format vectors (invalidating them in inner)
    // 4. Use global columns vector to construct MPI Kronecker Gather
    // 5. Use buffer index to construct local CooSparseBlockMat
    constexpr int invalid_idx = EllSparseBlockMat<real_type, thrust::host_vector>::invalid_index;
    int result;
    MPI_Comm_compare( g_rows.communicator(), g_cols.communicator(), &result);
    assert( result == MPI_CONGRUENT or result == MPI_IDENT);
    int rank;
    MPI_Comm_rank( g_cols.communicator(), &rank);
    int ndim;
    MPI_Cartdim_get( g_cols.communicator(), &ndim);
    assert( ndim == 1);
    int dim, period, coord;
    MPI_Cart_get( g_cols.communicator(), 1, &dim, &period, &coord); // coord of the calling process
    if( dim == 1)
    {
        return MPISparseBlockMat<thrust::host_vector, EllSparseBlockMat<real_type, thrust::host_vector>,
            CooSparseBlockMat<real_type, thrust::host_vector>>{ src, {}, {}};
    }
    unsigned n = src.n, rn = g_rows.n()/n, cn = g_cols.n()/n;
    int bpl = src.blocks_per_line;
    EllSparseBlockMat<real_type, thrust::host_vector> inner(g_rows.local().N()*rn, g_cols.local().N()*cn, // number of blocks
        bpl, src.data.size()/(src.n*src.n), src.n);
    inner.set_left_size( src.left_size);
    inner.set_right_size( src.right_size);
    //first copy data elements (even though not all might be needed it doesn't slow down things either)
    inner.data = src.data;
    std::vector<bool> local_row( inner.num_rows, true);
    //2. now grab the right rows of the cols and data indices (and mark communicating rows)
    for( int i=0; i<inner.num_rows; i++)
    for( int k=0; k<bpl; k++)
    {
        int gIdx=0;
        assert( g_rows.local2globalIdx( i*src.n, rank, gIdx)); // convert from block idx to idx !
        inner.data_idx[i*bpl + k] = src.data_idx[ gIdx/src.n*bpl + k];
        inner.cols_idx[i*bpl + k] = src.cols_idx[ gIdx/src.n*bpl + k];
        if( inner.cols_idx[i*bpl+k] == invalid_idx)
            continue;
        // find out if row is communicating
        int lIdx = 0, pid = 0;
        gIdx = inner.cols_idx[i*bpl + k]*src.n; // convert from block idx to idx !
        assert( g_cols.global2localIdx( gIdx, lIdx, pid));
        if ( pid != rank)
            local_row[i] = false;
    }
    // 3. Convert inner rows to local and move outer rows to Coo format vectors (invalidating them in inner)
    thrust::host_vector<std::array<int,2>> gColIdx;
    thrust::host_vector<int> rowIdx;
    thrust::host_vector<int> dataIdx;
    for( int i=0; i<inner.num_rows; i++)
    for( int k=0; k<bpl; k++)
    {
        int gIdx = inner.cols_idx[i*bpl + k];
        int lIdx = 0, pid = 0;
        if( gIdx != invalid_idx)
        {
            assert( g_cols.global2localIdx( gIdx*src.n, lIdx, pid));
            lIdx = lIdx/src.n; // convert from idx to block idx !
        }
        else
            lIdx = invalid_idx;
        if ( !local_row[i] )
        {
            // Simply skip invalid entries in Coo matrix
            if( gIdx != invalid_idx)
            {
                rowIdx.push_back( i);
                dataIdx.push_back( inner.data_idx[i*bpl+k]);
                gColIdx.push_back( {pid, lIdx});
                inner.cols_idx[i*bpl + k] = invalid_idx;
            }
        }
        else
        {
            inner.cols_idx[i*bpl + k] = lIdx;
        }
    }
    // Now make MPI Gather object
    thrust::host_vector<int> lColIdx;
    auto gather_map = dg::gIdx2unique_idx( gColIdx, lColIdx);
 
    MPIKroneckerGather<thrust::host_vector> mpi_gather( inner.left_size, gather_map, inner.n,
        inner.num_cols, inner.right_size, g_cols.communicator());
    CooSparseBlockMat<real_type, thrust::host_vector> outer(inner.num_rows, mpi_gather.buffer_size(),
        src.n, src.left_size, src.right_size);
    outer.data     = src.data; outer.cols_idx = lColIdx;
    outer.rows_idx = rowIdx;   outer.data_idx = dataIdx;
    outer.num_entries = rowIdx.size();
 
    return MPISparseBlockMat<thrust::host_vector, EllSparseBlockMat<real_type, thrust::host_vector>,
        CooSparseBlockMat<real_type, thrust::host_vector>>{ inner, outer, mpi_gather};
}
 
/*
// ///@cond
// namespace detail
// {
// // A combined Symv with a scatter of rows
// // Not good for cuda if number of columns is high so can be removed
// struct CSRSymvScatterFilter
// {
//     template<class real_type>
//     DG_DEVICE
//     void operator()( unsigned i, const int* scatter, const int* row_offsets,
//             const int* column_indices, const real_type* values,
//             const real_type* x, real_type* y)
//     {
//         for( int k=row_offsets[i]; k<row_offsets[i+1]; k++)
//             y[scatter[i]] += x[column_indices[k]]*values[k];
//     }
// };
// }// namespace detail
// /// @endcond
*/
enum dist_type
{
    row_dist=0, 
    allreduce=2 
};
 
 
 
template<template<class > class Vector, class LocalMatrixInner, class LocalMatrixOuter = LocalMatrixInner>
struct MPIDistMat
{
    MPIDistMat() = default;
 
    MPIDistMat( const LocalMatrixInner& m)
    : m_i(m), m_dist(row_dist){}
 
    template< template<class> class V, class LI, class LO>
    friend class MPIDistMat; // enable copy
    MPIDistMat( const LocalMatrixInner& inside,
                const LocalMatrixOuter& outside,
                const MPIGather<Vector>& mpi_gather,
                const Vector<int>& scatter // TODO scatter comp -> y from symv(outside, recv, comp);
                )
    : m_i(inside), m_o(outside), m_g(mpi_gather), m_scatter(scatter),
      m_dist( row_dist), m_reduce(mpi_gather.communicator()) { }
 
    MPIDistMat( const LocalMatrixInner& m, MPI_Comm comm)
    : m_i(m), m_dist(allreduce), m_reduce(comm){}
 
    template<template<class> class OtherVector, class OtherMatrixInner, class OtherMatrixOuter>
    MPIDistMat( const MPIDistMat<OtherVector, OtherMatrixInner, OtherMatrixOuter>& src)
    : m_i(src.inner_matrix()), m_o( src.outer_matrix()), m_g(src.mpi_gather()),
    m_scatter( src.m_scatter), m_dist( src.m_dist), m_reduce(src.communicator())
    { }
    const LocalMatrixInner& inner_matrix() const{return m_i;}
    const LocalMatrixOuter& outer_matrix() const{return m_o;}
    const MPIGather<Vector>& mpi_gather() const{return m_g;}
 
    MPI_Comm communicator() const{ return m_reduce.communicator();}
 
    const Vector<int>& scatter() const {return m_scatter;}
 
    template<class ContainerType1, class ContainerType2>
    void symv( const ContainerType1& x, ContainerType2& y) const
    {
        //the blas2 functions should make enough static assertions on tpyes
        if( !m_g.isCommunicating()) //no communication needed
        {
            dg::blas2::symv( m_i, x.data(), y.data());
            if( m_dist == allreduce)
            {
                m_reduce.reduce( y.data());
            }
            return;
 
        }
        assert( m_dist == row_dist);
        int result;
        MPI_Comm_compare( x.communicator(), y.communicator(), &result);
        assert( result == MPI_CONGRUENT || result == MPI_IDENT);
 
        using value_type = dg::get_value_type<ContainerType1>;
        m_recv_buffer.template set<value_type>( m_g.buffer_size());
        m_comp_buffer.template set<value_type>( m_o.num_rows());
        auto& recv_buffer = m_recv_buffer.template get<value_type>();
        auto& comp_buffer = m_comp_buffer.template get<value_type>();
 
        // 1 initiate communication
        m_g.global_gather_init( x.data(), recv_buffer);
        // 2 compute inner points
        dg::blas2::symv( m_i, x.data(), y.data());
        // 3 wait for communication to finish
        m_g.global_gather_wait( recv_buffer);
        if( comp_buffer.size() > 0)
        {
            // 4 compute and add outer points
            dg::blas2::symv( m_o, recv_buffer, comp_buffer);
            unsigned size = comp_buffer.size();
            dg::blas2::detail::doParallelFor( SharedVectorTag(),
                []DG_DEVICE( unsigned i, const value_type* x,
                    const int* idx, value_type* y) {
                    y[idx[i]] += x[i];
                }, size, comp_buffer, m_scatter, y.data());
            // What we need is a fused symv + scatter:
            // Something like
            // dg::blas2::parallel_for( detail::CSRSymvScatterFilter(),
            //        m_o.num_rows, m_scatter, m_o.row_offsets,
            //        m_o.column_indices, m_o.values, recv_buffer, y.data());
            // parallel_for is quite slow for ds_mpit because 50 entries per row
            // TODO Find out if cuSparse can do that and how fast
        }
    }
 
    template<class Functor, class ContainerType1, class ContainerType2>
    void stencil( const Functor f, const ContainerType1& x, ContainerType2& y) const
    {
        //the blas2 functions should make enough static assertions on tpyes
        if( !m_g.isCommunicating()) //no communication needed
        {
            dg::blas2::stencil( f, m_i, x.data(), y.data());
            if( m_dist == allreduce)
            {
                m_reduce.reduce( y.data());
            }
            return;
        }
        assert( m_dist == row_dist);
        int result;
        MPI_Comm_compare( x.communicator(), y.communicator(), &result);
        assert( result == MPI_CONGRUENT || result == MPI_IDENT);
 
        using value_type = dg::get_value_type<ContainerType1>;
        m_recv_buffer.template set<value_type>( m_g.buffer_size());
        m_comp_buffer.template set<value_type>( m_o.num_rows());
        auto& recv_buffer = m_recv_buffer.template get<value_type>();
        auto& comp_buffer = m_comp_buffer.template get<value_type>();
 
        // 1 initiate communication
        m_g.global_gather_init( x.data(), recv_buffer);
        // 2 compute inner points
        dg::blas2::stencil( f, m_i, x.data(), y.data());
        // 3 wait for communication to finish
        m_g.global_gather_wait( recv_buffer);
        // 4 compute and add outer points
        dg::blas2::stencil( f, m_o, recv_buffer, comp_buffer);
        unsigned size = comp_buffer.size();
        dg::blas2::detail::doParallelFor( SharedVectorTag(),
            [size]DG_DEVICE( unsigned i, const value_type* x,
                const int* idx, value_type* y) {
                y[idx[i]] += x[i];
            }, size, comp_buffer, m_scatter, y.data());
        // What we really need here is a sparse vector format
        // so we can fuse the symv with the scatter
        // OR y += Ax
    }
 
    private:
    LocalMatrixInner m_i;
    LocalMatrixOuter m_o;
    MPIGather<Vector> m_g;
    Vector<int> m_scatter;
    enum dist_type m_dist;
    detail::MPIAllreduce m_reduce;
    mutable detail::AnyVector<Vector>  m_recv_buffer, m_comp_buffer ;
};
 
 
template<template<class>class V, class LI, class LO>
struct TensorTraits<MPISparseBlockMat<V, LI,LO> >
{
    using value_type = get_value_type<LI>;
    using tensor_category = MPIMatrixTag;
};
template<template<class>class V, class LI, class LO>
struct TensorTraits<MPIDistMat<V, LI,LO> >
{
    using value_type = get_value_type<LI>;
    using tensor_category = MPIMatrixTag;
};
 
} //namespace dg