Add OMPI_BIGCOUNT_ARRAY_* macros for collective calls

Signed-off-by: Jake Tronge <[email protected]>
open-mpi · Feb 15, 2024 · e8a4be1 · e8a4be1
1 parent 4008ac8
commit e8a4be1
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 4 deletions.
diff --git a/ompi/mpi/c/alltoallv.c.in b/ompi/mpi/c/alltoallv.c.in
@@ -40,7 +40,12 @@ PROTOTYPE ERROR_CLASS alltoallv(BUFFER sendbuf, COUNT_ARRAY sendcounts,
                                 DATATYPE recvtype, COMM comm)
 {
     int i, size, err;
+    OMPI_BIGCOUNT_ARRAY_COUNT_IN_DECL(sendcounts);
+    OMPI_BIGCOUNT_ARRAY_DISPL_IN_DECL(sdispls);
+    OMPI_BIGCOUNT_ARRAY_COUNT_IN_DECL(recvcounts);
+    OMPI_BIGCOUNT_ARRAY_DISPL_IN_DECL(rdispls);
 
+    size = OMPI_COMM_IS_INTER(comm)?ompi_comm_remote_size(comm):ompi_comm_size(comm);
     SPC_RECORD(OMPI_SPC_ALLTOALLV, 1);
 
     MEMCHECKER(
@@ -56,7 +61,6 @@ PROTOTYPE ERROR_CLASS alltoallv(BUFFER sendbuf, COUNT_ARRAY sendcounts,
 
         memchecker_comm(comm);
 
-        size = OMPI_COMM_IS_INTER(comm)?ompi_comm_remote_size(comm):ompi_comm_size(comm);
         for ( i = 0; i < size; i++ ) {
             if (MPI_IN_PLACE != sendbuf) {
                 /* check if send chunks are defined. */
@@ -125,10 +129,21 @@ PROTOTYPE ERROR_CLASS alltoallv(BUFFER sendbuf, COUNT_ARRAY sendcounts,
     }
 #endif
 
+    /* TODO: MPI_IN_PLACE? */
+    OMPI_BIGCOUNT_ARRAY_COUNT_IN_PREPARE(sendcounts, size, i);
+    OMPI_BIGCOUNT_ARRAY_DISPL_IN_PREPARE(sdispls, size, i);
+    OMPI_BIGCOUNT_ARRAY_COUNT_IN_PREPARE(recvcounts, size, i);
+    OMPI_BIGCOUNT_ARRAY_DISPL_IN_PREPARE(rdispls, size, i);
     /* Invoke the coll component to perform the back-end operation */
-    err = comm->c_coll->coll_alltoallv(sendbuf, sendcounts, sdispls, sendtype,
-                                      recvbuf, recvcounts, rdispls, recvtype,
-                                      comm, comm->c_coll->coll_alltoallv_module);
+    err = comm->c_coll->coll_alltoallv(sendbuf, OMPI_BIGCOUNT_ARRAY_NAME(sendcounts),
+                                       OMPI_BIGCOUNT_ARRAY_NAME(sdispls), sendtype,
+                                       recvbuf, OMPI_BIGCOUNT_ARRAY_NAME(recvcounts),
+                                       OMPI_BIGCOUNT_ARRAY_NAME(rdispls), recvtype,
+                                       comm, comm->c_coll->coll_alltoallv_module);
+    OMPI_BIGCOUNT_ARRAY_COUNT_IN_POST(sendcounts);
+    OMPI_BIGCOUNT_ARRAY_DISPL_IN_POST(sdispls);
+    OMPI_BIGCOUNT_ARRAY_COUNT_IN_POST(recvcounts);
+    OMPI_BIGCOUNT_ARRAY_DISPL_IN_POST(rdispls);
     OMPI_ERRHANDLER_RETURN(err, comm, err, FUNC_NAME);
 }
 
diff --git a/ompi/mpi/c/bindings.h b/ompi/mpi/c/bindings.h
@@ -104,6 +104,55 @@ BEGIN_C_DECLS
         }                                                               \
     } while (0)
 
+/* Handling of MPI_Count/MPI_Aint and int array conversion */
+#define OMPI_BIGCOUNT_ARRAY_NAME(name) tmp_ ## name
+
+#define OMPI_BIGCOUNT_ARRAY_COUNT_IN_DECL(name) int *OMPI_BIGCOUNT_ARRAY_NAME(name)
+#define OMPI_BIGCOUNT_ARRAY_DISPL_IN_DECL(name) OMPI_BIGCOUNT_ARRAY_COUNT_IN_DECL(name)
+#define OMPI_BIGCOUNT_ARRAY_COUNT_OUT_DECL(name) OMPI_BIGCOUNT_ARRAY_COUNT_IN_DECL(name)
+#define OMPI_BIGCOUNT_ARRAY_DISPL_OUT_DECL(name) OMPI_BIGCOUNT_ARRAY_COUNT_IN_DECL(name)
+
+#define OMPI_BIGCOUNT_ARRAY_COUNT_IN_PREPARE(name, size, i) \
+    do { \
+        if (sizeof(*name) == sizeof(*OMPI_BIGCOUNT_ARRAY_NAME(name))) { \
+            OMPI_BIGCOUNT_ARRAY_NAME(name) = (int *) name; \
+            break; \
+        } \
+        OMPI_BIGCOUNT_ARRAY_NAME(name) = malloc(size * sizeof(*OMPI_BIGCOUNT_ARRAY_NAME(name))); \
+        for (i = 0; i < size; ++i) { \
+            OMPI_BIGCOUNT_ARRAY_NAME(name)[i] = name[i]; \
+        } \
+    } while (0)
+#define OMPI_BIGCOUNT_ARRAY_DISPL_IN_PREPARE(name, size, i) OMPI_BIGCOUNT_ARRAY_COUNT_IN_PREPARE(name, size, i)
+#define OMPI_BIGCOUNT_ARRAY_COUNT_OUT_PREPARE(name, size) \
+    do { \
+        if (sizeof(*name) == sizeof(*OMPI_BIGCOUNT_ARRAY_NAME(name))) { \
+            OMPI_BIGCOUNT_ARRAY_NAME(name) = (int *) name; \
+            break; \
+        } \
+        OMPI_BIGCOUNT_ARRAY_NAME(name) = malloc(size * sizeof(*OMPI_BIGCOUNT_ARRAY_NAME(name))); \
+    } while (0)
+#define OMPI_BIGCOUNT_ARRAY_DISPL_OUT_PREPARE(name, size) OMPI_BIGCOUNT_ARRAY_COUNT_OUT_PREPARE(name, size)
+
+#define OMPI_BIGCOUNT_ARRAY_COUNT_IN_POST(name) \
+    do { \
+        if (sizeof(*name) != sizeof(*OMPI_BIGCOUNT_ARRAY_NAME(name))) { \
+            free(OMPI_BIGCOUNT_ARRAY_NAME(name)); \
+        } \
+    } while (0)
+#define OMPI_BIGCOUNT_ARRAY_DISPL_IN_POST(name) OMPI_BIGCOUNT_ARRAY_COUNT_IN_POST(name)
+#define OMPI_BIGCOUNT_ARRAY_COUNT_OUT_POST(name, size, i) \
+    do { \
+        if (sizeof(*name) == sizeof(*OMPI_BIGCOUNT_ARRAY_NAME(name))) { \
+            break; \
+        } \
+        for (i = 0; i < size; ++i) { \
+            name[i] = OMPI_BIGCOUNT_ARRAY_NAME(name)[i]; \
+        } \
+        free(OMPI_BIGCOUNT_ARRAY_NAME(name)); \
+    } while (0)
+#define OMPI_BIGCOUNT_ARRAY_DISPL_OUT_POST(name, size, i) OMPI_BIGCOUNT_ARRAY_COUNT_OUT_POST(name, size, i)
+
 END_C_DECLS
 
 #endif /* OMPI_C_BINDINGS_H */