BMPI.hpp

#pragma once

// CLIB
#include <assert.h>   /* assert */

// STL
#include <array>     /* array */
#include <numeric>   /* std::accumulate */
#include <iostream>  /* std::cout */
#include <map>       /* std::map */
#include <exception> /* std::out_of_range */
#include <sstream>   /* std::stringstream */
#include <algorithm> /* std::transform */

// BOOST
#include <boost/mpi/environment.hpp>
#include <boost/mpi/communicator.hpp>
#include <boost/mpi/collectives.hpp>
#include <boost/mpi/datatype.hpp>
#include <boost/optional.hpp>
namespace mpi = boost::mpi;

// MPI
#include <mpi.h>     /* MPI_* */

// GRAYBAT
#include <graybat/utils/serialize_tuple.hpp>
#include <graybat/communicationPolicy/bmpi/Context.hpp> /* Context */
#include <graybat/communicationPolicy/bmpi/Event.hpp>   /* Event */
#include <graybat/communicationPolicy/bmpi/Config.hpp>   /* Config */
#include <graybat/communicationPolicy/Base.hpp> 
#include <graybat/communicationPolicy/Traits.hpp>

namespace graybat {
    
    namespace communicationPolicy {

    /************************************************************************/
        struct BMPI;
        
        namespace traits {

            template<>
            struct ContextType<BMPI> {
                using type = graybat::communicationPolicy::bmpi::Context;
            };

            template<>
            struct EventType<BMPI> {
                using type = graybat::communicationPolicy::bmpi::Event;
            };

            template<>
            struct ConfigType<BMPI> {
                using type = graybat::communicationPolicy::bmpi::Config;
            };

        }
        
    struct BMPI : Base<BMPI>{
            
            // Type defs
            using Tag       = typename graybat::communicationPolicy::Tag<BMPI>;
            using ContextID = typename graybat::communicationPolicy::ContextID<BMPI>;
            using MsgType   = typename graybat::communicationPolicy::MsgType<BMPI>;
            using VAddr     = typename graybat::communicationPolicy::VAddr<BMPI>;
            using Context   = typename graybat::communicationPolicy::Context<BMPI>;
            using Event     = typename graybat::communicationPolicy::Event<BMPI>;
            using Config    = typename graybat::communicationPolicy::Config<BMPI>;                        
            using Uri       = int;

        BMPI(Config const config) :contextCount(0),
            uriMap(0),
            initialContext(contextCount, mpi::communicator()){

        uriMap.push_back(std::vector<Uri>());
        
        for(unsigned i = 0; i < initialContext.size(); ++i){
            uriMap.back().push_back(i);
        }

        }

        // Destructor
        ~BMPI(){
        
        }
        /***********************************************************************/
        template <typename T_Send>
        void send(const VAddr destVAddr, const Tag tag, const Context context, const T_Send& sendData){
            Uri destUri = getVAddrUri(context, destVAddr);
        context.comm.send(destUri, tag, sendData.data(), sendData.size());
        }

            
        template <typename T_Send>
        Event asyncSend(const VAddr destVAddr, const Tag tag, const Context context, const T_Send& sendData){
        Uri destUri = getVAddrUri(context, destVAddr);
        mpi::request request = context.comm.isend(destUri, tag, sendData.data(), sendData.size());
            return Event(request);

        }
        template <typename T_Recv>
        void recv(const VAddr srcVAddr, const Tag tag, const Context context, T_Recv& recvData){
            Uri srcUri = getVAddrUri(context, srcVAddr);
        context.comm.recv(srcUri, tag, recvData.data(), recvData.size());

        }

