diff --git a/src/mpi/orbit_mpi.cc b/src/mpi/orbit_mpi.cc index 4965a9e1..26739276 100644 --- a/src/mpi/orbit_mpi.cc +++ b/src/mpi/orbit_mpi.cc @@ -28,6 +28,23 @@ 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 @@ -35,7 +52,7 @@ int ORBIT_MPI_Init(){ 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(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(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; diff --git a/src/mpi/orbit_mpi.hh b/src/mpi/orbit_mpi.hh index bff6ea5d..ee723990 100644 --- a/src/mpi/orbit_mpi.hh +++ b/src/mpi/orbit_mpi.hh @@ -1,4 +1,4 @@ -#include "Python.h" +#include #ifndef ORBIT_MPI_INCLUDE #define ORBIT_MPI_INCLUDE @@ -8,7 +8,7 @@ #endif #if USE_MPI > 0 - #include "mpi.h" + #include #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 *);