AlltoAll.cpp

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
 
// TODO: doxygen comments
//     make this a real unit test that gives helpful information if it fails
//     and uses different template values
 
#include <Zoltan2_Environment.hpp>   
#include <Zoltan2_AlltoAll.hpp>   
#include <Zoltan2_TestHelpers.hpp>   
 
#include <iostream>
#include <algorithm>
#include <vector>
#include <string>
 
#include <Teuchos_RCP.hpp>   
#include <Teuchos_ArrayRCP.hpp>   
#include <Teuchos_Comm.hpp>   
#include <Teuchos_DefaultComm.hpp>   
 
 
int main(int narg, char *arg[])
{
  Tpetra::ScopeGuard tscope(&narg, &arg);
  Teuchos::RCP<const Teuchos::Comm<int> > comm = Tpetra::getDefaultComm();
 
  int rank = comm->getRank();
  int nprocs = comm->getSize();
 
  Teuchos::RCP<const Zoltan2::Environment> envPtr = 
    Teuchos::rcp(new Zoltan2::Environment(comm));
 
  int errcode = 0;
 
  if (!errcode){
 
    // test of Zoltan2::AlltoAllv (different sized messages) using a Scalar type
 
    int myMsgSizeBase=rank*nprocs + 1;
    Array<int> sendCount(nprocs, 0);
    Array<int> recvCount(nprocs, 0);
    long totalOut = 0;
 
    for (int p=0; p < nprocs; p++){
      sendCount[p] = myMsgSizeBase + p;
      totalOut += sendCount[p];
    }
  
    Array<int> sendBuf(totalOut, 0);
  
    int *out = &(sendBuf[0]);
    for (int p=0; p < nprocs; p++){
      for (int i=0; i < sendCount[p]; i++){
        *out++ = p+rank;
      }
    }
  
    Teuchos::ArrayRCP<int> recvBuf;
    
    Zoltan2::AlltoAllv<int>(*comm, *envPtr,
                  sendBuf(),
                  sendCount(),
                  recvBuf,
                  recvCount());
 
    int *inBuf = recvBuf.get();
 
    for (int p=0; p < nprocs; p++){
      for (int i=0; i < recvCount[p]; i++){
        if (*inBuf++ != rank+p){
          errcode = 4;
          break;
        }
      }
    }
  }
 
  TEST_FAIL_AND_EXIT(*comm, errcode==0, "int", errcode);
 
  if (!errcode){
 
    // test of Zoltan2::AlltoAllv using strings - which can not
    //    be serialized by Teuchos.
    //  Rank p sends p messages to each process.
 
    int nstrings = nprocs * rank;
    string *sendStrings = NULL;
 
    if (nstrings > 0)
      sendStrings = new string [nstrings];
 
    std::ostringstream myMessage;
    myMessage << "message from process " << rank;
 
    for (int i=0; i < nstrings; i++)
      sendStrings[i] = myMessage.str();
 
    int *counts = new int [nprocs];
    for (int i=0; i < nprocs ; i++)
      counts[i] = rank;
 
    Teuchos::ArrayView<const string> sendBuf(sendStrings, nstrings);
    Teuchos::ArrayView<const int> sendCount(counts, nprocs);
    Teuchos::Array<int> recvCounts(nprocs, 0);
 
    Teuchos::ArrayRCP<string> recvBuf;
 
    Zoltan2::AlltoAllv<string>(*comm, *envPtr,
        sendBuf,    
        sendCount,   
        recvBuf,
        recvCounts());
 
    delete [] sendStrings;
    delete [] counts;
  
    int next = 0;
    for (int i=0; i < nprocs; i++){
      if (recvCounts[i] != i){
        errcode = 5;
        break;
      }
      std::ostringstream msg;
      msg << "message from process " << i;
      for (int j=0; j < recvCounts[i]; j++){
        if (recvBuf[next++] != msg.str()){
          errcode = 6;
          break;
        }
      }
    }
  }
 
  TEST_FAIL_AND_EXIT(*comm, errcode==0, "strings", errcode);
 
  if (rank == 0){
    if (errcode)
      std::cout << "FAIL" << std::endl;
    else
      std::cout << "PASS" << std::endl;
  }
 
  return errcode;
}