Merge pull request #9175 from karasevb/topic/v4.1.x/ucx_pml_dtype_retain

v4.1.x: pml/ucx: fix datatype refcount

Author: jsquyres

ompi/mca/pml/ucx/pml_ucx.c

@@ -564,7 +564,9 @@ int mca_pml_ucx_irecv_init(void *buf, size_t count, ompi_datatype_t *datatype,
     req->flags                    = 0;
     req->buffer                   = buf;
     req->count                    = count;
-    req->datatype.datatype        = mca_pml_ucx_get_datatype(datatype);
+    req->ompi_datatype            = datatype;
+    req->datatype                 = mca_pml_ucx_get_datatype(datatype);
+    OMPI_DATATYPE_RETAIN(datatype);
 
     PML_UCX_MAKE_RECV_TAG(req->tag, req->recv.tag_mask, tag, src, comm);
 
@@ -699,12 +701,13 @@ int mca_pml_ucx_isend_init(const void *buf, size_t count, ompi_datatype_t *datat
     req->tag                      = PML_UCX_MAKE_SEND_TAG(tag, comm);
     req->send.mode                = mode;
     req->send.ep                  = ep;
+    req->ompi_datatype            = datatype;
+    OMPI_DATATYPE_RETAIN(datatype);
 
     if (MCA_PML_BASE_SEND_BUFFERED == mode) {
-        req->datatype.ompi_datatype = datatype;
-        OBJ_RETAIN(datatype);
+        req->datatype = NULL;
     } else {
-        req->datatype.datatype = mca_pml_ucx_get_datatype(datatype);
+        req->datatype = mca_pml_ucx_get_datatype(datatype);
     }
 
     *request = &req->ompi;
@@ -1113,16 +1116,16 @@ int mca_pml_ucx_start(size_t count, ompi_request_t** requests)
             tmp_req = (ompi_request_t*)mca_pml_ucx_common_send(preq->send.ep,
                                                                preq->buffer,
                                                                preq->count,
-                                                               preq->datatype.ompi_datatype,
-                                                               preq->datatype.datatype,
+                                                               preq->ompi_datatype,
+                                                               preq->datatype,
                                                                preq->tag,
                                                                preq->send.mode,
                                                                mca_pml_ucx_psend_completion);
         } else {
             PML_UCX_VERBOSE(8, "start recv request %p", (void*)preq);
             tmp_req = (ompi_request_t*)ucp_tag_recv_nb(ompi_pml_ucx.ucp_worker,
                                                        preq->buffer, preq->count,
-                                                       preq->datatype.datatype,
+                                                       preq->datatype,
                                                        preq->tag,
                                                        preq->recv.tag_mask,
                                                        mca_pml_ucx_precv_completion);

ompi/mca/pml/ucx/pml_ucx_request.c

@@ -212,10 +212,7 @@ static int mca_pml_ucx_persistent_request_free(ompi_request_t **rptr)
         mca_pml_ucx_persistent_request_detach(preq, tmp_req);
         ucp_request_free(tmp_req);
     }
-    if ((preq->flags & MCA_PML_UCX_REQUEST_FLAG_SEND) &&
-         (MCA_PML_BASE_SEND_BUFFERED == preq->send.mode)) {
-        OBJ_RELEASE(preq->datatype.ompi_datatype);
-    }
+    OMPI_DATATYPE_RELEASE(preq->ompi_datatype);
     PML_UCX_FREELIST_RETURN(&ompi_pml_ucx.persistent_reqs, &preq->ompi.super);
     *rptr = MPI_REQUEST_NULL;
     return OMPI_SUCCESS;

ompi/mca/pml/ucx/pml_ucx_request.h

@@ -99,10 +99,8 @@ struct pml_ucx_persistent_request {
     unsigned                          flags;
     void                              *buffer;
     size_t                            count;
-    union {
-        ucp_datatype_t                datatype;
-        ompi_datatype_t              *ompi_datatype;
-    } datatype;
+    ucp_datatype_t                    datatype;
+    ompi_datatype_t                   *ompi_datatype;
     ucp_tag_t                         tag;
     struct {
         mca_pml_base_send_mode_t      mode;