ch4/ofi: use cq_data to carry data_sz if supported

If we have 8-byte cq_data, we can use the higher bits to carry data_sz. This saves a round trip for MPI_{Probe,Mprobe,Iprobe,Improbe}.
pmodels · Oct 21, 2024 · 67978a2 · 67978a2
1 parent 3f856d0
commit 67978a2
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 9 deletions.
diff --git a/src/mpid/ch4/netmod/ofi/ofi_rndv.c b/src/mpid/ch4/netmod/ofi/ofi_rndv.c
@@ -116,18 +116,26 @@ int MPIDI_OFI_peek_rndv_event(int vci, struct fi_cq_tagged_entry *wc, MPIR_Reque
     mpi_errno = rndv_event_common(vci, rreq, &vci_src, &vci_dst);
     MPIR_ERR_CHECK(mpi_errno);
 
-    /* prepare rndv_cts */
-    struct rndv_cts hdr;
-    hdr.rreq = rreq;
-    hdr.am_tag = -1;    /* don't issue am_tag_recv yet */
-    hdr.flag = MPIDI_OFI_CTS_FLAG__PROBE;
+    MPI_Aint data_sz;
+    data_sz = MPIDI_OFI_idata_get_size(wc->data);
 
-    /* send cts */
-    mpi_errno = MPIDI_OFI_send_ack(rreq, &hdr, sizeof(hdr));
-    MPIR_ERR_CHECK(mpi_errno);
+    if (data_sz > 0) {
+        /* complete probe */
+        MPIR_STATUS_SET_COUNT(rreq->status, hdr->data_sz);
+        MPL_atomic_release_store_int(&(MPIDI_OFI_REQUEST(rreq, peek_status)), MPIDI_OFI_PEEK_FOUND);
+    } else {
+        /* ask sender for data_sz */
+        struct rndv_cts hdr;
+        hdr.rreq = rreq;
+        hdr.am_tag = -1;        /* don't issue am_tag_recv yet */
+        hdr.flag = MPIDI_OFI_CTS_FLAG__PROBE;
 
-    MPIR_Request_add_ref(rreq);
+        mpi_errno = MPIDI_OFI_send_ack(rreq, &hdr, sizeof(hdr));
+        MPIR_ERR_CHECK(mpi_errno);
 
+        /* hold the ref until sender replies */
+        MPIR_Request_add_ref(rreq);
+    }
   fn_exit:
     MPIR_FUNC_EXIT;
     return mpi_errno;

diff --git a/src/mpid/ch4/netmod/ofi/ofi_send.h b/src/mpid/ch4/netmod/ofi/ofi_send.h
@@ -630,6 +630,7 @@ MPL_STATIC_INLINE_PREFIX int MPIDI_OFI_send(const void *buf, MPI_Aint count, MPI
         MPIR_ERR_CHECK(mpi_errno);
         /* inject a zero-size message with MPIDI_OFI_RNDV_SEND in match_bits */
         match_bits |= MPIDI_OFI_RNDV_SEND;
+        MPIDI_OFI_idata_set_size(&cq_data, data_sz);    /* optionally use cq_data to carry data_sz */
         mpi_errno = MPIDI_OFI_send_lightweight(NULL, 0, cq_data, dst_rank, tag, comm,
                                                match_bits, addr,
                                                vci_src, vci_dst, sender_nic, receiver_nic);

diff --git a/src/mpid/ch4/netmod/ofi/ofi_types.h b/src/mpid/ch4/netmod/ofi/ofi_types.h
@@ -98,6 +98,23 @@ static inline uint32_t MPIDI_OFI_idata_get_gpuchunk_bits(uint64_t idata)
     return (idata >> MPIDI_OFI_IDATA_GPUCHUNK_OFFSET);
 }
 
+static inline void MPIDI_OFI_idata_set_size(uint64_t * data_field, MPI_Aint data_sz)
+{
+    *data_field &= 0xffffffff;
+    if (MPIDI_OFI_global.cq_data_size == 8 && data_sz <= UINT32_MAX) {
+        *data_field |= (data_sz << 32);
+    }
+}
+
+static inline uint32_t MPIDI_OFI_idata_get_size(uint64_t idata)
+{
+    if (MPIDI_OFI_global.cq_data_size == 8) {
+        return idata >> 32;
+    } else {
+        return 0;
+    }
+}
+
 #define MPIDI_OFI_PROTOCOL_BITS (6)
 /* define protocol bits without MPIDI_OFI_PROTOCOL_SHIFT */
 #define MPIDI_OFI_ACK_SEND_0      1ULL