#include <mpi.h>

#include <iostream>
#include <sstream>
#include <memory>

// macro to get consistent output from MPI ranks
#define PRINT(STR,EXPR) {std::ostringstream oss; oss<<EXPR; (STR)<<oss.str()<<std::flush;}

// simplest case that fails:
// wait for all MPI requestw at once
// If not defined, a while-loop with
// MPI_Waitany() calls is used, which
// gives more insight into what's going on.
#define WAITALL

// this demo program transposes a nproc x npproc matrix,
// where nproc is the number of MPI processes. Initially,
// MPI rank r has row r of the matrix. Then nproc
// MPI_Igathers are dispatched to do the transpose, so that
// afterwards rank r has column r of the original matrix.
//
// The purpose of the code is to demonstrate that dispatching
// the Igather calls from multiple threads leads to a deadlock
// (with OpenMPI 4.0.3 at least). It looks like the requests
// are associated with different communication tasks on different processes.
int main(int argc, char* argv[])
{
  int provided;
  MPI_Init_thread(&argc,&argv, MPI_THREAD_MULTIPLE, &provided);
  MPI_Comm comm=MPI_COMM_WORLD;
  if (provided!=MPI_THREAD_MULTIPLE)
  {
    std::cerr << "Your MPI installation does not provide threading level MPI_THREAD_MULTIPLE, required for this program." << std::endl;
    MPI_Abort(MPI_COMM_WORLD, -1);
  }
  int rank, nproc;
  MPI_Comm_rank(comm,&rank);
  MPI_Comm_size(comm,&nproc);

  std::unique_ptr<double[]> M1(new double[nproc]);
  std::unique_ptr<double[]> M2(new double[nproc*nproc]);

  for (int j=0; j<nproc; j++)
    M1[j] = rank*nproc+j+1;

  for (int r=0; r<nproc; r++)
  {
    if (rank==r)
    {
      for (int j=0; j<nproc; j++)
      {
        std::cout << M1[j] << " ";
      }
      std::cout << std::endl;
      std::cout << std::flush;
    }
    MPI_Barrier(comm);
  }

MPI_Request* requests=new MPI_Request[nproc];

#pragma omp parallel for schedule(static) shared(requests)
for (int j=0; j<nproc; j++)
{
  int target_proc= 0;

  MPI_Igather(M1.get()+j,  1,  MPI_DOUBLE,
              M2.get()+j*nproc, 1, MPI_DOUBLE, target_proc,
              comm, &requests[j]);
}

PRINT(std::cerr,"P"<<rank<<" all igathers dispatched."<<std::endl);

// Variant 1: wait for all requests.
// This hits a deadlock if OMP_NUM_THREADS>1 for me.
#ifdef WAITALL

MPI_Waitall(nproc,requests,MPI_STATUSES_IGNORE);

#else

// instead we want to continue working as soon as the individual
// columns are gathered:
while(true)
{
  int index=0;
  MPI_Waitany(nproc,requests,&index,MPI_STATUS_IGNORE);
  if (index==MPI_UNDEFINED)
  {
    PRINT(std::cerr,"P"<<rank<<": no more work"<<std::endl);
    break;
  }
  int target_proc = index;
  PRINT(std::cerr, "P"<<rank<<": request "<<index<<" done"<<std::endl);
}
PRINT(std::cerr,"P"<<rank<<" after Waitany()'s\n");

#endif

  for (int r=0; r<nproc; r++)
  {
    if (rank==0)
    {
      for (int j=0; j<nproc; j++)
      {
        std::cout << M2[r*nproc+j] << " ";
      }
      std::cout << std::endl;
      std::cout << std::flush;
    }
    MPI_Barrier(MPI_COMM_WORLD);
  }
  MPI_Finalize();
  return 0;
}