            template <typename T_Recv>
        Event recv(const Context context, T_Recv& recvData){
                //std::cerr << mpi::any_source << " " << mpi::any_tag << std::endl;
                
                //auto status = context.comm.recv(mpi::any_source, mpi::any_tag, recvData.data(), recvData.size());
                auto status = context.comm.probe();
                context.comm.recv(status.source(), status.tag(), recvData.data(), recvData.size());
                return Event(status);
                

                //auto status = context.comm.recv(boost::mpi::any_source, boost::mpi::any_tag, recvData.data(), recvData.size());
                //auto status = context.comm.recv(boost::mpi::any_source, boost::mpi::any_tag);

        }


        template <typename T_Recv>
        Event asyncRecv(const VAddr srcVAddr, const Tag tag, const Context context, T_Recv& recvData){
             Uri srcUri = getVAddrUri(context, srcVAddr);
         mpi::request request = context.comm.irecv(srcUri, tag, recvData.data(), recvData.size());
             return Event(request);
        }


        /************************************************************************/
        template <typename T_Send, typename T_Recv>
        void gather(const VAddr rootVAddr, const Context context, const T_Send& sendData, T_Recv& recvData){
            Uri rootUri = getVAddrUri(context, rootVAddr);
            mpi::gather(context.comm, sendData.data(), sendData.size(), recvData, rootUri);
        }

    
        template <typename T_Send, typename T_Recv>
        void gatherVar(const VAddr rootVAddr, const Context context, const T_Send& sendData, T_Recv& recvData, std::vector<unsigned>& recvCount){
        // Retrieve number of elements each peer sends
        recvCount.resize(context.size());
        std::array<unsigned, 1> nElements{{(unsigned)sendData.size()}};
        allGather(context, nElements, recvCount);
        recvData.resize(std::accumulate(recvCount.begin(), recvCount.end(), 0U));

        Uri rootUri = getVAddrUri(context, rootVAddr);
        int rdispls[context.size()];

        // Create offset map 
        unsigned offset  = 0;
        for (unsigned i=0; i < context.size(); ++i) { 
            rdispls[i] = offset; 
            offset += recvCount[i];
        
        }
        
        // Gather data with varying size
        MPI_Gatherv(const_cast<typename T_Send::value_type*>(sendData.data()), sendData.size(),
                mpi::get_mpi_datatype<typename T_Send::value_type>(*(sendData.data())),
                const_cast<typename T_Recv::value_type*>(recvData.data()),
                const_cast<int*>((int*)recvCount.data()), rdispls,
                mpi::get_mpi_datatype<typename T_Recv::value_type>(*(recvData.data())), 
                rootUri, context.comm);
        
        
        }

    
        template <typename T_Send, typename T_Recv>
        void allGather(Context context, const T_Send& sendData, T_Recv& recvData){
        mpi::all_gather(context.comm, sendData.data(), sendData.size(), recvData.data());
        
        }

    
         template <typename T_Send, typename T_Recv>
         void allGatherVar(const Context context, const T_Send& sendData, T_Recv& recvData, std::vector<unsigned>& recvCount){
             // Retrieve number of elements each peer sends
             recvCount.resize(context.size());
             allGather(context, std::array<unsigned, 1>{{(unsigned)sendData.size()}}, recvCount);
             recvData.resize(std::accumulate(recvCount.begin(), recvCount.end(), 0U));

         int rdispls[context.size()];

         // Create offset map
         unsigned offset  = 0;
         for (unsigned i=0; i < context.size(); ++i) { 
             rdispls[i] = offset; 
             offset += recvCount[i];
        
         }
        
         // Gather data with varying size
         MPI_Allgatherv(const_cast<typename T_Send::value_type*>(sendData.data()), sendData.size(),
                mpi::get_mpi_datatype<typename T_Send::value_type>(*(sendData.data())),
                const_cast<typename T_Recv::value_type*>(recvData.data()),
                const_cast<int*>((int*)recvCount.data()), rdispls,
                mpi::get_mpi_datatype<typename T_Recv::value_type>(*(recvData.data())), 
                context.comm);
         
         
         }


