Merge pull request #9668 from bwbarrett/bugfix/alltoallv_in_place

Alltoall{,v,w} IN_PLACE fixes
open-mpi · Nov 16, 2021 · 7c9ab0a · 7c9ab0a
2 parents f8344de + 6802702
commit 7c9ab0a
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 19 deletions.
diff --git a/ompi/mca/coll/base/coll_base_alltoall.c b/ompi/mca/coll/base/coll_base_alltoall.c
@@ -15,6 +15,8 @@
  * Copyright (c) 2014-2017 Research Organization for Information Science
  *                         and Technology (RIST). All rights reserved.
  * Copyright (c) 2017      IBM Corporation. All rights reserved.
+ * Copyright (c) 2021      Amazon.com, Inc. or its affiliates.  All Rights
+ *                         reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -108,13 +110,13 @@ mca_coll_base_alltoall_intra_basic_inplace(const void *rbuf, int rcount,
 
         /* Receive data from the right */
         err = MCA_PML_CALL(irecv ((char *) rbuf + right * rcount * extent, rcount, rdtype,
-                                  right, MCA_COLL_BASE_TAG_ALLTOALLW, comm, &req));
+                                  right, MCA_COLL_BASE_TAG_ALLTOALL, comm, &req));
         if (MPI_SUCCESS != err) { goto error_hndl; }
 
         if( left != right ) {
             /* Send data to the left */
             err = MCA_PML_CALL(send ((char *) rbuf + left * rcount * extent, rcount, rdtype,
-                                     left, MCA_COLL_BASE_TAG_ALLTOALLW, MCA_PML_BASE_SEND_STANDARD,
+                                     left, MCA_COLL_BASE_TAG_ALLTOALL, MCA_PML_BASE_SEND_STANDARD,
                                      comm));
             if (MPI_SUCCESS != err) { goto error_hndl; }
 
@@ -123,13 +125,13 @@ mca_coll_base_alltoall_intra_basic_inplace(const void *rbuf, int rcount,
 
             /* Receive data from the left */
             err = MCA_PML_CALL(irecv ((char *) rbuf + left * rcount * extent, rcount, rdtype,
-                                      left, MCA_COLL_BASE_TAG_ALLTOALLW, comm, &req));
+                                      left, MCA_COLL_BASE_TAG_ALLTOALL, comm, &req));
             if (MPI_SUCCESS != err) { goto error_hndl; }
         }
 
         /* Send data to the right */
         err = MCA_PML_CALL(send ((char *) tmp_buffer,  packed_size, MPI_PACKED,
-                                 right, MCA_COLL_BASE_TAG_ALLTOALLW, MCA_PML_BASE_SEND_STANDARD,
+                                 right, MCA_COLL_BASE_TAG_ALLTOALL, MCA_PML_BASE_SEND_STANDARD,
                                  comm));
         if (MPI_SUCCESS != err) { goto error_hndl; }
 

diff --git a/ompi/mca/coll/base/coll_base_alltoallv.c b/ompi/mca/coll/base/coll_base_alltoallv.c
@@ -17,6 +17,8 @@
  * Copyright (c) 2014-2017 Research Organization for Information Science
  *                         and Technology (RIST). All rights reserved.
  * Copyright (c) 2017      IBM Corporation. All rights reserved.
+ * Copyright (c) 2021      Amazon.com, Inc. or its affiliates.  All Rights
+ *                         reserved.
  * $COPYRIGHT$
  *
  * Additional copyrights may follow
