Zoltan2_AlltoAll.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//   Zoltan2: A package of combinatorial algorithms for scientific computing
//
// Copyright 2012 NTESS and the Zoltan2 contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef _ZOLTAN2_ALLTOALL_HPP_
#define _ZOLTAN2_ALLTOALL_HPP_

#include <Zoltan2_Standards.hpp>
#include <Zoltan2_Environment.hpp>

#include <vector>
#include <climits>

namespace Zoltan2
{

extern void AlltoAllCount(const Comm<int> &comm, const Environment &env,
  const ArrayView<const int> &sendCount, const ArrayView<int> &recvCount);

template <typename T>
void AlltoAllv(const Comm<int> &comm,
              const Environment &env,
              const ArrayView<const T> &sendBuf,
              const ArrayView<const int> &sendCount,
              ArrayRCP<T> &recvBuf,      // output, allocated here
              const ArrayView<int> &recvCount   // output
)
{
  int nprocs = comm.getSize();
  int rank = comm.getRank();

  try{
    Zoltan2::AlltoAllCount(comm, env, sendCount, recvCount);
  }
  Z2_FORWARD_EXCEPTIONS;

  // Allocate the receive buffer.
  size_t totalrecv = 0;
  int maxMsg = 0;
  int nrecvranks = 0;
  for(int i = 0; i < nprocs; i++) {
    if (recvCount[i] > 0) {
      totalrecv += recvCount[i];
      nrecvranks++;
      if (recvCount[i] > maxMsg) maxMsg = recvCount[i];
    }
  }


  T *rbuf = NULL;
  if (totalrecv) rbuf = new T[totalrecv];

  if (nprocs > 1) {

    RCP<CommRequest<int> > *requests = new RCP<CommRequest<int> > [nrecvranks];

    // Error checking for memory and message size.
    int OK[2] = {1,1};
                // OK[0] -- true/false indicating whether each message size
                //          fits in an int (for MPI).
                // OK[1] -- true/false indicating whether memory allocs are OK
    int gOK[2]; // For global reduce of OK.

    if (size_t(maxMsg) * sizeof(T) > INT_MAX && nprocs > 1) OK[0] = false;
    if (totalrecv && !rbuf) OK[1] = 0;
    if (!requests) OK[1] = 0;

    // Post receives

    size_t offset = 0;
    size_t myrecvoffset = 0;
    size_t mysendoffset = 0;

    if (OK[0] && OK[1]) {
      int rcnt = 0;
      for (int i = 0; i < nprocs; i++) {
        if (i != rank && recvCount[i]) {
          try {
            requests[rcnt++] = Teuchos::ireceive<int,T>(comm,
                             Teuchos::arcp(&rbuf[offset],0,recvCount[i],false),
                             i);
          }
          Z2_THROW_OUTSIDE_ERROR(env);
        }
        else if (i == rank) {
          myrecvoffset = offset;
        }
        offset += recvCount[i];
      }
    }

    // Use barrier for error checking
    Teuchos::reduceAll<int>(comm, Teuchos::REDUCE_MIN, 2, OK, gOK);
    if (!gOK[0] || !gOK[1]) {
      delete [] rbuf;
      delete [] requests;
      if (!gOK[0])
        throw std::runtime_error("Max single message length exceeded");
      else
        throw std::bad_alloc();
    }

    // Send data; can use readySend since receives are posted.
    offset = 0;
    for (int i = 0; i < nprocs; i++) {
      if (i != rank && sendCount[i]) {
        try {
          Teuchos::readySend<int,T>(comm,
                            Teuchos::arrayView(&sendBuf[offset],sendCount[i]),
                            i);
        }
        Z2_THROW_OUTSIDE_ERROR(env);
      }
      else if (i == rank) {
        mysendoffset = offset;
      }
      offset += sendCount[i];
    }

    // Copy local data
    for (int j = 0; j < sendCount[rank]; j++)
      rbuf[myrecvoffset++] = sendBuf[mysendoffset++];

    // Wait for messages to return.
    try {
      Teuchos::waitAll<int>(comm, Teuchos::arrayView(requests, nrecvranks));
    }
    Z2_THROW_OUTSIDE_ERROR(env);

    delete [] requests;
  }
  else { // nprocs == 1; no communication needed

    if (totalrecv && !rbuf)
      throw std::bad_alloc();

    for (int j = 0; j < sendCount[0]; j++)
      rbuf[j] = sendBuf[j];
  }

  if (totalrecv)
    recvBuf = ArrayRCP<T>(rbuf, 0, totalrecv, true);
  else
    recvBuf = Teuchos::null;
}

/* \brief Specialization for std::string.

    For string of char. Number of chars in a string limited to SCHAR_MAX.
    Send as chars: 1 char for length of string, then chars in string,
     1 char for length of next string, and so on.
    \todo error checking
 */
template <>
void AlltoAllv(const Comm<int> &comm,
              const Environment &env,
              const ArrayView<const std::string> &sendBuf,
              const ArrayView<const int> &sendCount,
              ArrayRCP<std::string> &recvBuf,
              const ArrayView<int> &recvCount);

/* \brief Specialization for unsigned long long
 */
#ifdef HAVE_TPETRA_INT_LONG_LONG
template <>
void AlltoAllv(const Comm<int> &comm,
              const Environment &env,
              const ArrayView<const unsigned long long> &sendBuf,
              const ArrayView<const int> &sendCount,
              ArrayRCP<unsigned long long> &recvBuf,
              const ArrayView<int> &recvCount);
#endif

/* \brief Specialization for unsigned short
 */
template <>
void AlltoAllv(const Comm<int> &comm,
              const Environment &env,
              const ArrayView<const unsigned short> &sendBuf,
              const ArrayView<const int> &sendCount,
              ArrayRCP<unsigned short> &recvBuf,
              const ArrayView<int> &recvCount);

/* \brief For data type unsigned char (no Teuchos::DirectSerializationTraits)
 */
template <>
void AlltoAllv(const Comm<int> &comm,
              const Environment &env,
              const ArrayView<const unsigned char> &sendBuf,
              const ArrayView<const int> &sendCount,
              ArrayRCP<unsigned char> &recvBuf,
              const ArrayView<int> &recvCount);

}                   // namespace Z2
#endif