ch4/ucx: Unify AM completion logic

Since we already have allocated requests at the time we start an active
message send or recv operation, the completion logic is always the
same. In the receive path, we can simply force the completion callback
to execute even for immediately complete operations. On the send side,
forcing the callback execution appears to be broken in UCX, so we
manually call it instead.

Author: raffenet

src/mpid/ch4/netmod/ucx/ucx_am.c

@@ -73,6 +73,7 @@ void MPIDI_UCX_am_send_callback(void *request, ucs_status_t status)
 /* Called when recv buffer is posted */
 int MPIDI_UCX_do_am_recv(MPIR_Request * rreq)
 {
+    int mpi_errno = MPI_SUCCESS;
     void *recv_buf;
     bool is_contig;
     MPI_Aint data_sz, in_data_sz;
@@ -91,25 +92,19 @@ int MPIDI_UCX_do_am_recv(MPIR_Request * rreq)
     }
 
     MPIDI_UCX_ucp_request_t *ucp_request;
-    size_t received_length;
     ucp_request_param_t param = {
-        .op_attr_mask = UCP_OP_ATTR_FIELD_CALLBACK | UCP_OP_ATTR_FIELD_RECV_INFO | UCP_OP_ATTR_FIELD_USER_DATA,
+        .op_attr_mask = UCP_OP_ATTR_FIELD_CALLBACK | UCP_OP_ATTR_FIELD_USER_DATA | UCP_OP_ATTR_FLAG_NO_IMM_CMPL,
         .cb.recv_am = &MPIDI_UCX_am_recv_callback_nbx,
-        .recv_info.length = &received_length,
         .user_data = rreq,
     };
     void *data_desc = MPIDI_UCX_AM_RECV_REQUEST(rreq, data_desc);
     /* note: use in_data_sz to match promised data size */
     ucp_request = ucp_am_recv_data_nbx(MPIDI_UCX_global.ctx[vci].worker,
                                        data_desc, recv_buf, in_data_sz, &param);
-    if (ucp_request == NULL) {
-        /* completed immediately */
-        MPIDI_UCX_ucp_request_t tmp_ucp_request;
-        tmp_ucp_request.req = rreq;
-        MPIDI_UCX_am_recv_callback_nbx(&tmp_ucp_request, UCS_OK, received_length, NULL);
-    }
+    MPIDI_UCX_CHK_REQUEST(ucp_request);
 
-    return MPI_SUCCESS;
+  fn_fail:
+    return mpi_errno;
 }
 
 /* Am handler for messages sent from ucp_am_send_nbx. Registered with
@@ -175,7 +170,9 @@ void MPIDI_UCX_am_recv_callback_nbx(void *request, ucs_status_t status, size_t l
         MPIDIG_recv_done(length, rreq);
     }
     MPIDIG_REQUEST(rreq, req->target_cmpl_cb) (rreq);
-    ucp_request_release(request);
+    if (request) {
+        ucp_request_release(request);
+    }
 }
 
 void MPIDI_UCX_am_isend_callback_nbx(void *request, ucs_status_t status, void *user_data)
@@ -189,5 +186,8 @@ void MPIDI_UCX_am_isend_callback_nbx(void *request, ucs_status_t status, void *u
     MPL_free(MPIDI_UCX_AM_SEND_REQUEST(req, pack_buffer));
     MPIDI_UCX_AM_SEND_REQUEST(req, pack_buffer) = NULL;
     MPIDIG_global.origin_cbs[handler_id] (req);
+    if (request) {
+        ucp_request_release(request);
+    }
 }
 #endif

src/mpid/ch4/netmod/ucx/ucx_am.h

@@ -65,21 +65,17 @@ MPL_STATIC_INLINE_PREFIX int MPIDI_NM_am_isend(int rank,
         data_ptr = data;
         data_sz = count;
     }
+    MPIDI_UCX_AM_SEND_REQUEST(sreq, pack_buffer) = header;
+    MPIDI_UCX_AM_SEND_REQUEST(sreq, handler_id) = handler_id;
     ucp_request = (MPIDI_UCX_ucp_request_t *) ucp_am_send_nbx(ep, MPIDI_UCX_AM_NBX_HANDLER_ID,
                                                               header, header_size,
                                                               data_ptr, data_sz, &param);
     MPIDI_UCX_CHK_REQUEST(ucp_request);
-    /* if send is done, free all resources and complete the request */
     if (ucp_request == NULL) {
-        MPL_free(header);
-        MPIDIG_global.origin_cbs[handler_id] (sreq);
-        goto fn_exit;
+        /* call callback because UCP_OP_ATTR_FLAG_NO_IMM_CMPL is broken for ucp_am_send_nbx */
+        MPIDI_UCX_am_isend_callback_nbx(NULL, UCS_OK, sreq);
     }
 
-    MPIDI_UCX_AM_SEND_REQUEST(sreq, pack_buffer) = header;
-    MPIDI_UCX_AM_SEND_REQUEST(sreq, handler_id) = handler_id;
-    ucp_request_release(ucp_request);
-
 #else /* !HAVE_UCP_AM_NBX */
 
     MPL_pointer_attr_t attr;