Fix bug where Python interpreter exit status is consumed by ORBIT_MPI_Finalize hook. Add ORBIT_MPI_Gather/ORBIT_MPI_Allgather

src/mpi/orbit_mpi.cc

@@ -28,14 +28,31 @@ static std::size_t ORBIT_MPI_Type_size(MPI_Datatype data) {
 }
 #endif
 
+#if USE_MPI > 0
+// MPI_Finalize hook for Py_AtExit to propagate error codes.
+static void ORBIT_MPI_Finalize_AtExit() {
+  int initialized = 0;
+  if (MPI_Initialized(&initialized) != MPI_SUCCESS || !initialized) {
+    return;
+  }
+
+  int finalized = 0;
+  if (MPI_Finalized(&finalized) != MPI_SUCCESS || finalized) {
+    return;
+  }
+
+  MPI_Finalize();
+}
+#endif
+
 /** A C wrapper around MPI_Init. */
 int ORBIT_MPI_Init(){
 #if USE_MPI > 0
   // Ignoring result; if it fails, the proc is doomed anyway.
   MPI_Init(NULL, NULL);
 
   // Registering MPI finalize method at cleanup stage
-  Py_AtExit(ORBIT_MPI_Finalize);
+  Py_AtExit(ORBIT_MPI_Finalize_AtExit);
 #endif
   return MPI_SUCCESS;
 }
@@ -561,6 +578,39 @@ int ORBIT_MPI_Allreduce(void* ar1, void* ar2, int n, MPI_Datatype data, MPI_Op o
 #endif
 }
 
+int ORBIT_MPI_Gather(const void *sendbuf, int sendcount, MPI_Datatype sendtype,
+                     void *recvbuf, int recvcount, MPI_Datatype recvtype, int root, MPI_Comm comm) {
+#if USE_MPI > 0
+  return MPI_Gather(sendbuf, sendcount, sendtype, recvbuf, recvcount, recvtype, root, comm);
+#else
+  if (sendbuf == ORBIT_MPI_IN_PLACE || sendbuf == recvbuf) {
+    return MPI_SUCCESS;
+  }
+
+  if (sendcount > 0) {
+    std::memcpy(recvbuf, sendbuf, static_cast<std::size_t>(sendcount) * ORBIT_MPI_Type_size(sendtype));
+  }
+
+  return MPI_SUCCESS;
+#endif
+}
+
+int ORBIT_MPI_Allgather(const void *sendbuf, int sendcount, MPI_Datatype sendtype,
+                         void *recvbuf, int recvcount, MPI_Datatype recvtype, MPI_Comm comm) {
+#if USE_MPI > 0
+    return MPI_Allgather(sendbuf, sendcount, sendtype, recvbuf, recvcount, recvtype, comm);
+#else
+  if (sendbuf == ORBIT_MPI_IN_PLACE || sendbuf == recvbuf) {
+    return MPI_SUCCESS;
+  }
+
+  if (sendcount > 0) {
+    std::memcpy(recvbuf, sendbuf, static_cast<std::size_t>(sendcount) * ORBIT_MPI_Type_size(sendtype));
+  }
+  return MPI_SUCCESS;
+#endif
+}
+
 /** A C wrapper around MPI_Bcast. */
 int ORBIT_MPI_Bcast(void* ar, int n1, MPI_Datatype data, int n2, MPI_Comm comm ){
   int res = 0;

src/mpi/orbit_mpi.hh

@@ -1,4 +1,4 @@
-#include "Python.h"
+#include <Python.h>
 
 #ifndef ORBIT_MPI_INCLUDE
 #define ORBIT_MPI_INCLUDE
@@ -8,7 +8,7 @@
 #endif
 
 #if USE_MPI > 0
-  #include "mpi.h"
+  #include <mpi.h>
   #define ORBIT_MPI_IN_PLACE MPI_IN_PLACE
 #else
 //---------------------------------------------------------------
@@ -210,6 +210,10 @@ int ORBIT_MPI_Graph_neighbors(MPI_Comm comm, int rank, int maxneighbors, int *ne
 int ORBIT_MPI_Barrier(MPI_Comm comm);
 int ORBIT_MPI_Wait(MPI_Request  *request, MPI_Status *status);
 int ORBIT_MPI_Allreduce(void* buf_in, void* buf_out, int count, MPI_Datatype, MPI_Op, MPI_Comm);
+int ORBIT_MPI_Gather(const void *sendbuf, int sendcount, MPI_Datatype sendtype,
+                     void *recvbuf, int recvcount, MPI_Datatype recvtype, int root, MPI_Comm comm);
+int ORBIT_MPI_Allgather(const void *sendbuf, int sendcount, MPI_Datatype sendtype,
+                         void *recvbuf, int recvcount, MPI_Datatype recvtype, MPI_Comm comm);
 int ORBIT_MPI_Bcast(void* buf, int count, MPI_Datatype, int rank, MPI_Comm);
 int ORBIT_MPI_Send(void* buf, int count, MPI_Datatype, int dest,   int tag, MPI_Comm);
 int ORBIT_MPI_Recv(void* buf, int count, MPI_Datatype, int source, int tag, MPI_Comm, MPI_Status *);