@@ -52,7 +54,8 @@ mca_coll_base_alltoallv_intra_basic_inplace(const void *rbuf, const int *rcounts
                                             mca_coll_base_module_t *module)
 {
     int i, size, rank, left, right, err = MPI_SUCCESS, line;
-    ompi_request_t *req;
+    ptrdiff_t extent;
+    ompi_request_t *req = MPI_REQUEST_NULL;
     char *tmp_buffer;
     size_t packed_size = 0, max_size;
     opal_convertor_t convertor;
@@ -87,6 +90,8 @@ mca_coll_base_alltoallv_intra_basic_inplace(const void *rbuf, const int *rcounts
     }
 #endif  /* OPAL_ENABLE_HETEROGENEOUS_SUPPORT */
 
+    ompi_datatype_type_extent(rdtype, &extent);
+
     /* Allocate a temporary buffer */
     tmp_buffer = calloc (max_size, 1);
     if( NULL == tmp_buffer) { err = OMPI_ERR_OUT_OF_RESOURCE; line = __LINE__; goto error_hndl; }
@@ -102,43 +107,43 @@ mca_coll_base_alltoallv_intra_basic_inplace(const void *rbuf, const int *rcounts
             ompi_proc_t *right_proc = ompi_comm_peer_lookup(comm, right);
             opal_convertor_clone(right_proc->super.proc_convertor, &convertor, 0);
             opal_convertor_prepare_for_send(&convertor, &rdtype->super, rcounts[right],
-                                            (char *) rbuf + rdisps[right]);
+                                            (char *) rbuf + rdisps[right] * extent);
             packed_size = max_size;
             err = opal_convertor_pack(&convertor, &iov, &iov_count, &packed_size);
             if (1 != err) { goto error_hndl; }
 
             /* Receive data from the right */
-            err = MCA_PML_CALL(irecv ((char *) rbuf + rdisps[right], rcounts[right], rdtype,
-                                      right, MCA_COLL_BASE_TAG_ALLTOALLW, comm, &req));
+            err = MCA_PML_CALL(irecv ((char *) rbuf + rdisps[right] * extent, rcounts[right], rdtype,
+                                      right, MCA_COLL_BASE_TAG_ALLTOALLV, comm, &req));
             if (MPI_SUCCESS != err) { goto error_hndl; }
         }
 
         if( (left != right) && (0 != rcounts[left]) ) {
             /* Send data to the left */
-            err = MCA_PML_CALL(send ((char *) rbuf + rdisps[left], rcounts[left], rdtype,
-                                     left, MCA_COLL_BASE_TAG_ALLTOALLW, MCA_PML_BASE_SEND_STANDARD,
+            err = MCA_PML_CALL(send ((char *) rbuf + rdisps[left] * extent, rcounts[left], rdtype,
+                                     left, MCA_COLL_BASE_TAG_ALLTOALLV, MCA_PML_BASE_SEND_STANDARD,
                                      comm));
             if (MPI_SUCCESS != err) { goto error_hndl; }
 
             err = ompi_request_wait (&req, MPI_STATUSES_IGNORE);
             if (MPI_SUCCESS != err) { goto error_hndl; }
 
             /* Receive data from the left */
-            err = MCA_PML_CALL(irecv ((char *) rbuf + rdisps[left], rcounts[left], rdtype,
-                                      left, MCA_COLL_BASE_TAG_ALLTOALLW, comm, &req));
+            err = MCA_PML_CALL(irecv ((char *) rbuf + rdisps[left] * extent, rcounts[left], rdtype,
+                                      left, MCA_COLL_BASE_TAG_ALLTOALLV, comm, &req));
             if (MPI_SUCCESS != err) { goto error_hndl; }
         }
 
         if( 0 != rcounts[right] ) {  /* nothing to exchange with the peer on the right */
             /* Send data to the right */
             err = MCA_PML_CALL(send ((char *) tmp_buffer,  packed_size, MPI_PACKED,
-                                     right, MCA_COLL_BASE_TAG_ALLTOALLW, MCA_PML_BASE_SEND_STANDARD,
+                                     right, MCA_COLL_BASE_TAG_ALLTOALLV, MCA_PML_BASE_SEND_STANDARD,
                                      comm));
             if (MPI_SUCCESS != err) { goto error_hndl; }
-
-            err = ompi_request_wait (&req, MPI_STATUSES_IGNORE);
-            if (MPI_SUCCESS != err) { goto error_hndl; }
         }
+
+        err = ompi_request_wait (&req, MPI_STATUSES_IGNORE);
+        if (MPI_SUCCESS != err) { goto error_hndl; }
     }
 
  error_hndl:

diff --git a/ompi/mca/coll/basic/coll_basic_alltoallw.c b/ompi/mca/coll/basic/coll_basic_alltoallw.c
@@ -139,10 +139,10 @@ mca_coll_basic_alltoallw_intra_inplace(const void *rbuf, const int *rcounts, con
                                      right, MCA_COLL_BASE_TAG_ALLTOALLW, MCA_PML_BASE_SEND_STANDARD,
                                      comm));
             if (MPI_SUCCESS != err) { goto error_hndl; }
-
-            err = ompi_request_wait (&req, MPI_STATUSES_IGNORE);
-            if (MPI_SUCCESS != err) { goto error_hndl; }
         }
+
+        err = ompi_request_wait (&req, MPI_STATUSES_IGNORE);
+        if (MPI_SUCCESS != err) { goto error_hndl; }
     }
 
  error_hndl: