fei_CommUtils.hpp

/*--------------------------------------------------------------------*/
/*    Copyright 2007 Sandia Corporation.                              */
/*    Under the terms of Contract DE-AC04-94AL85000, there is a       */
/*    non-exclusive license for use of this work by or on behalf      */
/*    of the U.S. Government.  Export of this program may require     */
/*    a license from the United States Government.                    */
/*--------------------------------------------------------------------*/

#ifndef _fei_CommUtils_hpp_
#define _fei_CommUtils_hpp_

#include <fei_macros.hpp>
#include <fei_mpi.h>
#include <fei_mpiTraits.hpp>
#include <fei_chk_mpi.hpp>
#include <fei_iostream.hpp>
#include <fei_CommMap.hpp>
#include <fei_TemplateUtils.hpp>
#include <snl_fei_RaggedTable.hpp>

#include <vector>
#include <set>
#include <map>

#include <fei_ErrMacros.hpp>
#undef fei_file
#define fei_file "fei_CommUtils.hpp"

namespace fei {

int localProc(MPI_Comm comm);

int numProcs(MPI_Comm comm);

void Barrier(MPI_Comm comm);

int mirrorProcs(MPI_Comm comm, std::vector<int>& toProcs, std::vector<int>& fromProcs);

typedef snl_fei::RaggedTable<std::map<int,std::set<int>*>,std::set<int> > comm_map;

int mirrorCommPattern(MPI_Comm comm, comm_map* inPattern, comm_map*& outPattern);

int Allreduce(MPI_Comm comm, bool localBool, bool& globalBool);

int exchangeIntData(MPI_Comm comm,
                    const std::vector<int>& sendProcs,
                    std::vector<int>& sendData,
                    const std::vector<int>& recvProcs,
                    std::vector<int>& recvData);
 
//----------------------------------------------------------------------------
template<class T>
int GlobalMax(MPI_Comm comm, std::vector<T>& local, std::vector<T>& global)
{
#ifdef FEI_SER
  global = local;
#else

  MPI_Datatype mpi_dtype = fei::mpiTraits<T>::mpi_type();

  try {
    global.resize(local.size());
  }
  catch(std::runtime_error& exc) {
    fei::console_out() << exc.what()<<FEI_ENDL;
    return(-1);
  }

  CHK_MPI( MPI_Allreduce(&(local[0]), &(global[0]),
       local.size(), mpi_dtype, MPI_MAX, comm) );
#endif

  return(0);
}

//----------------------------------------------------------------------------
template<class T>
int GlobalMax(MPI_Comm comm, T local, T& global)
{
#ifdef FEI_SER
  global = local;
#else
  MPI_Datatype mpi_dtype = fei::mpiTraits<T>::mpi_type();

  CHK_MPI( MPI_Allreduce(&local, &global, 1, mpi_dtype, MPI_MAX, comm) );
#endif
  return(0);
}

//----------------------------------------------------------------------------
template<class T>
int GlobalMin(MPI_Comm comm, std::vector<T>& local, std::vector<T>& global)
{
#ifdef FEI_SER
  global = local;
#else

  MPI_Datatype mpi_dtype = fei::mpiTraits<T>::mpi_type();

  try {
    global.resize(local.size());
  }
  catch(std::runtime_error& exc) {
    fei::console_out() << exc.what()<<FEI_ENDL;
    return(-1);
  }

  CHK_MPI( MPI_Allreduce(&(local[0]), &(global[0]),
       local.size(), mpi_dtype, MPI_MIN, comm) );
#endif

  return(0);
}

//----------------------------------------------------------------------------
template<class T>
int GlobalMin(MPI_Comm comm, T local, T& global)
{
#ifdef FEI_SER
  global = local;
#else
  MPI_Datatype mpi_dtype = fei::mpiTraits<T>::mpi_type();

  CHK_MPI( MPI_Allreduce(&local, &global, 1, mpi_dtype, MPI_MIN, comm) );
#endif
  return(0);
}

//----------------------------------------------------------------------------
template<class T>
int GlobalSum(MPI_Comm comm, std::vector<T>& local, std::vector<T>& global)
{
#ifdef FEI_SER
  global = local;
#else
  global.resize(local.size());

  MPI_Datatype mpi_dtype = fei::mpiTraits<T>::mpi_type();

  CHK_MPI( MPI_Allreduce(&(local[0]), &(global[0]),
                      local.size(), mpi_dtype, MPI_SUM, comm) );
#endif
  return(0);
}

//----------------------------------------------------------------------------
template<class T>
int GlobalSum(MPI_Comm comm, T local, T& global)
{
#ifdef FEI_SER
  global = local;
#else
  MPI_Datatype mpi_dtype = fei::mpiTraits<T>::mpi_type();

  CHK_MPI( MPI_Allreduce(&local, &global, 1, mpi_dtype, MPI_SUM, comm) );
#endif
  return(0);
}


template<class T>
int Allgatherv(MPI_Comm comm,
               std::vector<T>& sendbuf,
               std::vector<int>& recvLengths,
               std::vector<T>& recvbuf)
{
#ifdef FEI_SER
  //If we're in serial mode, just copy sendbuf to recvbuf and return.

  recvbuf = sendbuf;
  recvLengths.resize(1);
  recvLengths[0] = sendbuf.size();
#else
  int numProcs = 1;
  MPI_Comm_size(comm, &numProcs);

  try {

  MPI_Datatype mpi_dtype = fei::mpiTraits<T>::mpi_type();

  std::vector<int> tmpInt(numProcs, 0);

  int len = sendbuf.size();
  int* tmpBuf = &tmpInt[0];

  recvLengths.resize(numProcs);
  int* recvLenPtr = &recvLengths[0];

  CHK_MPI( MPI_Allgather(&len, 1, MPI_INT, recvLenPtr, 1, MPI_INT, comm) );

  int displ = 0;
  for(int i=0; i<numProcs; i++) {
    tmpBuf[i] = displ;
    displ += recvLenPtr[i];
  }

  if (displ == 0) {
    recvbuf.resize(0);
    return(0);
  }

  recvbuf.resize(displ);

  T* sendbufPtr = sendbuf.size()>0 ? &sendbuf[0] : NULL;
  
  CHK_MPI( MPI_Allgatherv(sendbufPtr, len, mpi_dtype,
      &recvbuf[0], &recvLengths[0], tmpBuf,
      mpi_dtype, comm) );

  }
  catch(std::runtime_error& exc) {
    fei::console_out() << exc.what() << FEI_ENDL;
    return(-1);
  }
#endif

  return(0);
}

//------------------------------------------------------------------------
template<class T>
int Bcast(MPI_Comm comm, std::vector<T>& sendbuf, int sourceProc)
{
#ifndef FEI_SER
  MPI_Datatype mpi_dtype = fei::mpiTraits<T>::mpi_type();

  CHK_MPI(MPI_Bcast(&sendbuf[0], sendbuf.size(), mpi_dtype,
                    sourceProc, comm) );
#endif
  return(0);
}

//------------------------------------------------------------------------
template<typename T>
int exchangeCommMapData(MPI_Comm comm,
                        const typename CommMap<T>::Type& sendCommMap,
                        typename CommMap<T>::Type& recvCommMap,
                        bool recvProcsKnownOnEntry = false,
                        bool recvLengthsKnownOnEntry = false)
{
  if (!recvProcsKnownOnEntry) {
    recvCommMap.clear();
  }

#ifndef FEI_SER
  int tag = 11120;
  MPI_Datatype mpi_dtype = fei::mpiTraits<T>::mpi_type();

  std::vector<int> sendProcs;
  fei::copyKeysToVector(sendCommMap, sendProcs);
  std::vector<int> recvProcs;

  if (recvProcsKnownOnEntry) {
    fei::copyKeysToVector(recvCommMap, recvProcs);
  }
  else {
    mirrorProcs(comm, sendProcs, recvProcs);
    for(size_t i=0; i<recvProcs.size(); ++i) {
      addItemsToCommMap<T>(recvProcs[i], 0, NULL, recvCommMap);
    }
  }

  if (!recvLengthsKnownOnEntry) {
    std::vector<int> tmpIntData(sendProcs.size());
    std::vector<int> recvLengths(recvProcs.size());
    
    typename fei::CommMap<T>::Type::const_iterator
      s_iter = sendCommMap.begin(), s_end = sendCommMap.end();

    for(size_t i=0; s_iter != s_end; ++s_iter, ++i) {
      tmpIntData[i] = s_iter->second.size();
    }

    if ( exchangeIntData(comm, sendProcs, tmpIntData, recvProcs, recvLengths) != 0) {
      return(-1);
    }
    for(size_t i=0; i<recvProcs.size(); ++i) {
      std::vector<T>& rdata = recvCommMap[recvProcs[i]];
      rdata.resize(recvLengths[i]);
    }
  }

  //launch Irecv's for recv-data:
  std::vector<MPI_Request> mpiReqs;
  mpiReqs.resize(recvProcs.size());

  typename fei::CommMap<T>::Type::iterator
    r_iter = recvCommMap.begin(), r_end = recvCommMap.end();

  size_t req_offset = 0;
  for(; r_iter != r_end; ++r_iter) {
    int rproc = r_iter->first;
    std::vector<T>& recv_vec = r_iter->second;
    int len = recv_vec.size();
    T* recv_buf = len > 0 ? &recv_vec[0] : NULL;

    CHK_MPI( MPI_Irecv(recv_buf, len, mpi_dtype, rproc,
                       tag, comm, &mpiReqs[req_offset++]) );
  }

  //send the send-data:

  typename fei::CommMap<T>::Type::const_iterator
    s_iter = sendCommMap.begin(), s_end = sendCommMap.end();

  for(; s_iter != s_end; ++s_iter) {
    int sproc = s_iter->first;
    const std::vector<T>& send_vec = s_iter->second;
    int len = send_vec.size();
    T* send_buf = len>0 ? const_cast<T*>(&send_vec[0]) : NULL;

    CHK_MPI( MPI_Send(send_buf, len, mpi_dtype, sproc, tag, comm) );
  }

  //complete the Irecvs:
  for(size_t i=0; i<mpiReqs.size(); ++i) {
    int index;
    MPI_Status status;
    CHK_MPI( MPI_Waitany(mpiReqs.size(), &mpiReqs[0], &index, &status) );
  }

#endif
  return(0);
}


//------------------------------------------------------------------------
template<class T>
int exchangeData(MPI_Comm comm,
                 std::vector<int>& sendProcs,
                 std::vector<std::vector<T> >& sendData,
                 std::vector<int>& recvProcs,
                 bool recvDataLengthsKnownOnEntry,
                 std::vector<std::vector<T> >& recvData)
{
  if (sendProcs.size() == 0 && recvProcs.size() == 0) return(0);
  if (sendProcs.size() != sendData.size()) return(-1);
#ifndef FEI_SER
  std::vector<MPI_Request> mpiReqs;
  mpiReqs.resize(recvProcs.size());

  int tag = 11119;
  MPI_Datatype mpi_dtype = fei::mpiTraits<T>::mpi_type();

  if (!recvDataLengthsKnownOnEntry) {
    std::vector<int> tmpIntData(sendData.size());
    std::vector<int> recvLengths(recvProcs.size());
    for(unsigned i=0; i<sendData.size(); ++i) {
      tmpIntData[i] = sendData[i].size();
    }

    if ( exchangeIntData(comm, sendProcs, tmpIntData, recvProcs, recvLengths) != 0) {
      return(-1);
    }
    for(unsigned i=0; i<recvProcs.size(); ++i) {
      recvData[i].resize(recvLengths[i]);
    }
  }

  //launch Irecv's for recvData:

  size_t numRecvProcs = recvProcs.size();
  int req_offset = 0;
  int localProc = fei::localProc(comm);
  for(size_t i=0; i<recvProcs.size(); ++i) {
    if (recvProcs[i] == localProc) {--numRecvProcs; continue; }

    int len = recvData[i].size();
    std::vector<T>& recv_vec = recvData[i];
    T* recv_buf = len > 0 ? &recv_vec[0] : NULL;

    CHK_MPI( MPI_Irecv(recv_buf, len, mpi_dtype, recvProcs[i],
                       tag, comm, &mpiReqs[req_offset++]) );
  }

  //send the sendData:

  for(size_t i=0; i<sendProcs.size(); ++i) {
    if (sendProcs[i] == localProc) continue;

    std::vector<T>& send_buf = sendData[i];
    CHK_MPI( MPI_Send(&send_buf[0], sendData[i].size(), mpi_dtype,
                      sendProcs[i], tag, comm) );
  }

  //complete the Irecvs:
  for(size_t i=0; i<numRecvProcs; ++i) {
    if (recvProcs[i] == localProc) continue;
    int index;
    MPI_Status status;
    CHK_MPI( MPI_Waitany(numRecvProcs, &mpiReqs[0], &index, &status) );
  }

#endif
  return(0);
}

//------------------------------------------------------------------------
template<class T>
int exchangeData(MPI_Comm comm,
                 std::vector<int>& sendProcs,
                 std::vector<std::vector<T>*>& sendData,
                 std::vector<int>& recvProcs,
                 bool recvLengthsKnownOnEntry,
                 std::vector<std::vector<T>*>& recvData)
{
  if (sendProcs.size() == 0 && recvProcs.size() == 0) return(0);
  if (sendProcs.size() != sendData.size()) return(-1);
#ifndef FEI_SER
  int tag = 11115;
  MPI_Datatype mpi_dtype = fei::mpiTraits<T>::mpi_type();
  std::vector<MPI_Request> mpiReqs;

  try {
  mpiReqs.resize(recvProcs.size());

  if (!recvLengthsKnownOnEntry) {
    std::vector<int> tmpIntData;
    tmpIntData.resize(sendData.size());
    std::vector<int> recvLens(sendData.size());
    for(unsigned i=0; i<sendData.size(); ++i) {
      tmpIntData[i] = (int)sendData[i]->size();
    }

    if (exchangeIntData(comm, sendProcs, tmpIntData, recvProcs, recvLens) != 0) {
      return(-1);
    }

    for(unsigned i=0; i<recvLens.size(); ++i) {
      recvData[i]->resize(recvLens[i]);
    }
  }
  }
  catch(std::runtime_error& exc) {
    fei::console_out() << exc.what() << FEI_ENDL;
    return(-1);
  }

  //launch Irecv's for recvData:

  size_t numRecvProcs = recvProcs.size();
  int req_offset = 0;
  int localProc = fei::localProc(comm);
  for(unsigned i=0; i<recvProcs.size(); ++i) {
    if (recvProcs[i] == localProc) {--numRecvProcs; continue;}

    size_t len = recvData[i]->size();
    std::vector<T>& rbuf = *recvData[i];

    CHK_MPI( MPI_Irecv(&rbuf[0], (int)len, mpi_dtype,
                       recvProcs[i], tag, comm, &mpiReqs[req_offset++]) );
  }

  //send the sendData:

  for(unsigned i=0; i<sendProcs.size(); ++i) {
    if (sendProcs[i] == localProc) continue;

    std::vector<T>& sbuf = *sendData[i];
    CHK_MPI( MPI_Send(&sbuf[0], (int)sbuf.size(), mpi_dtype,
                      sendProcs[i], tag, comm) );
  }

  //complete the Irecv's:
  for(unsigned i=0; i<numRecvProcs; ++i) {
    if (recvProcs[i] == localProc) continue;
    int index;
    MPI_Status status;
    CHK_MPI( MPI_Waitany((int)numRecvProcs, &mpiReqs[0], &index, &status) );
  }

#endif
  return(0);
}


//------------------------------------------------------------------------
template<class T>
class MessageHandler {
public:
  virtual ~MessageHandler(){}

