feat: Support MPI_Open_port() with MPICH CH4:UCX

mpi4py · Dec 14, 2023 · 05316e5 · 05316e5
1 parent 5085a93
commit 05316e5
Showing 1 changed file with 50 additions and 0 deletions.
diff --git a/src/lib-mpi/compat/mpich.h b/src/lib-mpi/compat/mpich.h
@@ -4,6 +4,56 @@
 
 /* -------------------------------------------------------------------------- */
 
+/* https://github.com/pmodels/mpich/pull/5467 */
+
+#undef  MPI_MAX_PORT_NAME
+#define MPI_MAX_PORT_NAME 1024
+
+static int PyMPI_MPICH_port_info(MPI_Info info, MPI_Info *port_info)
+{
+  int ierr;
+# define pympi_str_(s) #s
+# define pympi_str(s) pympi_str_(s)
+  const char *key = "port_name_size";
+  const char *val = pympi_str(MPI_MAX_PORT_NAME);
+# undef pympi_str_
+# undef pympi_str
+  if (info == MPI_INFO_NULL) {
+    ierr = MPI_Info_create(port_info); if (ierr) return ierr;
+  } else {
+    ierr = MPI_Info_dup(info, port_info); if (ierr) return ierr;
+  }
+  ierr = MPI_Info_set(*port_info, key, val);
+  if (ierr) (void) MPI_Info_free(port_info);
+  return ierr;
+}
+
+static int PyMPI_MPICH_MPI_Open_port(MPI_Info info, char *port_name)
+{
+  int ierr;
+  ierr = PyMPI_MPICH_port_info(info, &info); if (ierr) return ierr;
+  ierr = MPI_Open_port(info, port_name);
+  (void) MPI_Info_free(&info);
+  return ierr;
+}
+#undef  MPI_Open_port
+#define MPI_Open_port PyMPI_MPICH_MPI_Open_port
+
+static int PyMPI_MPICH_MPI_Lookup_name(const char *service_name,
+                                       MPI_Info   info,
+                                       char       *port_name)
+{
+  int ierr;
+  ierr = PyMPI_MPICH_port_info(info, &info); if (ierr) return ierr;
+  ierr = MPI_Lookup_name(service_name, info, port_name);
+  (void) MPI_Info_free(&info);
+  return ierr;
+}
+#undef  MPI_Lookup_name
+#define MPI_Lookup_name PyMPI_MPICH_MPI_Lookup_name
+
+/* -------------------------------------------------------------------------- */
+
 /* https://github.com/pmodels/mpich/issues/5413 */
 /* https://github.com/pmodels/mpich/pull/6146   */