From d0f87df22d5e5f6898f4b0230b5eddf6646ee467 Mon Sep 17 00:00:00 2001 From: Thomas Vegas Date: Tue, 5 May 2026 14:14:21 +0300 Subject: [PATCH 1/3] UCP/RMA/BCOPY: Support non-host memory for get/bcopy --- src/ucp/rma/get_offload.c | 75 +++++++++++++++++++++++++++++++--- test/gtest/ucp/test_ucp_rma.cc | 26 +++++++++++- 2 files changed, 94 insertions(+), 7 deletions(-) diff --git a/src/ucp/rma/get_offload.c b/src/ucp/rma/get_offload.c index 7bec24804cf..71ea68ecb57 100644 --- a/src/ucp/rma/get_offload.c +++ b/src/ucp/rma/get_offload.c @@ -17,11 +17,28 @@ #include +typedef struct { + ucp_request_t *req; + void *dest; +} ucp_proto_get_bcopy_unpack_ctx_t; + + static void ucp_proto_get_offload_bcopy_unpack(void *arg, const void *data, size_t length) { - void *dest = arg; - ucs_memcpy_relaxed(dest, data, length, UCS_ARCH_MEMCPY_NT_SOURCE, length); + ucp_proto_get_bcopy_unpack_ctx_t *unpack_ctx = arg; + ucp_request_t *req = unpack_ctx->req; + + ucp_dt_contig_unpack(req->send.ep->worker, unpack_ctx->dest, data, length, + (ucs_memory_type_t)req->send.state.dt_iter.mem_info.type, + req->send.state.dt_iter.length); +} + +static UCS_F_ALWAYS_INLINE size_t +ucp_proto_get_offload_bcopy_num_frags(const ucp_request_t *req, + const ucp_proto_multi_lane_priv_t *lpriv) +{ + return ucs_div_round_up(req->send.state.dt_iter.length, lpriv->max_frag); } static UCS_F_ALWAYS_INLINE ucs_status_t @@ -30,16 +47,30 @@ ucp_proto_get_offload_bcopy_send_func(ucp_request_t *req, ucp_datatype_iter_t *next_iter, ucp_lane_index_t *lane_shift) { + ucp_proto_get_bcopy_unpack_ctx_t *unpack_ctx; uct_rkey_t tl_rkey = ucp_rkey_get_tl_rkey(req->send.rma.rkey, lpriv->super.rkey_index); size_t max_length, length; + size_t frag_idx; void *dest; max_length = ucp_proto_multi_max_payload(req, lpriv, 0); length = ucp_datatype_iter_next_ptr(&req->send.state.dt_iter, max_length, next_iter, &dest); + frag_idx = req->send.state.dt_iter.offset / lpriv->max_frag; + ucs_assertv(frag_idx < ucp_proto_get_offload_bcopy_num_frags(req, lpriv), + "frag_idx=%zu offset=%zu max_frag=%zu length=%zu", frag_idx, + req->send.state.dt_iter.offset, lpriv->max_frag, + req->send.state.dt_iter.length); + + unpack_ctx = req->send.buffer; + unpack_ctx[frag_idx].req = req; + unpack_ctx[frag_idx].dest = dest; + return uct_ep_get_bcopy(ucp_ep_get_lane(req->send.ep, lpriv->super.lane), - ucp_proto_get_offload_bcopy_unpack, dest, length, + ucp_proto_get_offload_bcopy_unpack, + &unpack_ctx[frag_idx], + length, req->send.rma.remote_addr + req->send.state.dt_iter.offset, tl_rkey, &req->send.state.uct_comp); @@ -49,6 +80,8 @@ static void ucp_proto_get_offload_bcopy_completion(uct_completion_t *self) { ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.state.uct_comp); + ucs_free(req->send.buffer); + req->send.buffer = NULL; ucp_datatype_iter_cleanup(&req->send.state.dt_iter, 0, UCS_BIT(UCP_DATATYPE_CONTIG)); ucp_request_complete_send(req, req->send.state.uct_comp.status); @@ -59,6 +92,8 @@ static ucs_status_t ucp_proto_get_offload_bcopy_progress(uct_pending_req_t *self ucp_request_t *req = ucs_container_of(self, ucp_request_t, send.uct); const ucp_proto_multi_priv_t *mpriv = req->send.proto_config->priv; + size_t ctx_count; + size_t num_frags; ucs_status_t status; if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) { @@ -70,6 +105,20 @@ static ucs_status_t ucp_proto_get_offload_bcopy_progress(uct_pending_req_t *self return UCS_OK; } + UCS_STATIC_ASSERT(UCP_PROTO_RMA_MAX_BCOPY_LANES == 1); + ucs_assertv(mpriv->num_lanes == 1, "num_lanes=%u", mpriv->num_lanes); + + num_frags = ucp_proto_get_offload_bcopy_num_frags(req, + &mpriv->lanes[0]); + ctx_count = ucs_max(num_frags, (size_t)1); + req->send.buffer = ucs_malloc(ctx_count * + sizeof(ucp_proto_get_bcopy_unpack_ctx_t), + "get_bcopy_unpack_ctx"); + if (req->send.buffer == NULL) { + ucp_proto_request_abort(req, UCS_ERR_NO_MEMORY); + return UCS_OK; + } + ucp_proto_completion_init(&req->send.state.uct_comp, ucp_proto_get_offload_bcopy_completion); req->flags |= UCP_REQUEST_FLAG_PROTO_INITIALIZED; @@ -101,7 +150,7 @@ ucp_proto_get_offload_bcopy_probe(const ucp_proto_init_params_t *init_params) .super.max_iov_offs = UCP_PROTO_COMMON_OFFSET_INVALID, .super.hdr_size = 0, .super.send_op = UCT_EP_OP_GET_BCOPY, - .super.memtype_op = UCT_EP_OP_LAST, + .super.memtype_op = UCT_EP_OP_PUT_SHORT, .super.flags = UCP_PROTO_COMMON_INIT_FLAG_RECV_ZCOPY | UCP_PROTO_COMMON_INIT_FLAG_REMOTE_ACCESS | UCP_PROTO_COMMON_INIT_FLAG_RESPONSE | @@ -133,8 +182,24 @@ static void ucp_proto_get_offload_reset(ucp_request_t *req) req->send.state.completed_size = 0; } +static void ucp_proto_get_offload_bcopy_abort(ucp_request_t *req, + ucs_status_t status) +{ + if (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED) { + ucs_free(req->send.buffer); + req->send.buffer = NULL; + } + + ucp_proto_request_bcopy_abort(req, status); +} + static ucs_status_t ucp_proto_get_offload_bcopy_reset(ucp_request_t *req) { + if (req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED) { + ucs_free(req->send.buffer); + req->send.buffer = NULL; + } + ucp_proto_get_offload_reset(req); return ucp_proto_request_bcopy_reset(req); } @@ -146,7 +211,7 @@ ucp_proto_t ucp_get_offload_bcopy_proto = { .probe = ucp_proto_get_offload_bcopy_probe, .query = ucp_proto_multi_query, .progress = {ucp_proto_get_offload_bcopy_progress}, - .abort = ucp_proto_request_bcopy_abort, + .abort = ucp_proto_get_offload_bcopy_abort, .reset = ucp_proto_get_offload_bcopy_reset }; diff --git a/test/gtest/ucp/test_ucp_rma.cc b/test/gtest/ucp/test_ucp_rma.cc index cec51afa787..2425e607f7c 100755 --- a/test/gtest/ucp/test_ucp_rma.cc +++ b/test/gtest/ucp/test_ucp_rma.cc @@ -292,11 +292,27 @@ UCS_TEST_SKIP_COND_P(test_ucp_rma, put_nonblocking_iov_zcopy, UCS_TEST_P(test_ucp_rma, get_blocking) { - test_mem_types(static_cast(&test_ucp_rma::get_b)); + test_mem_types(static_cast(&test_ucp_rma::get_b), 1); } UCS_TEST_P(test_ucp_rma, get_nonblocking) { - test_mem_types(static_cast(&test_ucp_rma::get_nbi)); + test_mem_types(static_cast(&test_ucp_rma::get_nbi), 1); +} + +UCS_TEST_P(test_ucp_rma, get_bcopy_forced_success, + "PROTOS=get/bcopy,reconfig") +{ + if (!mem_buffer::is_mem_type_supported(UCS_MEMORY_TYPE_CUDA)) { + UCS_TEST_SKIP_R("CUDA memory is not supported"); + } + + if (!check_reg_mem_types(sender(), UCS_MEMORY_TYPE_CUDA)) { + UCS_TEST_SKIP_R("CUDA memory registration is not supported"); + } + + test_message_sizes(static_cast(&test_ucp_rma::get_b), 1, + 4 * UCS_MBYTE, + UCS_MEMORY_TYPE_CUDA, UCS_MEMORY_TYPE_HOST, 0); } UCS_TEST_SKIP_COND_P(test_ucp_rma, get_nonblocking_iov_zcopy, @@ -444,6 +460,12 @@ UCS_TEST_P(test_ucp_proto_emulation_enable, get_zcopy_forced_success, test_forced_message_sizes(static_cast(&test_ucp_rma::get_b)); } +UCS_TEST_P(test_ucp_proto_emulation_enable, get_bcopy_forced_success, + "PROTOS=get/bcopy,reconfig") +{ + test_forced_message_sizes(static_cast(&test_ucp_rma::get_b)); +} + UCS_TEST_P(test_ucp_proto_emulation_enable, put_zcopy_forced_success, "PROTOS=put/offload/*,reconfig") { From 09a5f1534ea3d84c29610749c9fad43f5db8fb7f Mon Sep 17 00:00:00 2001 From: Thomas Vegas Date: Tue, 5 May 2026 17:20:16 +0300 Subject: [PATCH 2/3] UCP/RMA/BCOPY: Support non-host memory for get/bcopy --- src/ucp/rma/get_offload.c | 27 +++++++++++++++++++-------- 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/src/ucp/rma/get_offload.c b/src/ucp/rma/get_offload.c index 71ea68ecb57..fa4f8d855fb 100644 --- a/src/ucp/rma/get_offload.c +++ b/src/ucp/rma/get_offload.c @@ -36,9 +36,9 @@ static void ucp_proto_get_offload_bcopy_unpack(void *arg, const void *data, static UCS_F_ALWAYS_INLINE size_t ucp_proto_get_offload_bcopy_num_frags(const ucp_request_t *req, - const ucp_proto_multi_lane_priv_t *lpriv) + size_t max_frag) { - return ucs_div_round_up(req->send.state.dt_iter.length, lpriv->max_frag); + return ucs_div_round_up(req->send.state.dt_iter.length, max_frag); } static UCS_F_ALWAYS_INLINE ucs_status_t @@ -47,20 +47,25 @@ ucp_proto_get_offload_bcopy_send_func(ucp_request_t *req, ucp_datatype_iter_t *next_iter, ucp_lane_index_t *lane_shift) { - ucp_proto_get_bcopy_unpack_ctx_t *unpack_ctx; + size_t max_frag = lpriv->max_frag; uct_rkey_t tl_rkey = ucp_rkey_get_tl_rkey(req->send.rma.rkey, lpriv->super.rkey_index); + ucp_proto_get_bcopy_unpack_ctx_t *unpack_ctx; size_t max_length, length; size_t frag_idx; void *dest; + if (ucs_unlikely(max_frag == 0)) { + return UCS_ERR_INVALID_PARAM; + } + max_length = ucp_proto_multi_max_payload(req, lpriv, 0); length = ucp_datatype_iter_next_ptr(&req->send.state.dt_iter, max_length, next_iter, &dest); - frag_idx = req->send.state.dt_iter.offset / lpriv->max_frag; - ucs_assertv(frag_idx < ucp_proto_get_offload_bcopy_num_frags(req, lpriv), + frag_idx = req->send.state.dt_iter.offset / max_frag; + ucs_assertv(frag_idx < ucp_proto_get_offload_bcopy_num_frags(req, max_frag), "frag_idx=%zu offset=%zu max_frag=%zu length=%zu", frag_idx, - req->send.state.dt_iter.offset, lpriv->max_frag, + req->send.state.dt_iter.offset, max_frag, req->send.state.dt_iter.length); unpack_ctx = req->send.buffer; @@ -93,6 +98,7 @@ static ucs_status_t ucp_proto_get_offload_bcopy_progress(uct_pending_req_t *self send.uct); const ucp_proto_multi_priv_t *mpriv = req->send.proto_config->priv; size_t ctx_count; + size_t max_frag; size_t num_frags; ucs_status_t status; @@ -108,8 +114,13 @@ static ucs_status_t ucp_proto_get_offload_bcopy_progress(uct_pending_req_t *self UCS_STATIC_ASSERT(UCP_PROTO_RMA_MAX_BCOPY_LANES == 1); ucs_assertv(mpriv->num_lanes == 1, "num_lanes=%u", mpriv->num_lanes); - num_frags = ucp_proto_get_offload_bcopy_num_frags(req, - &mpriv->lanes[0]); + max_frag = mpriv->lanes[0].max_frag; + if (ucs_unlikely(max_frag == 0)) { + ucp_proto_request_abort(req, UCS_ERR_INVALID_PARAM); + return UCS_OK; + } + + num_frags = ucp_proto_get_offload_bcopy_num_frags(req, max_frag); ctx_count = ucs_max(num_frags, (size_t)1); req->send.buffer = ucs_malloc(ctx_count * sizeof(ucp_proto_get_bcopy_unpack_ctx_t), From f9645a69cfe57222869bc07103c3fe85eed7e6ec Mon Sep 17 00:00:00 2001 From: Thomas Vegas Date: Mon, 11 May 2026 14:07:00 +0300 Subject: [PATCH 3/3] UCP/RMA/BCOPY: Support non-host memory for get/bcopy --- test/gtest/ucp/test_ucp_rma.cc | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/test/gtest/ucp/test_ucp_rma.cc b/test/gtest/ucp/test_ucp_rma.cc index 2425e607f7c..3e7c5633968 100755 --- a/test/gtest/ucp/test_ucp_rma.cc +++ b/test/gtest/ucp/test_ucp_rma.cc @@ -299,15 +299,11 @@ UCS_TEST_P(test_ucp_rma, get_nonblocking) { test_mem_types(static_cast(&test_ucp_rma::get_nbi), 1); } -UCS_TEST_P(test_ucp_rma, get_bcopy_forced_success, - "PROTOS=get/bcopy,reconfig") +UCS_TEST_P(test_ucp_rma, get_bcopy_forced_success, "PROTOS=get/bcopy,reconfig") { - if (!mem_buffer::is_mem_type_supported(UCS_MEMORY_TYPE_CUDA)) { - UCS_TEST_SKIP_R("CUDA memory is not supported"); - } - - if (!check_reg_mem_types(sender(), UCS_MEMORY_TYPE_CUDA)) { - UCS_TEST_SKIP_R("CUDA memory registration is not supported"); + if (!sender().has_lane_with_caps(UCT_IFACE_FLAG_GET_BCOPY) || + !mem_buffer::is_mem_type_supported(UCS_MEMORY_TYPE_CUDA)) { + UCS_TEST_SKIP_R("get_bcopy or CUDA memory not supported"); } test_message_sizes(static_cast(&test_ucp_rma::get_b), 1,