  virtual std::vector<int>& getSendProcs() = 0;

  virtual std::vector<int>& getRecvProcs() = 0;

  virtual int getSendMessageLength(int destProc, int& messageLength) = 0;

  virtual int getSendMessage(int destProc, std::vector<T>& message) = 0;

  virtual int processRecvMessage(int srcProc, std::vector<T>& message) = 0;
};//class MessageHandler


//------------------------------------------------------------------------
template<class T>
int exchange(MPI_Comm comm, MessageHandler<T>* msgHandler)
{
#ifdef FEI_SER
  (void)msgHandler;
#else
  int numProcs = fei::numProcs(comm);
  if (numProcs < 2) return(0);

  std::vector<int>& sendProcs = msgHandler->getSendProcs();
  int numSendProcs = sendProcs.size();
  std::vector<int>& recvProcs = msgHandler->getRecvProcs();
  int numRecvProcs = recvProcs.size();
  int i;

  if (numSendProcs < 1 && numRecvProcs < 1) {
    return(0);
  }

  std::vector<int> sendMsgLengths(numSendProcs);

  for(i=0; i<numSendProcs; ++i) {
    CHK_ERR( msgHandler->getSendMessageLength(sendProcs[i], sendMsgLengths[i]) );
  }

  std::vector<std::vector<T> > recvMsgs(numRecvProcs);

  std::vector<std::vector<T> > sendMsgs(numSendProcs);
  for(i=0; i<numSendProcs; ++i) {
    CHK_ERR( msgHandler->getSendMessage(sendProcs[i], sendMsgs[i]) );
  }

  CHK_ERR( exchangeData(comm, sendProcs, sendMsgs,
                        recvProcs, false, recvMsgs) );

  for(i=0; i<numRecvProcs; ++i) {
    std::vector<T>& recvMsg = recvMsgs[i];
    CHK_ERR( msgHandler->processRecvMessage(recvProcs[i], recvMsg ) );
  }
#endif

  return(0);
}


} //namespace fei

#endif // _fei_CommUtils_hpp_