MinBMPI.hpp

#pragma once
#include <assert.h>   /* assert */
#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 */
#include <mpi.h>        /* MPI_* */


// Boost mpi stuff
#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;

namespace graybat {
    
    namespace communicationPolicy {
    
    /************************************************************************/
    struct MinBMPI {
        class Context {
        typedef unsigned ContextID;
        typedef unsigned VAddr;
        
        public:
        Context() :
            id(0),
            isValid(false){

        }

        Context(ContextID contextID, mpi::communicator comm) : 
            comm(comm),
            id(contextID),
            isValid(true){
        
        }

        Context& operator=(const Context& otherContext){
            id            = otherContext.getID();
            isValid       = otherContext.valid();
            comm          = otherContext.comm;
            return *this;

        }

        size_t size() const{
            return comm.size();
        }

        VAddr getVAddr() const {
            return comm.rank();
        }

        ContextID getID() const {
            return id;
        }

        bool valid() const{
            return isValid;
        }

        mpi::communicator comm;
    
        private:    
        ContextID id;
        bool      isValid;
        };

        class Event {
        public:
        Event(mpi::request request) : request(request){

        }

        ~Event(){

        }

        void wait(){
            request.wait();
    
        }

        bool ready(){
            boost::optional<mpi::status> status = request.test();

            if(status){
            return true;
            }
            else {
            return false;
            }

        }

        private:
        mpi::request request;
        };


        // Type defs
        typedef unsigned Tag;                                            
        typedef unsigned ContextID;
        typedef unsigned VAddr;

        typedef unsigned MsgType;
        typedef int      Uri;


        MinBMPI() :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
        ~MinBMPI(){
        
        }
        /***********************************************************************/
        // 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_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_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_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_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(const 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, const 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, const 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 createContext(const std::vector<VAddr> vAddrs, const Context oldContext){
            assert(vAddrs.size() > 0);
        
        VAddr myVAddr = oldContext.getVAddr();
        bool  isPartOfNewContext = false;
        
            for(VAddr vAddr : vAddrs){
            if(vAddr == myVAddr){
            isPartOfNewContext = true;
            }

            }

        mpi::communicator newComm = oldContext.comm.split(isPartOfNewContext);

            if(isPartOfNewContext){
                std::array<Uri, 1> uri;

                Context newContext(++contextCount, newComm);
            uri[0] = newContext.getVAddr();

                // Update UriMap
                uriMap.push_back(std::vector<Uri>(newContext.size()));
                std::vector<Uri> otherUris(newContext.size());

            allGather(newContext, uri, otherUris);

            std::copy(otherUris.begin(), otherUris.end(), uriMap[newContext.getID()].begin());

                return newContext;

            }
            else {
                // return invalid context
                // for peers not anymore included
                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