From e06a8e21c9d16589e97171312906d729a350eba1 Mon Sep 17 00:00:00 2001
From: "Wood, Tony" <woodap@ornl.gov>
Date: Tue, 8 Sep 2026 14:41:05 -0400
Subject: [PATCH 1/2] Introduce ORBIT_MPI_Gather/_Allgather. Don't clobber py
 exitcode.

---
 src/mpi/orbit_mpi.cc | 52 +++++++++++++++++++++++++++++++++++++++++++-
 src/mpi/orbit_mpi.hh |  4 ++++
 2 files changed, 55 insertions(+), 1 deletion(-)

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<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;
diff --git a/src/mpi/orbit_mpi.hh b/src/mpi/orbit_mpi.hh
index bff6ea5d..52f78db8 100644
--- a/src/mpi/orbit_mpi.hh
+++ b/src/mpi/orbit_mpi.hh
@@ -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 *);

From b9a72bf141b72c00259b1979e991e18e59b0b32f Mon Sep 17 00:00:00 2001
From: "Wood, Tony" <woodap@ornl.gov>
Date: Tue, 8 Sep 2026 14:54:42 -0400
Subject: [PATCH 2/2] Should be angle brackets.

---
 src/mpi/orbit_mpi.hh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/mpi/orbit_mpi.hh b/src/mpi/orbit_mpi.hh
index 52f78db8..ee723990 100644
--- a/src/mpi/orbit_mpi.hh
+++ b/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
 //---------------------------------------------------------------