Merge pull request open-mpi#10515 from bosilca/fix/retain_objects_for_topo_comm

awlauria · web-flow · commit 0a180c1fb096 · 2022-10-12T09:34:38.000-04:00
Correctly retain the datatypes for topological communicators.
diff --git a/ompi/mca/coll/base/coll_base_util.c b/ompi/mca/coll/base/coll_base_util.c
@@ -2,7 +2,7 @@
  * Copyright (c) 2004-2005 The Trustees of Indiana University and Indiana
  *                         University Research and Technology
  *                         Corporation.  All rights reserved.
- * Copyright (c) 2004-2021 The University of Tennessee and The University
+ * Copyright (c) 2004-2022 The University of Tennessee and The University
  *                         of Tennessee Research Foundation.  All rights
  *                         reserved.
  * Copyright (c) 2004-2005 High Performance Computing Center Stuttgart,
@@ -244,25 +244,21 @@ int ompi_coll_base_retain_datatypes( ompi_request_t *req, ompi_datatype_t *stype
     return OMPI_SUCCESS;
 }
 
-static void release_vecs_callback(ompi_coll_base_nbc_request_t *request) {
-    ompi_communicator_t *comm = request->super.req_mpi_object.comm;
-    int scount, rcount;
-    if (OMPI_COMM_IS_TOPO(comm)) {
-        (void)mca_topo_base_neighbor_count (comm, &rcount, &scount);
-    } else {
-        scount = rcount = OMPI_COMM_IS_INTER(comm)?ompi_comm_remote_size(comm):ompi_comm_size(comm);
-    }
+static void release_vecs_callback(ompi_coll_base_nbc_request_t *request)
+{
     if (NULL != request->data.refcounted.vecs.stypes) {
-        for (int i=0; i<scount; i++) {
-            if (NULL != request->data.refcounted.vecs.stypes[i]) {
+        for (int i = 0; i < request->data.refcounted.vecs.scount; i++) {
+            if (NULL != request->data.refcounted.vecs.stypes[i] &&
+                !ompi_datatype_is_predefined(request->data.refcounted.vecs.stypes[i])) {
                 OMPI_DATATYPE_RELEASE_NO_NULLIFY(request->data.refcounted.vecs.stypes[i]);
             }
         }
         request->data.refcounted.vecs.stypes = NULL;
     }
     if (NULL != request->data.refcounted.vecs.rtypes) {
-        for (int i=0; i<rcount; i++) {
-            if (NULL != request->data.refcounted.vecs.rtypes[i]) {
+        for (int i = 0; i < request->data.refcounted.vecs.rcount; i++) {
+            if (NULL != request->data.refcounted.vecs.rtypes[i] &&
+                !ompi_datatype_is_predefined(request->data.refcounted.vecs.rtypes[i])) {
                 OMPI_DATATYPE_RELEASE_NO_NULLIFY(request->data.refcounted.vecs.rtypes[i]);
             }
         }
@@ -292,35 +288,47 @@ static int free_vecs_callback(struct ompi_request_t **rptr) {
 }
 
 int ompi_coll_base_retain_datatypes_w( ompi_request_t *req,
-                                       ompi_datatype_t * const stypes[], ompi_datatype_t * const rtypes[]) {
+                                       ompi_datatype_t * const stypes[],
+                                       ompi_datatype_t * const rtypes[],
+                                       bool use_topo)
+{
     ompi_coll_base_nbc_request_t *request = (ompi_coll_base_nbc_request_t *)req;
-    bool retain = false;
     ompi_communicator_t *comm = request->super.req_mpi_object.comm;
     int scount, rcount;
+
     if (REQUEST_COMPLETE(req)) {
         return OMPI_SUCCESS;
     }
-    if (OMPI_COMM_IS_TOPO(comm)) {
+
+    if (use_topo && OMPI_COMM_IS_TOPO(comm)) {
         (void)mca_topo_base_neighbor_count (comm, &rcount, &scount);
     } else {
         scount = rcount = OMPI_COMM_IS_INTER(comm)?ompi_comm_remote_size(comm):ompi_comm_size(comm);
     }
 
-    for (int i=0; i<scount; i++) {
-        if (NULL != stypes && NULL != stypes[i] && !ompi_datatype_is_predefined(stypes[i])) {
-            OBJ_RETAIN(stypes[i]);
-            retain = true;
+    request->data.refcounted.vecs.scount = 0;  /* default value */
+    if (NULL != stypes) {
+        for (int i = 0; i < scount; i++) {
+            if (NULL != stypes[i] && !ompi_datatype_is_predefined(stypes[i])) {
+                OBJ_RETAIN(stypes[i]);
+                request->data.refcounted.vecs.scount = i;  /* last valid type */
+            }
         }
     }
-    for (int i=0; i<rcount; i++) {
-        if (NULL != rtypes && NULL != rtypes[i] && !ompi_datatype_is_predefined(rtypes[i])) {
-            OBJ_RETAIN(rtypes[i]);
-            retain = true;
+    request->data.refcounted.vecs.rcount = 0;  /* default value */
+    if (NULL != rtypes) {
+        for (int i = 0; i < rcount; i++) {
+            if (NULL != rtypes[i] && !ompi_datatype_is_predefined(rtypes[i])) {
+                OBJ_RETAIN(rtypes[i]);
+                request->data.refcounted.vecs.rcount = i;  /* last valid type */
+            }
         }
     }
-    if (OPAL_UNLIKELY(retain)) {
+    if (OPAL_LIKELY(request->data.refcounted.vecs.scount | request->data.refcounted.vecs.rcount) ) {
         request->data.refcounted.vecs.stypes = (ompi_datatype_t **) stypes;
         request->data.refcounted.vecs.rtypes = (ompi_datatype_t **) rtypes;
+        request->data.refcounted.vecs.scount = scount;
+        request->data.refcounted.vecs.rcount = rcount;
         if (req->req_persistent) {
             request->cb.req_free = req->req_free;
             req->req_free = free_vecs_callback;
diff --git a/ompi/mca/coll/base/coll_base_util.h b/ompi/mca/coll/base/coll_base_util.h
@@ -2,7 +2,7 @@
  * Copyright (c) 2004-2005 The Trustees of Indiana University and Indiana
  *                         University Research and Technology
  *                         Corporation.  All rights reserved.
- * Copyright (c) 2004-2021 The University of Tennessee and The University
+ * Copyright (c) 2004-2022 The University of Tennessee and The University
  *                         of Tennessee Research Foundation.  All rights
  *                         reserved.
  * Copyright (c) 2004-2007 High Performance Computing Center Stuttgart,
@@ -68,6 +68,8 @@ struct ompi_coll_base_nbc_request_t {
             struct {
                 ompi_datatype_t * const *stypes;
                 ompi_datatype_t * const *rtypes;
+                int scount;
+                int rcount;
             } vecs;
         } refcounted;
         void* release_arrays[OMPI_REQ_NB_RELEASE_ARRAYS];
@@ -175,7 +177,7 @@ int ompi_coll_base_retain_op( ompi_request_t *request,
  * (will be cast internally).
  */
 int ompi_coll_base_retain_datatypes( ompi_request_t *request,
-                                      ompi_datatype_t *stype,
+                                     ompi_datatype_t *stype,
                                      ompi_datatype_t *rtype);
 
 /**
@@ -185,7 +187,8 @@ int ompi_coll_base_retain_datatypes( ompi_request_t *request,
  */
 int ompi_coll_base_retain_datatypes_w( ompi_request_t *request,
                                        ompi_datatype_t * const stypes[],
-                                       ompi_datatype_t * const rtypes[]);
+                                       ompi_datatype_t * const rtypes[],
+                                       bool use_topo);
 
 /* File reading function */
 int ompi_coll_base_file_getnext_long(FILE *fptr, int *fileline, long* val);
diff --git a/ompi/mpi/c/alltoallw_init.c b/ompi/mpi/c/alltoallw_init.c
@@ -3,7 +3,7 @@
  * Copyright (c) 2004-2007 The Trustees of Indiana University and Indiana
  *                         University Research and Technology
  *                         Corporation.  All rights reserved.
- * Copyright (c) 2004-2020 The University of Tennessee and The University
+ * Copyright (c) 2004-2022 The University of Tennessee and The University
  *                         of Tennessee Research Foundation.  All rights
  *                         reserved.
  * Copyright (c) 2004-2008 High Performance Computing Center Stuttgart,
@@ -127,7 +127,7 @@ int MPI_Alltoallw_init(const void *sendbuf, const int sendcounts[], const int sd
                                             rdispls, recvtypes, comm, info, request,
                                             comm->c_coll->coll_alltoallw_init_module);
     if (OPAL_LIKELY(OMPI_SUCCESS == err)) {
-        ompi_coll_base_retain_datatypes_w(*request, (MPI_IN_PLACE==sendbuf)?NULL:sendtypes, recvtypes);
+        ompi_coll_base_retain_datatypes_w(*request, (MPI_IN_PLACE==sendbuf)?NULL:sendtypes, recvtypes, false);
     }
     OMPI_ERRHANDLER_RETURN(err, comm, err, FUNC_NAME);
 }
diff --git a/ompi/mpi/c/ialltoallw.c b/ompi/mpi/c/ialltoallw.c
@@ -3,7 +3,7 @@
  * Copyright (c) 2004-2007 The Trustees of Indiana University and Indiana
  *                         University Research and Technology
  *                         Corporation.  All rights reserved.
- * Copyright (c) 2004-2020 The University of Tennessee and The University
+ * Copyright (c) 2004-2022 The University of Tennessee and The University
  *                         of Tennessee Research Foundation.  All rights
  *                         reserved.
  * Copyright (c) 2004-2008 High Performance Computing Center Stuttgart,
@@ -121,7 +121,7 @@ int MPI_Ialltoallw(const void *sendbuf, const int sendcounts[], const int sdispl
                                        rdispls, recvtypes, comm, request,
                                        comm->c_coll->coll_ialltoallw_module);
     if (OPAL_LIKELY(OMPI_SUCCESS == err)) {
-        ompi_coll_base_retain_datatypes_w(*request, (MPI_IN_PLACE==sendbuf)?NULL:sendtypes, recvtypes);
+        ompi_coll_base_retain_datatypes_w(*request, (MPI_IN_PLACE==sendbuf)?NULL:sendtypes, recvtypes, false);
     }
     OMPI_ERRHANDLER_RETURN(err, comm, err, FUNC_NAME);
 }
diff --git a/ompi/mpi/c/ineighbor_alltoallw.c b/ompi/mpi/c/ineighbor_alltoallw.c
@@ -3,7 +3,7 @@
  * Copyright (c) 2004-2007 The Trustees of Indiana University and Indiana
  *                         University Research and Technology
  *                         Corporation.  All rights reserved.
- * Copyright (c) 2004-2020 The University of Tennessee and The University
+ * Copyright (c) 2004-2022 The University of Tennessee and The University
  *                         of Tennessee Research Foundation.  All rights
  *                         reserved.
  * Copyright (c) 2004-2008 High Performance Computing Center Stuttgart,
@@ -146,7 +146,7 @@ int MPI_Ineighbor_alltoallw(const void *sendbuf, const int sendcounts[], const M
                                                 recvbuf, recvcounts, rdispls, recvtypes, comm, request,
                                                 comm->c_coll->coll_ineighbor_alltoallw_module);
     if (OPAL_LIKELY(OMPI_SUCCESS == err)) {
-        ompi_coll_base_retain_datatypes_w(*request, sendtypes, recvtypes);
+        ompi_coll_base_retain_datatypes_w(*request, sendtypes, recvtypes, true);
     }
     OMPI_ERRHANDLER_RETURN(err, comm, err, FUNC_NAME);
 }
diff --git a/ompi/mpi/c/neighbor_alltoallw_init.c b/ompi/mpi/c/neighbor_alltoallw_init.c
@@ -3,7 +3,7 @@
  * Copyright (c) 2004-2007 The Trustees of Indiana University and Indiana
  *                         University Research and Technology
  *                         Corporation.  All rights reserved.
- * Copyright (c) 2004-2020 The University of Tennessee and The University
+ * Copyright (c) 2004-2022 The University of Tennessee and The University
  *                         of Tennessee Research Foundation.  All rights
  *                         reserved.
  * Copyright (c) 2004-2008 High Performance Computing Center Stuttgart,
@@ -148,7 +148,7 @@ int MPI_Neighbor_alltoallw_init(const void *sendbuf, const int sendcounts[], con
                                                      info, request,
                                                      comm->c_coll->coll_neighbor_alltoallw_init_module);
     if (OPAL_LIKELY(OMPI_SUCCESS == err)) {
-        ompi_coll_base_retain_datatypes_w(*request, sendtypes, recvtypes);
+        ompi_coll_base_retain_datatypes_w(*request, sendtypes, recvtypes, true);
     }
     OMPI_ERRHANDLER_RETURN(err, comm, err, FUNC_NAME);
 }
