diff --git a/src/ucp/Makefile.am b/src/ucp/Makefile.am index 5193bc942ca..993fd777267 100644 --- a/src/ucp/Makefile.am +++ b/src/ucp/Makefile.am @@ -71,6 +71,8 @@ noinst_HEADERS = \ proto/proto.h \ rma/rma.h \ rma/rma.inl \ + rma/rma_bw.h \ + rma/rma_bw.inl \ rma/rma_rndv.h \ rndv/proto_rndv.h \ rndv/proto_rndv.inl \ @@ -149,6 +151,7 @@ libucp_la_SOURCES = \ rma/get_offload.c \ rma/put_am.c \ rma/put_offload.c \ + rma/rma_bw.c \ rma/rma_rndv.c \ rma/rma_send.c \ rma/rma_sw.c \ diff --git a/src/ucp/core/ucp_context.c b/src/ucp/core/ucp_context.c index 6f13d421423..d2ea253f613 100644 --- a/src/ucp/core/ucp_context.c +++ b/src/ucp/core/ucp_context.c @@ -586,6 +586,12 @@ static ucs_config_field_t ucp_context_config_table[] = { " 'n' : Select RMA/AMO lanes according to performance charasteristics", ucs_offsetof(ucp_context_config_t, prefer_offload), UCS_CONFIG_TYPE_BOOL}, + {"RMA_BW_MEASURE", "n", + "Sample completed payload throughput of multi-lane PUT/GET zcopy operations. " + "This experimental option records measurements but does not change lane " + "weights.", + ucs_offsetof(ucp_context_config_t, rma_bw_measure), UCS_CONFIG_TYPE_BOOL}, + {"PROTO_OVERHEAD", "single:5ns,multi:10ns,rndv_offload:40ns,rndv_rtr:40ns," "rndv_rts:275ns,sw:40ns,rkey_ptr:0", "Protocol overhead", 0, diff --git a/src/ucp/core/ucp_context.h b/src/ucp/core/ucp_context.h index 7d19b9bf163..49c8a473c71 100644 --- a/src/ucp/core/ucp_context.h +++ b/src/ucp/core/ucp_context.h @@ -230,6 +230,8 @@ typedef struct ucp_context_config { int reg_nb_fallback; /** Prefer native RMA transports for RMA/AMO protocols */ int prefer_offload; + /** Sample completed multi-rail RMA zcopy operations */ + int rma_bw_measure; /** RMA zcopy segment size */ size_t rma_zcopy_max_seg_size; /** Enable global VA MR */ diff --git a/src/ucp/core/ucp_request.c b/src/ucp/core/ucp_request.c index 49a2290cda6..34a9b9f73e6 100644 --- a/src/ucp/core/ucp_request.c +++ b/src/ucp/core/ucp_request.c @@ -55,6 +55,7 @@ static const char *ucp_request_flag_names[] = { [ucs_ilog2(UCP_REQUEST_FLAG_RNDV_GET_REQ)] = "rndv_get_req", [ucs_ilog2(UCP_REQUEST_FLAG_RNDV_FLUSH)] = "rndv_flush", [ucs_ilog2(UCP_REQUEST_FLAG_RNDV_START_FLUSH)] = "rndv_start_flush", + [ucs_ilog2(UCP_REQUEST_FLAG_RMA_BW_SAMPLE)] = "rma_bw_sample", }; static ucs_memory_type_t ucp_request_get_mem_type(ucp_request_t *req) diff --git a/src/ucp/core/ucp_request.h b/src/ucp/core/ucp_request.h index c0674d0fab3..b994184635e 100644 --- a/src/ucp/core/ucp_request.h +++ b/src/ucp/core/ucp_request.h @@ -69,7 +69,8 @@ enum { UCP_REQUEST_FLAG_RNDV_SEND_INTERNAL = UCS_BIT(26), UCP_REQUEST_FLAG_RNDV_GET_REQ = UCS_BIT(27), UCP_REQUEST_FLAG_RNDV_FLUSH = UCS_BIT(28), - UCP_REQUEST_FLAG_RNDV_START_FLUSH = UCS_BIT(29) + UCP_REQUEST_FLAG_RNDV_START_FLUSH = UCS_BIT(29), + UCP_REQUEST_FLAG_RMA_BW_SAMPLE = UCS_BIT(30) }; @@ -245,10 +246,14 @@ struct ucp_request { uint64_t remote_addr; /* Remote address */ ucp_rkey_h rkey; /* Remote memory key */ - struct { - const uint64_t *remote_addrs; - ucp_rkey_h const *rkeys; - } sgl; + union { + struct { + const uint64_t *remote_addrs; + ucp_rkey_h const *rkeys; + } sgl; + /* Used only by sampled contiguous PUT/GET zcopy */ + ucp_rma_bw_sample_t *bw_sample; + }; } rma; struct { diff --git a/src/ucp/core/ucp_types.h b/src/ucp/core/ucp_types.h index 41f13675aad..543c8b98721 100644 --- a/src/ucp/core/ucp_types.h +++ b/src/ucp/core/ucp_types.h @@ -64,6 +64,7 @@ typedef uint16_t ucp_worker_cfg_index_t; /* Forward declarations */ typedef struct ucp_request ucp_request_t; +typedef struct ucp_rma_bw_sample ucp_rma_bw_sample_t; typedef struct ucp_recv_desc ucp_recv_desc_t; typedef struct ucp_address_iface_attr ucp_address_iface_attr_t; typedef struct ucp_address_entry ucp_address_entry_t; diff --git a/src/ucp/core/ucp_worker.c b/src/ucp/core/ucp_worker.c index 27909730a3f..2c3b9d3d3ea 100644 --- a/src/ucp/core/ucp_worker.c +++ b/src/ucp/core/ucp_worker.c @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -2884,6 +2885,16 @@ ucs_status_t ucp_worker_create(ucp_context_h context, goto err_tag_match_cleanup; } + if (context->config.ext.rma_bw_measure) { + worker->rma_bw_samples = ucs_calloc(UCP_RMA_BW_MAX_ACTIVE, + sizeof(*worker->rma_bw_samples), + "rma_bw_samples"); + if (worker->rma_bw_samples == NULL) { + status = UCS_ERR_NO_MEMORY; + goto err_am_cleanup; + } + } + /* Select atomic resources */ ucp_worker_init_atomic_tls(worker); @@ -2910,6 +2921,7 @@ ucs_status_t ucp_worker_create(ucp_context_h context, return UCS_OK; err_am_cleanup: + ucs_free(worker->rma_bw_samples); ucp_am_cleanup(worker); err_tag_match_cleanup: ucp_tag_match_cleanup(&worker->tm); @@ -3186,6 +3198,7 @@ void ucp_worker_destroy(ucp_worker_h worker) ucs_conn_match_cleanup(&worker->conn_match_ctx); ucp_worker_wakeup_cleanup(worker); uct_worker_destroy(worker->uct); + ucs_free(worker->rma_bw_samples); ucs_async_context_cleanup(&worker->async); UCS_STATS_NODE_FREE(worker->tm_offload_stats); UCS_STATS_NODE_FREE(worker->stats); diff --git a/src/ucp/core/ucp_worker.h b/src/ucp/core/ucp_worker.h index 32ae9d1cac6..d8778dc9491 100644 --- a/src/ucp/core/ucp_worker.h +++ b/src/ucp/core/ucp_worker.h @@ -314,6 +314,8 @@ typedef struct ucp_worker { uct_worker_h uct; /* UCT worker handle */ ucs_mpool_t req_mp; /* Memory pool for requests */ ucs_mpool_t rkey_mp; /* Pool for small memory keys */ + ucp_rma_bw_sample_t *rma_bw_samples; /* Sample pool */ + ucs_time_t rma_bw_next_sample; /* Next sample admission */ ucp_tl_bitmap_t atomic_tls; /* Which resources can be used for atomics */ int inprogress; diff --git a/src/ucp/rma/get_offload.c b/src/ucp/rma/get_offload.c index 0830039f2ef..e9ab2d3f9e3 100644 --- a/src/ucp/rma/get_offload.c +++ b/src/ucp/rma/get_offload.c @@ -10,6 +10,7 @@ #endif #include "rma.inl" +#include "rma_bw.inl" #include #include @@ -157,16 +158,18 @@ ucp_proto_t ucp_get_offload_bcopy_proto = { }; static UCS_F_ALWAYS_INLINE ucs_status_t -ucp_proto_get_offload_zcopy_send_func(ucp_request_t *req, - const ucp_proto_multi_lane_priv_t *lpriv, - ucp_datatype_iter_t *next_iter, - ucp_lane_index_t *lane_shift) +ucp_proto_get_offload_zcopy_send_common( + ucp_request_t *req, const ucp_proto_multi_lane_priv_t *lpriv, + ucp_datatype_iter_t *next_iter, int measure) { uct_rkey_t tl_rkey = ucp_rkey_get_tl_rkey(req->send.rma.rkey, lpriv->super.rkey_index); - size_t offset = req->send.state.dt_iter.offset; + size_t offset = req->send.state.dt_iter.offset; + uct_completion_t *comp = &req->send.state.uct_comp; const ucp_proto_multi_priv_t *mpriv; + ucp_rma_bw_frag_t *frag; uct_iov_t iov; + ucs_status_t status; ucp_datatype_iter_next_iov(&req->send.state.dt_iter, ucp_proto_multi_max_payload(req, lpriv, 0), @@ -176,18 +179,60 @@ ucp_proto_get_offload_zcopy_send_func(ucp_request_t *req, mpriv = req->send.proto_config->priv; ucp_proto_common_zcopy_adjust_min_frag(req, mpriv->min_frag, iov.length, &iov, 1, &offset); - return uct_ep_get_zcopy(ucp_ep_get_lane(req->send.ep, lpriv->super.lane), - &iov, 1, req->send.rma.remote_addr + offset, - tl_rkey, &req->send.state.uct_comp); + if (measure) { + comp = ucp_rma_bw_frag_start(req, req->send.multi_lane_idx, &frag); + } + + status = uct_ep_get_zcopy(ucp_ep_get_lane(req->send.ep, lpriv->super.lane), + &iov, 1, req->send.rma.remote_addr + offset, + tl_rkey, comp); + if (measure) { + ucp_rma_bw_frag_posted(frag, iov.length, status); + } + + return status; +} + +static UCS_F_ALWAYS_INLINE ucs_status_t +ucp_proto_get_offload_zcopy_send_func(ucp_request_t *req, + const ucp_proto_multi_lane_priv_t *lpriv, + ucp_datatype_iter_t *next_iter, + ucp_lane_index_t *lane_shift) +{ + return ucp_proto_get_offload_zcopy_send_common(req, lpriv, next_iter, 0); +} + +static UCS_F_ALWAYS_INLINE ucs_status_t +ucp_proto_get_offload_zcopy_sampled_send_func( + ucp_request_t *req, const ucp_proto_multi_lane_priv_t *lpriv, + ucp_datatype_iter_t *next_iter, ucp_lane_index_t *lane_shift) +{ + return ucp_proto_get_offload_zcopy_send_common(req, lpriv, next_iter, 1); } static ucs_status_t ucp_proto_get_offload_zcopy_progress(uct_pending_req_t *self) { ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); + /* coverity[tainted_data_downcast] */ + const ucp_proto_multi_priv_t *mpriv = req->send.proto_config->priv; + + if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { + ucp_rma_bw_sample_try_start(req, mpriv->num_lanes); + } + + if (ucs_unlikely(req->flags & UCP_REQUEST_FLAG_RMA_BW_SAMPLE)) { + /* coverity[tainted_data_downcast] */ + return ucp_proto_multi_zcopy_progress( + req, mpriv, ucp_proto_multi_rma_init_func, + UCT_MD_MEM_ACCESS_LOCAL_WRITE, UCP_DT_MASK_CONTIG_IOV, + ucp_proto_get_offload_zcopy_sampled_send_func, + ucp_request_invoke_uct_completion_success, + ucp_rma_bw_sample_complete); + } /* coverity[tainted_data_downcast] */ return ucp_proto_multi_zcopy_progress( - req, req->send.proto_config->priv, ucp_proto_multi_rma_init_func, + req, mpriv, ucp_proto_multi_rma_init_func, UCT_MD_MEM_ACCESS_LOCAL_WRITE, UCP_DT_MASK_CONTIG_IOV, ucp_proto_get_offload_zcopy_send_func, ucp_request_invoke_uct_completion_success, @@ -245,6 +290,7 @@ ucp_proto_get_offload_zcopy_probe(const ucp_proto_init_params_t *init_params) static ucs_status_t ucp_proto_get_offload_zcopy_reset(ucp_request_t *req) { + ucp_rma_bw_sample_detach(req); ucp_proto_get_offload_reset(req); return ucp_proto_offload_zcopy_reset(req); } diff --git a/src/ucp/rma/put_offload.c b/src/ucp/rma/put_offload.c index 6311ef702ee..ea3199a1cea 100644 --- a/src/ucp/rma/put_offload.c +++ b/src/ucp/rma/put_offload.c @@ -10,6 +10,7 @@ #include "rma.h" #include "rma.inl" +#include "rma_bw.inl" #include #include @@ -245,10 +246,9 @@ ucp_proto_t ucp_put_offload_bcopy_proto = { }; static UCS_F_ALWAYS_INLINE ucs_status_t -ucp_proto_put_offload_zcopy_send_func(ucp_request_t *req, - const ucp_proto_multi_lane_priv_t *lpriv, - ucp_datatype_iter_t *next_iter, - ucp_lane_index_t *lane_shift) +ucp_proto_put_offload_zcopy_send_common( + ucp_request_t *req, const ucp_proto_multi_lane_priv_t *lpriv, + ucp_datatype_iter_t *next_iter, int measure) { ucp_ep_h ep = req->send.ep; uct_ep_h uct_ep = ucp_ep_get_lane(ep, lpriv->super.lane); @@ -256,6 +256,8 @@ ucp_proto_put_offload_zcopy_send_func(ucp_request_t *req, req->send.state.dt_iter.offset; uct_rkey_t tl_rkey = ucp_rkey_get_tl_rkey(req->send.rma.rkey, lpriv->super.rkey_index); + uct_completion_t *comp = &req->send.state.uct_comp; + ucp_rma_bw_frag_t *frag; uct_iov_t iov; ucs_status_t status; @@ -263,8 +265,15 @@ ucp_proto_put_offload_zcopy_send_func(ucp_request_t *req, ucp_proto_multi_max_payload(req, lpriv, 0), lpriv->super.md_index, UCP_DT_MASK_CONTIG_IOV, next_iter, &iov, 1); - status = uct_ep_put_zcopy(uct_ep, &iov, 1, address, tl_rkey, - &req->send.state.uct_comp); + if (measure) { + comp = ucp_rma_bw_frag_start(req, req->send.multi_lane_idx, &frag); + } + + status = uct_ep_put_zcopy(uct_ep, &iov, 1, address, tl_rkey, comp); + if (measure) { + ucp_rma_bw_frag_posted(frag, iov.length, status); + } + if (!UCS_STATUS_IS_ERR(status)) { ucp_proto_put_offload_update_remote_flush(ep, lpriv->flush_sys_dev_mask, tl_rkey, uct_ep, address); @@ -273,14 +282,47 @@ ucp_proto_put_offload_zcopy_send_func(ucp_request_t *req, return status; } +static UCS_F_ALWAYS_INLINE ucs_status_t +ucp_proto_put_offload_zcopy_send_func(ucp_request_t *req, + const ucp_proto_multi_lane_priv_t *lpriv, + ucp_datatype_iter_t *next_iter, + ucp_lane_index_t *lane_shift) +{ + return ucp_proto_put_offload_zcopy_send_common(req, lpriv, next_iter, 0); +} + +static UCS_F_ALWAYS_INLINE ucs_status_t +ucp_proto_put_offload_zcopy_sampled_send_func( + ucp_request_t *req, const ucp_proto_multi_lane_priv_t *lpriv, + ucp_datatype_iter_t *next_iter, ucp_lane_index_t *lane_shift) +{ + return ucp_proto_put_offload_zcopy_send_common(req, lpriv, next_iter, 1); +} + static ucs_status_t ucp_proto_put_offload_zcopy_progress(uct_pending_req_t *self) { ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); + /* coverity[tainted_data_downcast] */ + const ucp_proto_multi_priv_t *mpriv = req->send.proto_config->priv; + + if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { + ucp_rma_bw_sample_try_start(req, mpriv->num_lanes); + } + + if (ucs_unlikely(req->flags & UCP_REQUEST_FLAG_RMA_BW_SAMPLE)) { + /* coverity[tainted_data_downcast] */ + return ucp_proto_multi_zcopy_progress( + req, mpriv, ucp_proto_multi_rma_init_func, + UCT_MD_MEM_ACCESS_LOCAL_READ, UCP_DT_MASK_CONTIG_IOV, + ucp_proto_put_offload_zcopy_sampled_send_func, + ucp_request_invoke_uct_completion_success, + ucp_rma_bw_sample_complete); + } /* coverity[tainted_data_downcast] */ return ucp_proto_multi_zcopy_progress( - req, req->send.proto_config->priv, ucp_proto_multi_rma_init_func, + req, mpriv, ucp_proto_multi_rma_init_func, UCT_MD_MEM_ACCESS_LOCAL_READ, UCP_DT_MASK_CONTIG_IOV, ucp_proto_put_offload_zcopy_send_func, ucp_request_invoke_uct_completion_success, @@ -345,6 +387,12 @@ ucp_proto_put_offload_zcopy_probe(const ucp_proto_init_params_t *init_params) init_params, 30, UCP_PROTO_COMMON_INIT_FLAG_FAILOVER, 0); } +static ucs_status_t ucp_proto_put_offload_zcopy_reset(ucp_request_t *req) +{ + ucp_rma_bw_sample_detach(req); + return ucp_proto_offload_zcopy_reset(req); +} + ucp_proto_t ucp_put_offload_zcopy_proto = { .name = "put/offload/zcopy", .desc = UCP_PROTO_ZCOPY_DESC, @@ -355,7 +403,7 @@ ucp_proto_t ucp_put_offload_zcopy_proto = { .query = ucp_proto_multi_query, .progress = {ucp_proto_put_offload_zcopy_progress}, .abort = ucp_proto_request_zcopy_abort, - .reset = ucp_proto_offload_zcopy_reset + .reset = ucp_proto_put_offload_zcopy_reset }; static void diff --git a/src/ucp/rma/rma_bw.c b/src/ucp/rma/rma_bw.c new file mode 100644 index 00000000000..02d9f36fb3f --- /dev/null +++ b/src/ucp/rma/rma_bw.c @@ -0,0 +1,122 @@ +/** + * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2026. ALL RIGHTS RESERVED. + * + * See file LICENSE for terms. + */ + +#ifdef HAVE_CONFIG_H +# include "config.h" +#endif + +#include "rma_bw.inl" + +#include +#include + +void ucp_rma_bw_sample_start(ucp_request_t *req, + ucp_lane_index_t num_lanes) +{ + ucp_worker_h worker = req->send.ep->worker; + ucp_rma_bw_sample_t *sample; + ucs_time_t now; + unsigned i; + + if (ucs_likely(!worker->context->config.ext.rma_bw_measure) || + (num_lanes < 2) || + (num_lanes > UCP_MAX_LANES) || + (req->send.state.dt_iter.length < + UCP_RMA_BW_MIN_LANE_LENGTH * num_lanes)) { + return; + } + + now = ucs_get_time(); + if (now < worker->rma_bw_next_sample) { + return; + } + + for (i = 0; i < UCP_RMA_BW_MAX_ACTIVE; ++i) { + sample = &worker->rma_bw_samples[i]; + if (ucp_rma_bw_sample_is_free(sample)) { + break; + } + } + if (i == UCP_RMA_BW_MAX_ACTIVE) { + return; + } + + memset(sample, 0, sizeof(*sample)); + sample->req = req; + sample->num_lanes = num_lanes; + req->send.rma.bw_sample = sample; + req->flags |= UCP_REQUEST_FLAG_RMA_BW_SAMPLE; + worker->rma_bw_next_sample = now + + ucs_time_from_sec(UCP_RMA_BW_SAMPLE_INTERVAL); +} + +void ucp_rma_bw_sample_complete(uct_completion_t *comp) +{ + ucp_request_t *req = ucs_container_of(comp, ucp_request_t, + send.state.uct_comp); + ucp_rma_bw_sample_t *sample = ucp_rma_bw_sample_get(req); + ucp_rma_bw_lane_t *lane; + unsigned i; + + if ((comp->status == UCS_OK) && !sample->invalid) { + for (i = 0; i < sample->num_lanes; ++i) { + lane = &sample->lanes[i]; + if ((lane->bytes == 0) || (lane->num_async == 0) || + (lane->last_comp <= lane->first_post)) { + break; + } + } + + if (i == sample->num_lanes) { + for (i = 0; i < sample->num_lanes; ++i) { + ucp_trace_req(req, "rma bw sample lane %u/%u: %zu bytes / " + "%.3f us", i, sample->num_lanes, + sample->lanes[i].bytes, + ucs_time_to_sec(sample->lanes[i].last_comp - + sample->lanes[i].first_post) * + 1e6); + } + } + } + + req->flags &= ~UCP_REQUEST_FLAG_RMA_BW_SAMPLE; + req->send.rma.bw_sample = NULL; + sample->req = NULL; + ucs_assert(sample->pending == 0); + ucp_proto_request_zcopy_complete(req, comp->status); +} + +void ucp_rma_bw_frag_complete(uct_completion_t *comp) +{ + ucp_rma_bw_frag_t *frag = ucs_container_of(comp, ucp_rma_bw_frag_t, comp); + ucp_rma_bw_sample_t *sample = frag->sample; + ucp_rma_bw_lane_t *lane = &sample->lanes[frag->lane_idx]; + + lane->last_comp = ucs_get_time(); + sample->invalid |= (comp->status != UCS_OK); + ucs_assert(sample->pending > 0); + --sample->pending; + if (sample->req == NULL) { + return; + } + + ucp_invoke_uct_completion(&sample->req->send.state.uct_comp, comp->status); +} + +void ucp_rma_bw_sample_detach(ucp_request_t *req) +{ + ucp_rma_bw_sample_t *sample = ucp_rma_bw_sample_get(req); + + if (sample == NULL) { + return; + } + + sample->req = NULL; + /* Overflow fragments may still use the request completion after reset. */ + req->send.state.uct_comp.func = ucp_proto_request_zcopy_completion; + req->flags &= ~UCP_REQUEST_FLAG_RMA_BW_SAMPLE; + req->send.rma.bw_sample = NULL; +} diff --git a/src/ucp/rma/rma_bw.h b/src/ucp/rma/rma_bw.h new file mode 100644 index 00000000000..e722c640a1d --- /dev/null +++ b/src/ucp/rma/rma_bw.h @@ -0,0 +1,48 @@ +/** + * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2026. ALL RIGHTS RESERVED. + * + * See file LICENSE for terms. + */ + +#ifndef UCP_RMA_BW_H_ +#define UCP_RMA_BW_H_ + +#include + +/* The first measurement stage records completed payload throughput only. */ +#define UCP_RMA_BW_MAX_FRAGS 128 +#define UCP_RMA_BW_MAX_ACTIVE 4 +#define UCP_RMA_BW_SAMPLE_INTERVAL 1.0 +#define UCP_RMA_BW_MIN_LANE_LENGTH (128 * UCS_KBYTE) + +typedef struct { + uct_completion_t comp; + ucp_rma_bw_sample_t *sample; + ucs_time_t post_time; + uint8_t lane_idx; +} ucp_rma_bw_frag_t; + +typedef struct { + size_t bytes; + ucs_time_t first_post; + ucs_time_t last_comp; + unsigned num_frags; + unsigned num_async; +} ucp_rma_bw_lane_t; + +struct ucp_rma_bw_sample { + ucp_request_t *req; + unsigned num_frags; + unsigned pending; + unsigned invalid; + ucp_lane_index_t num_lanes; + ucp_rma_bw_lane_t lanes[UCP_MAX_LANES]; + ucp_rma_bw_frag_t frags[UCP_RMA_BW_MAX_FRAGS]; +}; + +void ucp_rma_bw_sample_start(ucp_request_t *req, ucp_lane_index_t num_lanes); +void ucp_rma_bw_sample_complete(uct_completion_t *comp); +void ucp_rma_bw_frag_complete(uct_completion_t *comp); +void ucp_rma_bw_sample_detach(ucp_request_t *req); + +#endif diff --git a/src/ucp/rma/rma_bw.inl b/src/ucp/rma/rma_bw.inl new file mode 100644 index 00000000000..a9f7a6cb2cb --- /dev/null +++ b/src/ucp/rma/rma_bw.inl @@ -0,0 +1,97 @@ +/** + * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2026. ALL RIGHTS RESERVED. + * + * See file LICENSE for terms. + */ + +#ifndef UCP_RMA_BW_INL_ +#define UCP_RMA_BW_INL_ + +#include "rma_bw.h" + +/* A detached request keeps its slot until every UCT fragment completes. */ +static UCS_F_ALWAYS_INLINE int +ucp_rma_bw_sample_is_free(const ucp_rma_bw_sample_t *sample) +{ + return (sample->req == NULL) && (sample->pending == 0); +} + +static UCS_F_ALWAYS_INLINE ucp_rma_bw_sample_t * +ucp_rma_bw_sample_get(ucp_request_t *req) +{ + return (req->flags & UCP_REQUEST_FLAG_RMA_BW_SAMPLE) ? + req->send.rma.bw_sample : NULL; +} + +static UCS_F_ALWAYS_INLINE void +ucp_rma_bw_sample_try_start(ucp_request_t *req, + ucp_lane_index_t num_lanes) +{ + if (ucs_unlikely(req->send.ep->worker->context->config.ext.rma_bw_measure)) { + ucp_rma_bw_sample_start(req, num_lanes); + } +} + +static UCS_F_ALWAYS_INLINE uct_completion_t * +ucp_rma_bw_frag_start(ucp_request_t *req, ucp_lane_index_t lane_idx, + ucp_rma_bw_frag_t **frag_p) +{ + ucp_rma_bw_sample_t *sample = ucp_rma_bw_sample_get(req); + ucp_rma_bw_frag_t *frag; + + *frag_p = NULL; + if (sample == NULL) { + return &req->send.state.uct_comp; + } + + ucs_assert(lane_idx < sample->num_lanes); + if (sample->num_frags == UCP_RMA_BW_MAX_FRAGS) { + sample->invalid = 1; + return &req->send.state.uct_comp; + } + + frag = &sample->frags[sample->num_frags]; + frag->sample = sample; + frag->lane_idx = lane_idx; + frag->post_time = ucs_get_time(); + frag->comp.func = ucp_rma_bw_frag_complete; + frag->comp.count = 1; + frag->comp.status = UCS_OK; + *frag_p = frag; + return &frag->comp; +} + +static UCS_F_ALWAYS_INLINE void +ucp_rma_bw_frag_posted(ucp_rma_bw_frag_t *frag, size_t bytes, + ucs_status_t status) +{ + ucp_rma_bw_sample_t *sample; + ucp_rma_bw_lane_t *lane; + ucs_time_t now; + + if ((frag == NULL) || UCS_STATUS_IS_ERR(status)) { + return; + } + + sample = frag->sample; + lane = &sample->lanes[frag->lane_idx]; + if (lane->num_frags == 0) { + lane->first_post = frag->post_time; + } + lane->bytes += bytes; + ++lane->num_frags; + if (ucs_likely(status == UCS_INPROGRESS)) { + ++lane->num_async; + ++sample->num_frags; + ++sample->pending; + return; + } + + ucs_assert(status == UCS_OK); + now = ucs_get_time(); + if (now > lane->last_comp) { + lane->last_comp = now; + } +} + +#endif diff --git a/test/gtest/ucp/test_ucp_rma.cc b/test/gtest/ucp/test_ucp_rma.cc index 8faa3e73b46..6d868dbe908 100755 --- a/test/gtest/ucp/test_ucp_rma.cc +++ b/test/gtest/ucp/test_ucp_rma.cc @@ -13,6 +13,7 @@ extern "C" { #include #include #include +#include #include #include } @@ -1697,3 +1698,84 @@ UCS_TEST_SKIP_COND_P(test_ucp_rma_sgl, put_without_proto, } UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_rma_sgl, all, "all") + + +class test_ucp_rma_bw : public test_ucp_rma { +public: + static void get_test_variants(std::vector& variants) { + add_variant_with_value(variants, UCP_FEATURE_RMA, 0, ""); + } + + test_ucp_rma_bw() + { + modify_config("RMA_BW_MEASURE", "y"); + modify_config("ZCOPY_THRESH", "0"); + } +}; + +UCS_TEST_P(test_ucp_rma_bw, put_get_enabled) +{ + constexpr size_t size = 512 * UCS_KBYTE; + mem_buffer sendbuf(size, UCS_MEMORY_TYPE_HOST); + mem_buffer recvbuf(size, UCS_MEMORY_TYPE_HOST); + mapped_buffer rbuf(size, receiver()); + ucs::handle rkey = rbuf.rkey(sender()); + ucp_request_param_t param = {0}; + + mem_buffer::pattern_fill(sendbuf.ptr(), size, ucs::rand()); + ASSERT_UCS_OK(request_wait(ucp_put_nbx( + sender().ep(), sendbuf.ptr(), size, (uint64_t)rbuf.ptr(), rkey, + ¶m))); + flush_ep(sender()); + ASSERT_TRUE(mem_buffer::compare(sendbuf.ptr(), rbuf.ptr(), size, + UCS_MEMORY_TYPE_HOST)); + + ASSERT_UCS_OK(request_wait(ucp_get_nbx( + sender().ep(), recvbuf.ptr(), size, (uint64_t)rbuf.ptr(), rkey, + ¶m))); + ASSERT_TRUE(mem_buffer::compare(sendbuf.ptr(), recvbuf.ptr(), size, + UCS_MEMORY_TYPE_HOST)); + for (unsigned i = 0; i < UCP_RMA_BW_MAX_ACTIVE; ++i) { + const auto &sample = sender().worker()->rma_bw_samples[i]; + EXPECT_EQ(nullptr, sample.req); + EXPECT_EQ(0u, sample.pending); + } +} + +UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_rma_bw, all, "all") + + +TEST(test_ucp_rma_bw_helpers, detach_pending) +{ + ucp_request_t req = {}; + ucp_rma_bw_sample_t sample = {}; + ucp_rma_bw_frag_t *frag; + uct_completion_t *comps[3]; + + req.flags = UCP_REQUEST_FLAG_RMA_BW_SAMPLE; + req.send.rma.bw_sample = &sample; + req.send.state.uct_comp.count = 4; + req.send.state.uct_comp.func = ucp_rma_bw_sample_complete; + sample.req = &req; + sample.num_lanes = 3; + + for (unsigned lane = 0; lane < sample.num_lanes; ++lane) { + comps[lane] = ucp_rma_bw_frag_start(&req, lane, &frag); + ucp_rma_bw_frag_posted(frag, 4096, UCS_INPROGRESS); + } + ucp_proto_request_zcopy_abort(&req, UCS_ERR_CANCELED); + ucp_rma_bw_sample_detach(&req); + + EXPECT_FALSE(req.flags & UCP_REQUEST_FLAG_RMA_BW_SAMPLE); + EXPECT_TRUE(req.send.state.uct_comp.func == + ucp_proto_request_zcopy_completion); + EXPECT_EQ(nullptr, req.send.rma.bw_sample); + EXPECT_EQ(nullptr, sample.req); + EXPECT_EQ(UCS_ERR_CANCELED, req.send.state.uct_comp.status); + EXPECT_FALSE(ucp_rma_bw_sample_is_free(&sample)); + + for (uct_completion_t *comp : comps) { + comp->func(comp); + } + EXPECT_TRUE(ucp_rma_bw_sample_is_free(&sample)); +}