mpi_gather.h

Go to the documentation of this file.

#pragma once
#include <cassert>
#include <thrust/host_vector.h>
#include <thrust/copy.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"
#ifdef DG_WITH_NCCL
#include "nccl.h"
namespace dg
{
namespace detail
{
 
// Note from ncclCommGetAsyncError documentation: Query progress and/or errors
// from asynchronous nccl function
//
// nccl functions without stream argument are complete as soon as the return
// value is ncclSuccess. nccl functions with stream arguments return
// ncclSuccess as soon as the operation kernel is submitted to the stream. If
// the state is ncclInProgress no other kernels are allowed to be issued on the
// stream used by nccl (MW: unless it is another nccl function within ncclGroup
// calls I guess?) If there is an error on a communicator it should be
// destroyed using ncclCommAbort
void ncclCommSynchronize(ncclComm_t comm)
{
    // Wait untill all Kernels are submitted, or function returns
    ncclResult_t state;
    do
    {
        ncclResult_t result = ncclCommGetAsyncError( comm, &state);
        if( result != ncclSuccess)
            throw dg::Error(dg::Message(_ping_)<<ncclGetErrorString(result));
    }
    while( state == ncclInProgress);
    if( state != ncclSuccess)
        throw dg::Error(dg::Message(_ping_)<<ncclGetErrorString(state));
}
 
// Similar to exblas/mpi_accumulate, this keeps a static map of comms we already created
inline void getNcclComm( MPI_Comm comm, ncclComm_t * ncclcomm, cudaStream_t * stream )
{
    static std::map<MPI_Comm, std::pair<ncclComm_t, cudaStream_t>> comms;
    if( comms.count(comm) == 1 )
    {
        *ncclcomm = std::get<0>(comms[comm]);
        *stream = std::get<1>(comms[comm]);
        return;
    }
    int rank, size;
    MPI_Comm_rank( comm, &rank);
    MPI_Comm_size( comm, &size);
 
    ncclUniqueId ncclid;
    // Check the nccl user-guide for options
    ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
    config.blocking = 0; // important! We want asynchronous communication (MW but has seg-fault on nccl-2.22 and super slow on nccl-2.28)
    // If blocking = 1 the ncclCommSynchronize calls should be deleted
 
    if( rank == 0)
        ncclGetUniqueId(&ncclid);
    MPI_Bcast((void *)&ncclid, sizeof(ncclid), MPI_BYTE, 0, comm);
    ncclCommInitRankConfig( ncclcomm, size, ncclid, rank, &config);
    ncclCommSynchronize( *ncclcomm);
 
 
    cudaStreamCreate( stream);
 
    comms[comm] = {*ncclcomm, *stream};
}
template<class value_type>
inline ncclDataType_t getNcclDataType(){ assert( false && "Type not supported!\n" ); return ncclDataType_t{}; }
template<> inline ncclDataType_t getNcclDataType<char>(){ return ncclChar;}
template<> inline ncclDataType_t getNcclDataType<signed char>(){ return ncclChar;}
template<> inline ncclDataType_t getNcclDataType<unsigned char>(){ return ncclUint8;}
template<> inline ncclDataType_t getNcclDataType<int>(){ return ncclInt;}
template<> inline ncclDataType_t getNcclDataType<unsigned int>(){ return ncclUint32;}
template<> inline ncclDataType_t getNcclDataType<signed long int>(){ return ncclInt64;}
template<> inline ncclDataType_t getNcclDataType<unsigned long int>(){ return ncclUint64;}
template<> inline ncclDataType_t getNcclDataType<float>(){ return ncclFloat;}
template<> inline ncclDataType_t getNcclDataType<double>(){ return ncclDouble;}
 
// We need to separately catch the complex nature of the array though
template<> inline ncclDataType_t getNcclDataType<std::complex<float>>(){ return ncclFloat;}
template<> inline ncclDataType_t getNcclDataType<std::complex<double>>(){ return ncclDouble;}
template<> inline ncclDataType_t getNcclDataType<thrust::complex<float>>(){ return ncclFloat;}
template<> inline ncclDataType_t getNcclDataType<thrust::complex<double>>(){ return ncclDouble;}
 
}// namespace detail
}// namespace dg
#endif // DG_WITH_NCCL
 
namespace dg{
 
namespace detail{
inline void cuda_check_error_and_sync_device()
{
#ifdef __CUDACC__ // g++ does not know cuda code (happens for device=cpu, code must be valid)
    // cuda - sync device
    cudaError_t code = cudaGetLastError( );
    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));
#endif // __CUDACC__
}
 
