ch4/ucx: Use UCX datatypes in AM path

The nbx active message interfaces can support noncontig data, so use
it. We already do the same in the tagged send and recv path.

Author: raffenet

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

@@ -78,18 +78,6 @@ int MPIDI_UCX_do_am_recv(MPIR_Request * rreq)
     MPI_Aint data_sz, in_data_sz;
     int vci = MPIDI_Request_get_vci(rreq);
 
-    MPIDIG_get_recv_buffer(&recv_buf, &data_sz, &is_contig, &in_data_sz, rreq);
-    if (!is_contig || in_data_sz > data_sz) {
-        /* non-contig datatype, need receive into pack buffer */
-        /* ucx will error out if buffer size is less than the promised data size,
-         * also use a pack buffer in this case */
-        recv_buf = MPL_malloc(in_data_sz, MPL_MEM_OTHER);
-        MPIR_Assert(recv_buf);
-        MPIDI_UCX_AM_RECV_REQUEST(rreq, pack_buffer) = recv_buf;
-    } else {
-        MPIDI_UCX_AM_RECV_REQUEST(rreq, pack_buffer) = NULL;
-    }
-
     MPIDI_UCX_ucp_request_t *ucp_request;
     size_t received_length;
     ucp_request_param_t param = {
@@ -99,6 +87,26 @@ int MPIDI_UCX_do_am_recv(MPIR_Request * rreq)
         .recv_info.length = &received_length,
         .user_data = rreq,
     };
+
+    MPIDIG_get_recv_buffer(&recv_buf, &data_sz, &is_contig, &in_data_sz, rreq);
+    if (in_data_sz > data_sz) {
+        /* ucx will error out if buffer size is less than the promised data size,
+         * use a pack buffer in this case */
+        /* FIXME: what? how does UCX know the buffer size differs? */
+        recv_buf = MPL_malloc(in_data_sz, MPL_MEM_OTHER);
+        MPIR_Assert(recv_buf);
+        MPIDI_UCX_AM_RECV_REQUEST(rreq, pack_buffer) = recv_buf;
+    } else {
+        MPIDI_UCX_AM_RECV_REQUEST(rreq, pack_buffer) = NULL;
+        if (!is_contig) {
+            MPIR_Datatype *dt_ptr;
+            MPIR_Datatype_get_ptr(MPIDIG_REQUEST(rreq, datatype), dt_ptr);
+            param.op_attr_mask |= UCP_OP_ATTR_FIELD_DATATYPE;
+            param.datatype = dt_ptr->dev.netmod.ucx.ucp_datatype;
+            MPIR_Datatype_ptr_add_ref(dt_ptr);
+        }
+    }
+
     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,

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

@@ -42,52 +42,41 @@ MPL_STATIC_INLINE_PREFIX int MPIDI_NM_am_isend(int rank,
 
 #ifdef HAVE_UCP_AM_NBX
     size_t header_size = sizeof(ucx_hdr) + am_hdr_sz;
-    void *send_buf, *header, *data_ptr;
-    /* note: since we are not copying large contig gpu data, it is less useful
-     * to use MPIR_gpu_malloc_host */
-    if (dt_contig) {
-        /* only need copy headers */
-        send_buf = MPL_malloc(header_size, MPL_MEM_OTHER);
-        MPIR_Assert(send_buf);
-        header = send_buf;
-
-        MPIR_Memcpy(header, &ucx_hdr, sizeof(ucx_hdr));
-        MPIR_Memcpy((char *) header + sizeof(ucx_hdr), am_hdr, am_hdr_sz);
-
-        data_ptr = (char *) data + dt_true_lb;
-    } else {
-        /* need copy headers and pack data */
-        send_buf = MPL_malloc(header_size + data_sz, MPL_MEM_OTHER);
-        MPIR_Assert(send_buf);
-        header = send_buf;
-        data_ptr = (char *) send_buf + header_size;
-
-        MPIR_Memcpy(header, &ucx_hdr, sizeof(ucx_hdr));
-        MPIR_Memcpy((char *) header + sizeof(ucx_hdr), am_hdr, am_hdr_sz);
-
-        MPI_Aint actual_pack_bytes;
-        mpi_errno = MPIR_Typerep_pack(data, count, datatype, 0, data_ptr, data_sz,
-                                      &actual_pack_bytes, MPIR_TYPEREP_FLAG_NONE);
-        MPIR_ERR_CHECK(mpi_errno);
-        MPIR_Assert(actual_pack_bytes == data_sz);
-    }
+    void *header;
+    const void *data_ptr;
     ucp_request_param_t param = {
         .op_attr_mask = UCP_OP_ATTR_FIELD_CALLBACK | UCP_OP_ATTR_FIELD_USER_DATA,
         .cb.send = &MPIDI_UCX_am_isend_callback_nbx,
         .user_data = sreq,
     };
+
+    header = MPL_malloc(header_size, MPL_MEM_OTHER);
+    MPIR_Assert(header);
+
+    MPIR_Memcpy(header, &ucx_hdr, sizeof(ucx_hdr));
+    MPIR_Memcpy((char *) header + sizeof(ucx_hdr), am_hdr, am_hdr_sz);
+
+    if (dt_contig) {
+        data_ptr = (char *) data + dt_true_lb;
+    } else {
+        param.op_attr_mask |= UCP_OP_ATTR_FIELD_DATATYPE;
+        param.datatype = dt_ptr->dev.netmod.ucx.ucp_datatype;
+        MPIR_Datatype_ptr_add_ref(dt_ptr);
+        data_ptr = data;
+        data_sz = count;
+    }
     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(send_buf);
+        MPL_free(header);
         MPIDIG_global.origin_cbs[handler_id] (sreq);
         goto fn_exit;
     }
 
-    MPIDI_UCX_AM_SEND_REQUEST(sreq, pack_buffer) = send_buf;
+    MPIDI_UCX_AM_SEND_REQUEST(sreq, pack_buffer) = header;
     MPIDI_UCX_AM_SEND_REQUEST(sreq, handler_id) = handler_id;
     ucp_request_release(ucp_request);