mpi_gather.h

Go to the documentation of this file.

#pragma once
#include <cassert>
#include <thrust/host_vector.h>
 
#include "exceptions.h"
#include "config.h"
#include "mpi_datatype.h"
#include "mpi_permutation.h"
#include "tensor_traits.h"
#include "memory.h"
#include "index.h"
 
namespace dg{
 
namespace detail{
 
// Used for Average operation
struct MPIAllreduce
{
    MPIAllreduce( MPI_Comm comm = MPI_COMM_NULL) : m_comm(comm){}
    MPI_Comm communicator() const{ return m_comm;}
    template<class ContainerType> // a Shared Vector
    void reduce( ContainerType& y) const
    {
        using value_type = dg::get_value_type<ContainerType>;
        void * send_ptr = thrust::raw_pointer_cast(y.data());
        if constexpr (dg::has_policy_v<ContainerType, dg::CudaTag>)
        {
#ifdef __CUDACC__ // g++ does not know cuda code
            // cuda - sync device
            cudaError_t code = cudaGetLastError( );
            if( code != cudaSuccess)
                throw dg::Error(dg::Message(_ping_)<<cudaGetErrorString(code));
            if constexpr ( not dg::cuda_aware_mpi)
            {
                m_h_buffer.template set<value_type>( y.size());
                auto& h_buffer = m_h_buffer.template get<value_type>();
                code = cudaMemcpy( &h_buffer[0], send_ptr,
                    y.size()*sizeof(value_type), cudaMemcpyDeviceToHost);
                send_ptr = // re-point pointer
                    thrust::raw_pointer_cast(&h_buffer[0]);
                if( code != cudaSuccess)
                    throw dg::Error(dg::Message(_ping_)<<cudaGetErrorString(code));
            }
            // We have to wait that all kernels are finished and values are
            // ready to be sent
            code = cudaDeviceSynchronize();
            if( code != cudaSuccess)
                throw dg::Error(dg::Message(_ping_)<<cudaGetErrorString(code));
#else
            assert( false && "Something is wrong! This should never execute!");
#endif
        }
        MPI_Allreduce( MPI_IN_PLACE, send_ptr, y.size(),
            getMPIDataType<value_type>(), MPI_SUM, m_comm);
        if constexpr (dg::has_policy_v<ContainerType, dg::CudaTag>
            and not dg::cuda_aware_mpi)
            y = m_h_buffer.template get<value_type>();
    }
    private:
    MPI_Comm m_comm;
    mutable detail::AnyVector<thrust::host_vector>  m_h_buffer;
};
 
struct MPIContiguousGather
{
    MPIContiguousGather( MPI_Comm comm = MPI_COMM_NULL)
    : m_comm(comm), m_communicating(false){ }
    MPIContiguousGather(
        const std::map<int, thrust::host_vector<MsgChunk>>& recvMsg,
        MPI_Comm comm)
    : m_comm(comm), m_recvMsg( recvMsg)
    {
        m_sendMsg = mpi_permute ( recvMsg, comm);
        m_communicating = is_communicating( recvMsg, comm);
        // if cuda unaware we need to send messages through host
        for( auto& chunks : m_sendMsg)
        for( auto& chunk : chunks.second)
            m_store_size += chunk.size;
        resize_rqst();
    }
 
    static const std::map<int, thrust::host_vector<MsgChunk>> make_chunks(
        const std::map<int, thrust::host_vector<int> > &recvIdx, int chunk_size = 1)
    {
        std::map<int, thrust::host_vector<MsgChunk>> recvChunk;
        for( auto& idx: recvIdx)
        {
            auto chunks = detail::find_contiguous_chunks( idx.second);
            for( auto& chunk : chunks)
            {
                recvChunk[idx.first].push_back( {chunk.idx*chunk_size,
                    chunk.size*chunk_size});
            }
        }
        return recvChunk;
    }
 
    MPI_Comm communicator() const{return m_comm;}
    unsigned buffer_size( bool self_communication = true) const
    {
        unsigned buffer_size = 0;
        int rank;
        MPI_Comm_rank( m_comm, &rank);
        // We need to find out the minimum amount of memory we need to allocate
        for( auto& chunks : m_recvMsg) // first is PID, second is vector of chunks
        for( auto& chunk : chunks.second)
        {
            if( chunks.first == rank and not self_communication)
                continue;
            buffer_size += chunk.size;
        }
        return buffer_size;
    }
 