        template <typename T_Send, typename T_Recv>
        void scatter(const VAddr rootVAddr, const Context context, const T_Send& sendData, T_Recv& recvData){
             Uri rootUri = getVAddrUri(context, rootVAddr);
         mpi::scatter(context.comm, sendData.data(), recvData.data(), recvData.size(), rootUri);

        }

    
        template <typename T_Send, typename T_Recv>
        void allToAll(const Context context, const T_Send& sendData, T_Recv& recvData){
            unsigned elementsPerPeer = sendData.size() / context.size();

        mpi::all_to_all(context.comm, sendData.data(), elementsPerPeer, recvData.data());

        }

    
        template <typename T_Send, typename T_Recv, typename T_Op>
        void reduce(const VAddr rootVAddr, const Context context, const T_Op op, const T_Send& sendData, T_Recv& recvData){
             Uri rootUri = getVAddrUri(context, rootVAddr);
         mpi::reduce(context.comm, sendData.data(), sendData.size(), recvData.data(), op, rootUri);

        }

        template <typename T_Send, typename T_Recv, typename T_Op>
        void allReduce(const Context context, T_Op op, const T_Send& sendData, T_Recv& recvData){
        mpi::all_reduce(context.comm, sendData.data(), sendData.size(), recvData.data(), op);
         
        }

    
        template <typename T_SendRecv>
        void broadcast(const VAddr rootVAddr, const Context context, T_SendRecv& data){
             Uri rootUri = uriMap.at(context.getID()).at(rootVAddr);
         mpi::broadcast(context.comm, data.data(), data.size(), rootUri);
        }

    
         void synchronize(const Context context){
         context.comm.barrier();
         }

    
         void synchronize(){
             synchronize(getGlobalContext());
         }
        /*************************************************************************/
            Context splitContext(const bool isMember, const Context oldContext){
                mpi::communicator newComm = oldContext.comm.split(isMember);

                if(isMember){
                Context newContext(++contextCount, newComm);
                    std::array<Uri, 1> uri {{ (int) newContext.getVAddr() }};
                    uriMap.push_back(std::vector<Uri>(newContext.size()));

                    std::vector<Event> events;
                    
                    for(unsigned i = 0; i < newContext.size(); ++i){
                        events.push_back(Event(newContext.comm.isend(i, 0, uri.data(), 1)));
                        
                    }

                    for(unsigned i = 0; i < newContext.size(); ++i){
                        std::array<Uri, 1> otherUri {{ 0 }};
                        newContext.comm.recv(i, 0, otherUri.data(), 1);
                        uriMap.at(newContext.getID()).at(i) = otherUri[0];
                        
                    }

                    for(unsigned i = 0; i < events.size(); ++i){
                        events.back().wait();
                        events.pop_back();
                    }

                    return newContext;

                    
                }
                else {
                    return Context();
                }

                
            }

    
        Context getGlobalContext(){
            return initialContext;
        }
            private:

        /***************************************************************************
         *
         * @name Private Member
         *
         ***************************************************************************/
        ContextID                      contextCount;
        std::vector<std::vector<Uri>>  uriMap;
        Context                        initialContext;
        mpi::environment               env;

        /***************************************************************************
         *
         * @name Helper Functions
         *
         ***************************************************************************/
        
        
        void error(VAddr vAddr, std::string msg){
        std::cout << "[" << vAddr << "] " << msg;

        }

        template <typename T_Context>
        inline Uri getVAddrUri(const T_Context context, const VAddr vAddr){
            Uri uri  = 0;
            try {
                uri = uriMap.at(context.getID()).at(vAddr);

            } catch(const std::out_of_range& e){
                std::stringstream errorStream;
                errorStream << "MPI::getVAddrUri::" << e.what()<< " : Communicator with ID " << vAddr << " is not part of the context " << context.getID() << std::endl;
                error(context.getID(), errorStream.str());
                exit(1);
            }

            return uri;
        }

    };

    } // namespace communicationPolicy
    
} // namespace graybat