ch4/shm: add fast path for host buffer

Check if source buffer is on host and choose typerep fast path for H2H.
pmodels · Aug 1, 2024 · c4d1e02 · c4d1e02
1 parent 174033a
commit c4d1e02
Show file tree

Hide file tree

Showing 6 changed files with 47 additions and 12 deletions.
diff --git a/src/mpid/ch4/shm/posix/eager/iqueue/iqueue_send.h b/src/mpid/ch4/shm/posix/eager/iqueue/iqueue_send.h
@@ -95,9 +95,11 @@ MPIDI_POSIX_eager_send(int grank, MPIDI_POSIX_am_header_t * msg_hdr, const void
         cell->type = MPIDI_POSIX_EAGER_IQUEUE_CELL_TYPE_HDR;
         /* send am_hdr if this is the first segment */
         if (is_topo_local) {
-            MPIR_Typerep_copy(payload, am_hdr, am_hdr_sz, MPIR_TYPEREP_FLAG_NONE);
+            MPIR_Typerep_copy(payload, am_hdr, am_hdr_sz,
+                              MPIR_TYPEREP_FLAG_H2H | MPIR_TYPEREP_FLAG_NONE);
         } else {
-            MPIR_Typerep_copy(payload, am_hdr, am_hdr_sz, MPIR_TYPEREP_FLAG_STREAM);
+            MPIR_Typerep_copy(payload, am_hdr, am_hdr_sz,
+                              MPIR_TYPEREP_FLAG_H2H | MPIR_TYPEREP_FLAG_STREAM);
         }
         /* make sure the data region starts at the boundary of MAX_ALIGNMENT */
         payload = payload + resized_am_hdr_sz;
@@ -114,13 +116,16 @@ MPIDI_POSIX_eager_send(int grank, MPIDI_POSIX_am_header_t * msg_hdr, const void
      * not reliable because the derived datatype could have zero block size which contains no
      * data. */
     if (bytes_sent) {
-        if (is_topo_local) {
-            MPIR_Typerep_pack(buf, count, datatype, offset, payload, available, &packed_size,
-                              MPIR_TYPEREP_FLAG_NONE);
-        } else {
-            MPIR_Typerep_pack(buf, count, datatype, offset, payload, available, &packed_size,
-                              MPIR_TYPEREP_FLAG_STREAM);
+        int typerep_flags = MPIR_TYPEREP_NONE;
+        if (msg_hdr == MPIDI_POSIX_AM_TYPE__SHORT_HOST
+            || msg_hdr == MPIDI_POSIX_AM_TYPE__PIPELINE_HOST) {
+            typerep_flags |= MPIR_TYPEREP_FLAG_H2H;
+        }
+        if (!is_topo_local) {
+            typerep_flags |= MPIR_TYPEREP_FLAG_STREAM;
         }
+        MPIR_Typerep_pack(buf, count, datatype, offset, payload, available, &packed_size,
+                          typerep_flags);
         cell->payload_size += packed_size;
         *bytes_sent = packed_size;
     }

diff --git a/src/mpid/ch4/shm/posix/posix_am.h b/src/mpid/ch4/shm/posix/posix_am.h
@@ -11,6 +11,10 @@
 #include "posix_eager.h"
 #include "mpidu_genq.h"
 
+#undef IS_HOST
+#define IS_HOST(attr) \
+    ((attr).type & (MPL_GPU_POINTER_UNREGISTERED_HOST | MPL_GPU_POINTER_REGISTERED_HOST))
+
 MPL_STATIC_INLINE_PREFIX MPI_Aint MPIDI_POSIX_am_eager_limit(void)
 {
     return MPIDI_POSIX_eager_payload_limit() - MAX_ALIGNMENT;
@@ -284,9 +288,17 @@ MPL_STATIC_INLINE_PREFIX int MPIDI_POSIX_do_am_isend(int grank,
 
         msg_hdr_p = msg_hdr;
         if (data_sz + am_hdr_sz <= MPIDI_POSIX_am_eager_limit()) {
-            msg_hdr_p->am_type = MPIDI_POSIX_AM_TYPE__SHORT;
+            if (IS_HOST(MPIDIG_REQUEST(sreq, buf_attr))) {
+                msg_hdr_p->am_type = MPIDI_POSIX_AM_TYPE__SHORT_HOST;
+            } else {
+                msg_hdr_p->am_type = MPIDI_POSIX_AM_TYPE__SHORT;
+            }
         } else {
-            msg_hdr_p->am_type = MPIDI_POSIX_AM_TYPE__PIPELINE;
+            if (IS_HOST(MPIDIG_REQUEST(sreq, buf_attr))) {
+                msg_hdr_p->am_type = MPIDI_POSIX_AM_TYPE__PIPELINE_HOST;
+            } else {
+                msg_hdr_p->am_type = MPIDI_POSIX_AM_TYPE__PIPELINE;
+            }
         }
 
         MPIDIG_am_send_async_init(sreq, datatype, data_sz);

diff --git a/src/mpid/ch4/shm/posix/posix_pre.h b/src/mpid/ch4/shm/posix/posix_pre.h
@@ -22,7 +22,9 @@ typedef enum {
 typedef enum {
     MPIDI_POSIX_AM_TYPE__HDR,
     MPIDI_POSIX_AM_TYPE__SHORT,
-    MPIDI_POSIX_AM_TYPE__PIPELINE
+    MPIDI_POSIX_AM_TYPE__SHORT_HOST,
+    MPIDI_POSIX_AM_TYPE__PIPELINE,
+    MPIDI_POSIX_AM_TYPE__PIPELINE_HOST
 } MPIDI_POSIX_am_type_t;
 
 struct MPIR_Request;

diff --git a/src/mpid/ch4/shm/posix/posix_progress.h b/src/mpid/ch4/shm/posix/posix_progress.h
@@ -57,11 +57,13 @@ MPL_STATIC_INLINE_PREFIX int MPIDI_POSIX_progress_recv(int vci, int *made_progre
         switch (msg_hdr->am_type) {
             case MPIDI_POSIX_AM_TYPE__HDR:
             case MPIDI_POSIX_AM_TYPE__SHORT:
+            case MPIDI_POSIX_AM_TYPE__SHORT_HOST:
                 MPIDIG_global.target_msg_cbs[msg_hdr->handler_id] (am_hdr, payload, payload_left,
                                                                    attr, NULL);
                 MPIDI_POSIX_eager_recv_commit(&transaction);
                 goto fn_exit;
             case MPIDI_POSIX_AM_TYPE__PIPELINE:
+            case MPIDI_POSIX_AM_TYPE__PIPELINE_HOST:
                 MPIDIG_global.target_msg_cbs[msg_hdr->handler_id] (am_hdr, NULL, payload_left,
                                                                    attr | MPIDIG_AM_ATTR__IS_ASYNC,
                                                                    &rreq);

diff --git a/src/mpid/ch4/shm/posix/posix_send.h b/src/mpid/ch4/shm/posix/posix_send.h
@@ -16,6 +16,10 @@
 
 #include "posix_impl.h"
 
+#undef IS_HOST
+#define IS_HOST(attr) \
+    ((attr).type & (MPL_GPU_POINTER_UNREGISTERED_HOST | MPL_GPU_POINTER_REGISTERED_HOST))
+
 #define MPIDI_POSIX_SEND_VSIS(vci_src_, vci_dst_) \
     do { \
         MPIDI_EXPLICIT_VCIS(comm, attr, comm->rank, rank, vci_src_, vci_dst_); \
@@ -52,7 +56,13 @@ MPL_STATIC_INLINE_PREFIX int MPIDI_POSIX_mpi_isend(const void *buf, MPI_Aint cou
             MPIDI_POSIX_am_header_t msg_hdr;
             msg_hdr.handler_id = MPIDIG_SEND;
             msg_hdr.am_hdr_sz = sizeof(MPIDIG_hdr_t);
-            msg_hdr.am_type = MPIDI_POSIX_AM_TYPE__SHORT;
+            MPL_pointer_attr_t attr;
+            MPIR_GPU_query_pointer_attr(buf, &attr);
+            if (IS_HOST(attr)) {
+                msg_hdr.am_type = MPIDI_POSIX_AM_TYPE__SHORT_HOST;
+            } else {
+                msg_hdr.am_type = MPIDI_POSIX_AM_TYPE__SHORT;
+            }
 
             MPIDIG_hdr_t am_hdr;
             am_hdr.src_rank = comm->rank;

diff --git a/src/mpid/ch4/src/mpidig_send_utils.h b/src/mpid/ch4/src/mpidig_send_utils.h
@@ -6,6 +6,10 @@
 #ifndef MPIDIG_SEND_UTILS_H_INCLUDED
 #define MPIDIG_SEND_UTILS_H_INCLUDED
 
+#undef IS_HOST
+#define IS_HOST(attr) \
+    ((attr).type & (MPL_GPU_POINTER_UNREGISTERED_HOST | MPL_GPU_POINTER_REGISTERED_HOST))
+
 /* This file is for supporting routines used for pipelined data send. These routines mainly is for
  * managing the send request counters, completion counters and DT refcount */