diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 205a11feac3..29d0d6b3387 100644 --- a/src/ucp/core/ucp_context.c +++ b/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 */ diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index aea1d5de6c2..4ff39eb416f 100644 --- a/src/ucp/core/ucp_context.h +++ b/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; diff --git a/src/ucp/core/ucp_request.h b/src/ucp/core/ucp_request.h index 4f4d8720ac4..3a536ae2656 100644 --- a/src/ucp/core/ucp_request.h +++ b/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 */ diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 76548ac1a88..a5281fa74f4 100644 --- a/src/ucp/core/ucp_worker.c +++ b/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); diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 28e5013079f..8d0088c9de4 100644 --- a/src/ucp/core/ucp_worker.h +++ b/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; diff --git a/src/ucp/tag/rndv.c b/src/ucp/tag/rndv.c index 2c7cccf35e9..924bbe722ed 100644 --- a/src/ucp/tag/rndv.c +++ b/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); } diff --git a/src/ucp/tag/tag_send.c b/src/ucp/tag/tag_send.c index 956e168e731..e864385da93 100644 --- a/src/ucp/tag/tag_send.c +++ b/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,