Topic/rndv pipeline modify

src/ucp/core/ucp_context.c

@@ -178,6 +178,10 @@ static ucs_config_field_t ucp_config_table[] = {
    "to ucp_init()",
    ucs_offsetof(ucp_config_t, ctx.estimated_num_eps), UCS_CONFIG_TYPE_ULUNITS},
 
+  {"RNDV_FRAG_SIZE", "65536",
+   "RNDV fragment size \n",
+   ucs_offsetof(ucp_config_t, ctx.rndv_frag_size), UCS_CONFIG_TYPE_MEMUNITS},
+
   {NULL}
 };
 
@@ -860,13 +864,6 @@ static ucs_status_t ucp_fill_config(ucp_context_h context,
     ucs_debug("Estimated number of endpoints is %d",
               context->config.est_num_eps);
 
-    if (context->config.ext.rndv_mode == UCP_RNDV_MODE_AUTO) {
-        /* TODO: currently UCP_RNDV_MODE_AUTO == UCP_RNDV_MODE_GET_ZCOPY,
-         * after memory type support is added, will add tru UCP_RNDV_MODE_AUTO
-         * implementation */
-        context->config.ext.rndv_mode = UCP_RNDV_MODE_GET_ZCOPY;
-    }
-
     /* always init MT lock in context even though it is disabled by user,
      * because we need to use context lock to protect ucp_mm_ and ucp_rkey_
      * routines */

src/ucp/core/ucp_context.h

@@ -49,6 +49,8 @@ typedef struct ucp_context_config {
     size_t                                 bcopy_bw;
     /** Segment size in the worker pre-registered memory pool */
     size_t                                 seg_size;
+    /** RNDV pipeline fragment size */
+    size_t                                 rndv_frag_size;
     /** Threshold for using tag matching offload capabilities. Smaller buffers
      *  will not be posted to the transport. */
     size_t                                 tm_thresh;

src/ucp/core/ucp_request.h

@@ -95,7 +95,9 @@ struct ucp_request {
             const void            *buffer;  /* Send buffer */
             ucp_datatype_t        datatype; /* Send type */
             size_t                length;   /* Total length, in bytes */
+            uct_memory_type_t     mem_type; /* Memory type for send buffer */
             ucp_send_callback_t   cb;       /* Completion callback */
+            ucp_ep_h              dummy_ep;
 
             union {
                 ucp_tag_t         tag;      /* Tagged send */

src/ucp/core/ucp_worker.c

@@ -981,8 +981,18 @@ static ucs_status_t ucp_worker_init_mpools(ucp_worker_h worker)
         goto err_release_am_mpool;
     }
 
+    status = ucs_mpool_init(&worker->rndv_frag_mp, 0,
+                            context->config.ext.rndv_frag_size,
+                            0, 128, 128, UINT_MAX,
+                            &ucp_reg_mpool_ops, "ucp_rndv_frags");
+    if (status != UCS_OK) {
+        goto err_release_reg_mpool;
+    }
+
     return UCS_OK;
 
+err_release_reg_mpool:
+    ucs_mpool_cleanup(&worker->reg_mp, 0);
 err_release_am_mpool:
     ucs_mpool_cleanup(&worker->am_mp, 0);
 out:
@@ -1200,6 +1210,7 @@ void ucp_worker_destroy(ucp_worker_h worker)
     ucp_worker_destroy_eps(worker);
     ucs_mpool_cleanup(&worker->am_mp, 1);
     ucs_mpool_cleanup(&worker->reg_mp, 1);
+    ucs_mpool_cleanup(&worker->rndv_frag_mp, 1);
     ucp_worker_close_ifaces(worker);
     ucp_tag_match_cleanup(&worker->tm);
     ucp_worker_wakeup_cleanup(worker);

src/ucp/core/ucp_worker.h

@@ -164,6 +164,7 @@ typedef struct ucp_worker {
     uint64_t                      uuid;          /* Unique ID for wireup */
     uct_worker_h                  uct;           /* UCT worker handle */
     ucs_mpool_t                   req_mp;        /* Memory pool for requests */
+    ucs_mpool_t                   rndv_frag_mp;  /* Memory pool for RNDV fragments */
     uint64_t                      atomic_tls;    /* Which resources can be used for atomics */
 
     int                           inprogress;

src/ucp/tag/rndv.c

@@ -19,14 +19,21 @@ size_t ucp_tag_rndv_rts_pack(void *dest, void *arg)
     ucp_rndv_rts_hdr_t *rndv_rts_hdr = dest;
     ucp_worker_h worker              = sreq->send.ep->worker;
     ssize_t packed_rkey_size;
+    ucp_rndv_mode_t rndv_mode;
 
     rndv_rts_hdr->super.tag        = sreq->send.tag;
     rndv_rts_hdr->sreq.reqptr      = (uintptr_t)sreq;
     rndv_rts_hdr->sreq.sender_uuid = worker->uuid;
     rndv_rts_hdr->size             = sreq->send.length;
 
+    rndv_mode = worker->context->config.ext.rndv_mode;
+
     /* Pack remote keys (which can be empty list) */
-    if (UCP_DT_IS_CONTIG(sreq->send.datatype)) {
+    if (UCP_DT_IS_CONTIG(sreq->send.datatype) &&
+        (rndv_mode == UCP_RNDV_MODE_GET_ZCOPY ||
+         (rndv_mode == UCP_RNDV_MODE_AUTO &&
+          sreq->send.mem_type == UCT_MD_MEM_TYPE_HOST))) {
+        /* pack rkey, ask target to do get_zcopy */
         rndv_rts_hdr->address = (uintptr_t)sreq->send.buffer;
         packed_rkey_size = ucp_rkey_pack_uct(worker->context,
                                              sreq->send.state.dt.dt.contig.md_map,
@@ -37,6 +44,12 @@ size_t ucp_tag_rndv_rts_pack(void *dest, void *arg)
                       ucs_status_string(packed_rkey_size));
         }
     } else {
+        if (UCP_DT_IS_CONTIG(sreq->send.datatype) &&
+            rndv_mode == UCP_RNDV_MODE_AUTO &&
+            sreq->send.mem_type != UCT_MD_MEM_TYPE_HOST) {
+            ucp_ep_create_stub(sreq->send.ep->worker, sreq->send.ep->worker->uuid,
+                               NULL, "", "for rndv pipeline", &sreq->send.dummy_ep);
+        }
         rndv_rts_hdr->address = 0;
         packed_rkey_size      = 0;
     }
@@ -722,6 +735,84 @@ static ucs_status_t ucp_rndv_progress_am_zcopy_multi(uct_pending_req_t *self)
                                  ucp_rndv_am_zcopy_send_req_complete);
 }
 
+UCS_PROFILE_FUNC_VOID(ucp_rndv_frag_put_completion, (self, status),
+                      uct_completion_t *self, ucs_status_t status)
+{
+    ucp_request_t *frag_req = ucs_container_of(self, ucp_request_t,
+                                               send.state.uct_comp);
+    ucp_request_t *sreq = frag_req->send.rndv_get.rreq;
+
+    ucp_request_send_buffer_dereg(frag_req);
+
+    ucs_mpool_put((void *)frag_req->send.buffer);
+
+    sreq->send.state.uct_comp.count--;
+    if (0 == sreq->send.state.uct_comp.count) {
+        ucp_rndv_send_atp(frag_req, sreq->send.rndv_put.remote_request);
+        ucp_ep_destroy(sreq->send.dummy_ep);
+        ucp_request_complete_send(sreq, status);
+    } else {
+        ucp_request_put(frag_req);
+    }
+}
+
+UCS_PROFILE_FUNC_VOID(ucp_rndv_frag_get_completion, (self, status),
+                      uct_completion_t *self, ucs_status_t status)
+{
+    ucp_request_t *frag_req = ucs_container_of(self, ucp_request_t, send.state.uct_comp);
+    ucp_request_t *sreq = frag_req->send.rndv_get.rreq;
+
+    frag_req->send.ep                   = sreq->send.ep;
+    frag_req->send.uct.func             = ucp_rndv_progress_rma_put_zcopy;
+    frag_req->send.datatype             = ucp_dt_make_contig(1);
+    ucp_request_send_state_reset(frag_req, ucp_rndv_frag_put_completion,
+                                 UCP_REQUEST_SEND_PROTO_RNDV_PUT);
+    frag_req->send.lane                 = ucp_rkey_get_rma_bw_lane(sreq->send.rndv_put.rkey, sreq->send.ep,
+                                                                   &sreq->send.rndv_put.uct_rkey);
+
+    frag_req->send.buffer = (void *)frag_req->send.rndv_get.remote_address;
+    ucp_request_send_buffer_reg(frag_req, sreq->send.lane);
+
+    frag_req->send.rndv_put.remote_address    = sreq->send.rndv_put.remote_address + frag_req->send.state.dt.offset;
+    frag_req->send.rndv_put.rkey              = sreq->send.rndv_put.rkey;
+    frag_req->send.rndv_put.uct_rkey          = sreq->send.rndv_put.uct_rkey;
+
+    ucp_request_send(frag_req);
+}
+
+static void ucp_rndv_pipeline(ucp_request_t *sreq, ucp_rndv_rtr_hdr_t *rndv_rtr_hdr)
+{
+    int num_frags, i;
+    size_t frag_size, length, offset;
+    ucp_request_t *frag_req;
+    ucp_ep_h ep = sreq->send.ep;
+
+    frag_size = ep->worker->context->config.ext.rndv_frag_size;
+    num_frags = (sreq->send.length + frag_size - 1) / frag_size;
+    sreq->send.state.uct_comp.count = num_frags;
+
+    offset = 0;
+    for (i = 0; i < num_frags; i++) {
+        length = (i != (num_frags - 1)) ? frag_size : (sreq->send.length - offset);
+
+        frag_req = ucp_request_get(ep->worker);
+
+        frag_req->send.uct.func = ucp_rndv_progress_rma_get_zcopy;
+        frag_req->send.buffer = sreq->send.buffer + offset;
+        frag_req->send.datatype = sreq->send.datatype;
+        frag_req->send.length = length;
+        frag_req->send.lane = ucp_rkey_get_rma_bw_lane(sreq->send.rndv_put.rkey, sreq->send.dummy_ep,
+                                                   &sreq->send.rndv_put.uct_rkey);
+        frag_req->send.rndv_get.remote_address = (uintptr_t) ucs_mpool_get(&ep->worker->rndv_frag_mp);
+        frag_req->send.rndv_get.rreq = sreq;
+        frag_req->send.state.dt.offset = offset;
+        ucp_request_send_state_reset(frag_req, ucp_rndv_frag_get_completion,
+                                     UCP_REQUEST_SEND_PROTO_RNDV_GET);
+        ucp_request_send(frag_req);
+        offset += length;
+    }
+}
+
 UCS_PROFILE_FUNC(ucs_status_t, ucp_rndv_rtr_handler,
                  (arg, data, length, flags),
                  void *arg, void *data, size_t length, unsigned flags)
@@ -754,10 +845,15 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_rndv_rtr_handler,
         if (sreq->send.lane != UCP_NULL_LANE) {
             ucp_request_send_state_reset(sreq, ucp_rndv_put_completion,
                                          UCP_REQUEST_SEND_PROTO_RNDV_PUT);
-            sreq->send.uct.func                = ucp_rndv_progress_rma_put_zcopy;
-            sreq->send.rndv_put.remote_request = rndv_rtr_hdr->rreq_ptr;
-            sreq->send.rndv_put.remote_address = rndv_rtr_hdr->address;
-            goto out_send;
+            if (sreq->send.mem_type == UCT_MD_MEM_TYPE_HOST) {
+                sreq->send.uct.func                = ucp_rndv_progress_rma_put_zcopy;
+                sreq->send.rndv_put.remote_request = rndv_rtr_hdr->rreq_ptr;
+                sreq->send.rndv_put.remote_address = rndv_rtr_hdr->address;
+                goto out_send;
+            } else {
+                ucp_rndv_pipeline(sreq, rndv_rtr_hdr);
+                return UCS_OK;
+            }
         } else {
             ucp_rkey_destroy(sreq->send.rndv_put.rkey);
         }

src/ucp/tag/tag_send.c

@@ -55,6 +55,7 @@ ucp_tag_send_req(ucp_request_t *req, size_t count,
     size_t zcopy_thresh = ucp_proto_get_zcopy_threshold(req, msg_config, count,
                                                         rndv_thresh);
     ssize_t max_short   = ucp_proto_get_short_max(req, msg_config);
+    ucp_rndv_mode_t rndv_mode;
     ucs_status_t status;
 
     ucs_trace_req("select tag request(%p) progress algorithm datatype=%lx "
@@ -63,6 +64,12 @@ ucp_tag_send_req(ucp_request_t *req, size_t count,
                   req, req->send.datatype, req->send.buffer, req->send.length,
                   max_short, rndv_thresh, zcopy_thresh);
 
+    rndv_mode = req->send.ep->worker->context->config.ext.rndv_mode;
+    if (UCP_DT_IS_CONTIG(req->send.datatype) && rndv_mode == UCP_RNDV_MODE_AUTO) {
+        ucp_memory_type_detect_mds(req->send.ep->worker->context, (void *)req->send.buffer,
+                                   req->send.length, &req->send.mem_type);
+    }
+
     status = ucp_request_send_start(req, max_short, zcopy_thresh, seg_size,
                                     rndv_thresh, proto);
     if (ucs_unlikely(status != UCS_OK)) {
@@ -119,6 +126,7 @@ ucp_tag_send_req_init(ucp_request_t* req, ucp_ep_h ep, const void* buffer,
                                            req->send.buffer,
                                            &req->send.state.dt);
     req->send.lane         = ucp_ep_config(ep)->tag.lane;
+    req->send.mem_type     = UCT_MD_MEM_TYPE_HOST;
 }
 
 UCS_PROFILE_FUNC(ucs_status_ptr_t, ucp_tag_send_nb,