    bool isCommunicating() const{
        return m_communicating;
    }
    // if not self_communication  then buffer can be smaller
    template<class ContainerType0, class ContainerType1>
    void global_gather_init( const ContainerType0& gatherFrom, ContainerType1& buffer,
        bool self_communication = true) const
    {
        // TODO only works if m_recvMsg is non-overlapping
        using value_type = dg::get_value_type<ContainerType0>;
        static_assert( std::is_same_v<value_type,
                get_value_type<ContainerType1>>);
        int rank;
        MPI_Comm_rank( m_comm, &rank);
        // BugFix: buffer value_type must be set even if no messages are sent
        // so that global_gather_wait works
        if constexpr (dg::has_policy_v<ContainerType1, dg::CudaTag>
            and not dg::cuda_aware_mpi)
        {
            m_h_store.template set<value_type>( m_store_size);
            m_h_buffer.template set<value_type>( buffer_size(self_communication));
        }
        // Receives (we implicitly receive chunks in the order)
        unsigned start = 0;
        unsigned rqst_counter = 0;
        for( auto& msg : m_recvMsg) // first is PID, second is vector of chunks
        for( unsigned u=0; u<msg.second.size(); u++)
        {
            if( msg.first == rank and not self_communication)
                continue;
            auto chunk = msg.second[u];
            void * recv_ptr;
            assert( buffer.size() >= unsigned(start + chunk.size - 1));
            if constexpr (dg::has_policy_v<ContainerType1, dg::CudaTag>
                and not dg::cuda_aware_mpi)
            {
                auto& h_buffer = m_h_buffer.template get<value_type>();
                recv_ptr = thrust::raw_pointer_cast( h_buffer.data())
                   + start;
            }
            else
                recv_ptr = thrust::raw_pointer_cast( buffer.data())
                   + start;
            MPI_Irecv( recv_ptr, chunk.size,
                   getMPIDataType<value_type>(),  //receiver
                   msg.first, u, m_comm, &m_rqst[rqst_counter]);  //source
            rqst_counter ++;
            start += chunk.size;
        }
 
        // Send
        start = 0;
        for( auto& msg : m_sendMsg) // first is PID, second is vector of chunks
        for( unsigned u=0; u<msg.second.size(); u++)
        {
            if( msg.first == rank and not self_communication)
                continue;
            auto chunk = msg.second[u];
            const void * send_ptr = thrust::raw_pointer_cast(gatherFrom.data()) + chunk.idx;
            assert( gatherFrom.size() >= unsigned(chunk.idx + chunk.size - 1));
            if constexpr (dg::has_policy_v<ContainerType0, dg::CudaTag>)
            {
#ifdef __CUDACC__ // g++ does not know cuda code
                // cuda - sync device
                cudaError_t code = cudaGetLastError( );
                if( code != cudaSuccess)
                    throw dg::Error(dg::Message(_ping_)<<cudaGetErrorString(code));
                if constexpr ( not dg::cuda_aware_mpi)
                {
                    auto& h_store = m_h_store.template get<value_type>();
                    code = cudaMemcpy( &h_store[start], send_ptr,
                        chunk.size*sizeof(value_type), cudaMemcpyDeviceToHost);
                    send_ptr = // re-point pointer
                        thrust::raw_pointer_cast(&h_store[start]);
                    if( code != cudaSuccess)
                        throw dg::Error(dg::Message(_ping_)<<cudaGetErrorString(code));
                }
                // We have to wait that all kernels are finished and values are
                // ready to be sent
                code = cudaDeviceSynchronize();
                if( code != cudaSuccess)
                    throw dg::Error(dg::Message(_ping_)<<cudaGetErrorString(code));
#else
                assert( false && "Something is wrong! This should never execute!");
#endif
            }
            MPI_Isend( send_ptr, chunk.size,
                   getMPIDataType<value_type>(),  //sender
                   msg.first, u, m_comm, &m_rqst[rqst_counter]);  //destination
            rqst_counter ++;
            start+= chunk.size;
        }
    }
 
    template<class ContainerType>
    void global_gather_wait( ContainerType& buffer) const
    {
        using value_type = dg::get_value_type<ContainerType>;
        MPI_Waitall( m_rqst.size(), &m_rqst[0], MPI_STATUSES_IGNORE );
        if constexpr (dg::has_policy_v<ContainerType, dg::CudaTag>
                and not dg::cuda_aware_mpi)
            buffer = m_h_buffer.template get<value_type>();
    }
    private:
    MPI_Comm m_comm; // from constructor
    bool m_communicating = false;
    std::map<int,thrust::host_vector<MsgChunk>> m_sendMsg;
    std::map<int,thrust::host_vector<MsgChunk>> m_recvMsg; // from constructor
 
    mutable detail::AnyVector<thrust::host_vector>  m_h_buffer;
 
    unsigned m_store_size = 0;
    mutable detail::AnyVector<thrust::host_vector>  m_h_store;
 
    mutable std::vector<MPI_Request> m_rqst;
    void resize_rqst()
    {
        unsigned rqst_size = 0;
        // number of messages to send and receive
        for( auto& msg : m_recvMsg)
            rqst_size += msg.second.size();
        for( auto& msg : m_sendMsg)
            rqst_size += msg.second.size();
        m_rqst.resize( rqst_size, MPI_REQUEST_NULL);
    }
};
}//namespace detail
 
template< template <class> class Vector>
struct MPIGather
{
 
    MPIGather( MPI_Comm comm = MPI_COMM_NULL) : m_mpi_gather(comm){ }
 