// 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
    {
 
        if constexpr (dg::has_policy_v<ContainerType, dg::CudaTag>)
        {
            detail::cuda_check_error_and_sync_device();
            if constexpr ( dg::nccl_mpi)
                nccl_reduce( y);
            else if constexpr ( dg::cuda_aware_mpi)
                mpi_reduce( y);
            else
                cuda_unaware_reduce( y);
        }
        else
            mpi_reduce(y);
    }
    private:
    template<class ContainerType>
    void mpi_reduce( ContainerType& y) const
    {
        using value_type = dg::get_value_type<ContainerType>;
        void * send_ptr = thrust::raw_pointer_cast(y.data());
        MPI_Allreduce( MPI_IN_PLACE, send_ptr, y.size(),
            getMPIDataType<value_type>(), MPI_SUM, m_comm);
    }
    template<class ContainerType>
    void nccl_reduce( ContainerType& y) const
    {
#ifdef DG_WITH_NCCL
        using value_type = dg::get_value_type<ContainerType>;
        ncclComm_t ncclcomm;
        cudaStream_t stream;
        detail::getNcclComm( m_comm, &ncclcomm, &stream);
        unsigned ysize = y.size();
        if constexpr( std::is_same_v<dg::get_tensor_category<value_type>, dg::ComplexTag>)
            ysize *= 2;
        // Since this operation is put on stream != default_stream this will execute asynchronously
        void * send_ptr = thrust::raw_pointer_cast(y.data());
        ncclAllReduce( send_ptr, send_ptr, ysize,
            detail::getNcclDataType<value_type>(), ncclSum, ncclcomm, stream);
    // Wait for Kernel to be issued to stream
    ncclCommSynchronize( ncclcomm);
    // Sync the stream
    cudaStreamSynchronize(stream);
#endif // DG_WITH_NCCL
    }
    template<class ContainerType>
    void cuda_unaware_reduce( ContainerType& y) const
    {
        using value_type = dg::get_value_type<ContainerType>;
        m_h_buffer.template set<value_type>( y.size());
        auto& h_buffer = m_h_buffer.template get<value_type>();
        thrust::copy( y.begin(), y.end(), h_buffer.begin() ); // works even if y is a view
        mpi_reduce( h_buffer);
        thrust::copy( h_buffer.begin(), h_buffer.end(), y.begin() ); // works even if y is a view
    }
    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);
        resize_rqst();
        m_gatherFrom_minSize = 0;
        for( auto& chunks : m_sendMsg) // first is PID, second is vector of chunks
        for( auto& chunk : chunks.second)
            m_gatherFrom_minSize = std::max( m_gatherFrom_minSize,
                (unsigned)(chunk.idx+chunk.size));
    }
 
    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
    {
        return msg_size( m_recvMsg, self_communication);
    }
 
    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
    {
        if( gatherFrom.size() < m_gatherFrom_minSize)
            throw dg::Error(dg::Message(_ping_)<<"In MPIGather GatherFrom size "
                <<gatherFrom.size() << " must be at least "<<m_gatherFrom_minSize);
        if( buffer.size() < buffer_size( self_communication))
            throw dg::Error(dg::Message(_ping_)<<"In MPIGather buffer     size "
                <<buffer.size() << " must be at least "<<buffer_size( self_communication));
        int rank;
        MPI_Comm_rank( m_comm, &rank);
        // 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>>);
        if constexpr (dg::has_policy_v<ContainerType0, dg::CudaTag>)
        {
            detail::cuda_check_error_and_sync_device();
            if constexpr ( dg::nccl_mpi)
                nccl_global_gather_init( gatherFrom, buffer,
                    self_communication, rank, m_recvMsg, m_sendMsg);
            else if constexpr ( dg::cuda_aware_mpi)
                mpi_global_gather_init( gatherFrom, buffer,
                    self_communication, rank, m_recvMsg, m_sendMsg);
            else
                cuda_unaware_global_gather_init( gatherFrom, buffer,
                    self_communication, rank, m_recvMsg, m_sendMsg);
        }
        else
            mpi_global_gather_init( gatherFrom, buffer,
                self_communication, rank, m_recvMsg, m_sendMsg);
    }
 
    template<class ContainerType>
    void global_gather_wait( ContainerType& buffer) const
    {
        using value_type = dg::get_value_type<ContainerType>;
        if constexpr (dg::has_policy_v<ContainerType, dg::CudaTag>)
        {
            if constexpr ( dg::nccl_mpi)
            {
#ifdef DG_WITH_NCCL
    // Make sure all Send/Recv Kernels are actually issued to the stream
                ncclComm_t ncclcomm;
                cudaStream_t stream;
                detail::getNcclComm( m_comm, &ncclcomm, &stream);
        ncclCommSynchronize( ncclcomm);
                cudaStreamSynchronize(stream);
#endif // DG_WITH_NCCL
            }
            else if constexpr ( dg::cuda_aware_mpi)
                MPI_Waitall( m_rqst.size(), &m_rqst[0], MPI_STATUSES_IGNORE );
            else
            {
                MPI_Waitall( m_rqst.size(), &m_rqst[0], MPI_STATUSES_IGNORE );
                auto & h_buffer = m_h_buffer.template get<value_type>();
                thrust::copy( h_buffer.begin(), h_buffer.end(), buffer.begin() );
            }
        }
        else
            MPI_Waitall( m_rqst.size(), &m_rqst[0], MPI_STATUSES_IGNORE );
 
    }
    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_gatherFrom_minSize = 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);
    }
    unsigned msg_size( const std::map<int,thrust::host_vector<MsgChunk>>& msg,
        bool self_communication = true) const
    {
        unsigned msg_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 : msg) // first is PID, second is vector of chunks
        for( auto& chunk : chunks.second)
        {
            if( chunks.first == rank and not self_communication)
                continue;
            msg_size += chunk.size;
        }
        return msg_size;
    }
    // Size of the h_store
    unsigned store_size( bool self_communication = true) const
    {
        return msg_size( m_sendMsg, self_communication);
    }
    template<class ContainerType0, class ContainerType1>
    void mpi_global_gather_init( const ContainerType0& gatherFrom, ContainerType1& buffer,
        bool self_communication, int rank,
        const std::map<int,thrust::host_vector<MsgChunk>>& recvMsg,
        const std::map<int,thrust::host_vector<MsgChunk>>& sendMsg
        ) const
    {
        // TODO only works if recvMsg is non-overlapping
        // MW: not sure what exactly can go wrong ...
        //
        using value_type = dg::get_value_type<ContainerType0>;
        // Receives (we implicitly receive chunks in the order)
        unsigned start = 0;
        unsigned rqst_counter = 0;
        for( auto& msg : 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 = 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
        for( auto& msg : 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;
            MPI_Isend( send_ptr, chunk.size,
               getMPIDataType<value_type>(),  //sender
               msg.first, u, m_comm, &m_rqst[rqst_counter]);  //destination
            rqst_counter ++;
        }
    }
    template<class ContainerType0, class ContainerType1>
    void nccl_global_gather_init( const ContainerType0& gatherFrom, ContainerType1& buffer,
        bool self_communication, int rank,
        const std::map<int,thrust::host_vector<MsgChunk>>& recvMsg,
        const std::map<int,thrust::host_vector<MsgChunk>>& sendMsg
        ) const
    {
#ifdef DG_WITH_NCCL
        using value_type = dg::get_value_type<ContainerType0>;
        ncclComm_t ncclcomm;
        cudaStream_t stream;
        detail::getNcclComm( m_comm, &ncclcomm, &stream);
        unsigned mod_size = 1;
        if constexpr( std::is_same_v<dg::get_tensor_category<value_type>, dg::ComplexTag>)
            mod_size = 2;
        ncclGroupStart();
 
        unsigned start = 0;
        for( auto& msg : 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 = thrust::raw_pointer_cast( buffer.data()) + start;
            ncclRecv( recv_ptr, chunk.size*mod_size,
                detail::getNcclDataType<value_type>(), msg.first, ncclcomm, stream);
            start += chunk.size;
        }
        for( auto& msg : 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;
            ncclSend( send_ptr, chunk.size*mod_size,
                detail::getNcclDataType<value_type>(), msg.first, ncclcomm, stream);
        }
    // GroupEnd does not yet mean that all Kernels are submitted, just that grouping is done
        ncclGroupEnd();
#endif // DG_WITH_NCCL
    }
    template<class ContainerType0, class ContainerType1>
    void cuda_unaware_global_gather_init( const ContainerType0& gatherFrom, ContainerType1& buffer,
        bool self_communication, int rank,
        const std::map<int,thrust::host_vector<MsgChunk>>& recvMsg,
        const std::map<int,thrust::host_vector<MsgChunk>>& sendMsg
        ) const
    {
#ifdef __CUDACC__ // g++ does not know cuda code
        // if cuda unaware we need to send messages through host
        using value_type = dg::get_value_type<ContainerType0>;
        m_h_store.template set<value_type>( store_size( self_communication));
        // BugFix: buffer value_type must be set even if no messages are sent
        // so that global_gather_wait works
        m_h_buffer.template set<value_type>( buffer.size()); // global_gather_wait assigns same size
        auto& h_buffer = m_h_buffer.template get<value_type>();
        auto& h_store = m_h_store.template get<value_type>();
        //
        std::map<int,thrust::host_vector<MsgChunk>> h_sendMsg = sendMsg;
        unsigned start = 0;
        for( auto& msg : 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;
            thrust::copy( gatherFrom.begin()+chunk.idx,
                gatherFrom.begin()+chunk.idx+chunk.size, h_store.begin() + start);
 
            h_sendMsg[msg.first][u].idx = start;
            start += chunk.size;
        }
        mpi_global_gather_init( h_store, h_buffer, self_communication, rank,
            recvMsg, h_sendMsg);
#endif // __CUDACC__
 
    }
 
};
}//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 struct 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