    MPIGather(
        const std::map<int, thrust::host_vector<int>>& recvIdx, // should be unique global indices (->gIdx2unique)
        MPI_Comm comm) : MPIGather( recvIdx, 1, comm)
    { }
    MPIGather(
        const std::map<int, thrust::host_vector<int>>& recvIdx, // in units of chunk size
        unsigned chunk_size, // can be 1 (contiguous indices in recvIdx are concatenated)
        MPI_Comm comm)
    {
        // TODO Catch wrong size of recvIdx
        static_assert( dg::is_vector_v<Vector<double>, SharedVectorTag>,
                "Only Shared vectors allowed");
        // The idea is that recvIdx and sendIdx completely define the communication pattern
        // and we can choose an optimal implementation
        // Actually the MPI library should do this but general gather indices
        // don't seem to be supported
        // So for now let's just use MPI_Alltoallv
        // Possible optimization includes
        // - check for duplicate messages to save internal buffer memory
        // - check for contiguous messages to avoid internal buffer memory
        // - check if an MPI_Type could fit the index map
        //v -> store -> buffer -> w
        // G_1 P G_2 v
        // G_2^T P^T G_1^T w
        // We need some logic to determine if we should pre-gather messages into a store
        unsigned num_messages = 0;
        auto recvChunks = detail::MPIContiguousGather::make_chunks(
            recvIdx, chunk_size);
        for( auto& chunks : recvChunks)
            num_messages+= chunks.second.size(); // size of vector of chunks
        unsigned size = recvChunks.size(); // number of pids involved
        double avg_msg_per_pid = (double)num_messages/(double)size; // > 1
        MPI_Allreduce( MPI_IN_PLACE, &avg_msg_per_pid, 1, MPI_DOUBLE, MPI_MAX, comm);
        m_contiguous = ( avg_msg_per_pid < 10); // 10 is somewhat arbitrary
        if( not m_contiguous) // messages are too fractioned
        {
            m_contiguous = false;
            // bootstrap communication pattern
            auto sendIdx = mpi_permute ( recvIdx, comm);
            auto lIdx = detail::flatten_map(sendIdx);
            thrust::host_vector<int> g2;
            unsigned start = 0;
            for( auto& idx : sendIdx)
            {
                for( unsigned l=0; l<idx.second.size(); l++)
                {
                    for( unsigned k=0; k<chunk_size; k++)
                    {
                        g2.push_back(idx.second[l] + k);
                    }
                    idx.second[l] = start + l; // repoint index to store index in units of chunk_size
                }
                start += idx.second.size();
            }
            m_g2 = g2;
            // bootstrap communication pattern back to recvIdx
            // (so that recvIdx has index into same store)
            auto store_recvIdx = mpi_permute( sendIdx, comm);
            auto store_recvChunks = detail::MPIContiguousGather::make_chunks(
                store_recvIdx, chunk_size);
            m_mpi_gather = detail::MPIContiguousGather( store_recvChunks, comm);
        }
        else
            m_mpi_gather = detail::MPIContiguousGather( recvChunks, comm);
 
    }
    template<class ArrayVec = thrust::host_vector<std::array<int,2>>,
        class IntVec = thrust::host_vector<int>>
    MPIGather(
        const ArrayVec& gather_map,
        IntVec& bufferIdx,
        MPI_Comm comm)
        : MPIGather( gIdx2unique_idx ( gather_map, bufferIdx), comm)
    {
    }
 
    template< template<class> class OtherVector>
    friend class MPIGather; // enable copy
 
    template< template<typename > typename OtherVector>
    MPIGather( const MPIGather<OtherVector>& src)
    :   m_contiguous( src.m_contiguous),
        m_g2( src.m_g2),
        m_mpi_gather( src.m_mpi_gather)
        // we don't need to copy memory buffers (they are just workspace) or the request
    {
    }
 
 
    MPI_Comm communicator() const{return m_mpi_gather.communicator();}
 
    bool isContiguous() const { return m_contiguous;}
    unsigned buffer_size() const { return m_mpi_gather.buffer_size();}
 
    bool isCommunicating() const
    {
        return m_mpi_gather.isCommunicating();
    }
    template<class ContainerType0, class ContainerType1>
    void global_gather_init( const ContainerType0& gatherFrom, ContainerType1& buffer) const
    {
        using value_type = dg::get_value_type<ContainerType0>;
        if( not m_contiguous)
        {
            m_store.template set<value_type>( m_g2.size());
            auto& store = m_store.template get<value_type>();
            thrust::gather( m_g2.begin(), m_g2.end(), gatherFrom.begin(),
                store.begin());
            m_mpi_gather.global_gather_init( store, buffer);
        }
        else
            m_mpi_gather.global_gather_init( gatherFrom, buffer);
    }
    template<class ContainerType>
    void global_gather_wait( ContainerType& buffer) const
    {
        m_mpi_gather.global_gather_wait( buffer);
    }
 
    private:
 
    bool m_contiguous = false;
    Vector<int> m_g2;
    dg::detail::MPIContiguousGather m_mpi_gather;
    // These are mutable and we never expose them to the user
    //unsigned m_store_size;// not needed because g2.index_map.size()
    mutable detail::AnyVector< Vector> m_store;
 
};
 
}//namespace dg