Skip to content

UCP/RMA/BCOPY: Support non-host memory for get/bcopy #11416

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tvegas1 wants to merge 5 commits into
base: master
Choose a base branch
from
Open

UCP/RMA/BCOPY: Support non-host memory for get/bcopy #11416

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
d0f87df
UCP/RMA/BCOPY: Support non-host memory for get/bcopy
tvegas1 May 5, 2026
09a5f15
UCP/RMA/BCOPY: Support non-host memory for get/bcopy
tvegas1 May 5, 2026
f9645a6
UCP/RMA/BCOPY: Support non-host memory for get/bcopy
tvegas1 May 11, 2026
8a7bf58
Merge branch 'master' into get_bcopy_to_non_host
brminich Jun 17, 2026
e1756bc
Merge branch 'master' into get_bcopy_to_non_host
brminich Sep 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 81 additions & 5 deletions src/ucp/rma/get_offload.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,28 @@
#include <ucp/proto/proto_multi.inl>


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,
size_t max_frag)
{
return ucs_div_round_up(req->send.state.dt_iter.length, max_frag);
}

static UCS_F_ALWAYS_INLINE ucs_status_t
Expand All @@ -30,16 +47,35 @@ ucp_proto_get_offload_bcopy_send_func(ucp_request_t *req,
ucp_datatype_iter_t *next_iter,
ucp_lane_index_t *lane_shift)
{
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)) {

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

max_frag is always > 0 (asserted in ucp_proto_multi_init_priv()), pls drop this check and the one in the progress function (:119); this one also adds a branch on the send path.

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 / 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, max_frag,
req->send.state.dt_iter.length);

unpack_ctx = req->send.buffer;

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why keep the unpack ctx array in req->send.buffer? it is the user send buffer field and get_am sets it for get_reply processing - pls use a dedicated field or document the aliasing.

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);
Expand All @@ -49,6 +85,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_proto_request_restart_or_complete(req, req->send.state.uct_comp.status,
Expand All @@ -60,6 +98,9 @@ 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 max_frag;
size_t num_frags;
ucs_status_t status;

if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) {
Expand All @@ -71,6 +112,25 @@ 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);

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);

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: num_frags is always >= 1 (zero-length get returns early in ucp_get_nbx()), and if it were 0 the frag_idx < num_frags assert in the send function would fail anyway - pls use num_frags directly instead of ucs_max(num_frags, 1).

req->send.buffer = ucs_malloc(ctx_count *

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this adds malloc/free per GET request plus O(length/max_frag) memory (~200KB for a 100MB get on an 8KB frag lane), and get/bcopy is selected for small messages too, where docs/OptimizationStyle.md requires no malloc/free on the data path. can we keep the host path allocation-free (pass dest with the plain memcpy unpack cb) and use a pooled ctx only when the local memory is not CPU-accessible?

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;
Expand Down Expand Up @@ -102,7 +162,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 |
Expand Down Expand Up @@ -134,8 +194,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) {

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can reset/abort run while fragments posted earlier are still in flight (pending replay from ucp_wireup_replay_pending_requests(), or a send error on a later fragment)? in that case UCT is left with a dangling unpack arg - IMO the array should be released only in the completion callback. (also applies to the free at :201)

ucs_free(req->send.buffer);
req->send.buffer = NULL;
}

ucp_proto_get_offload_reset(req);
/* GET is idempotent, so a restarted request replays from the beginning
* instead of resuming from an offset which counts posted rather than
Expand All @@ -152,7 +228,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
};

Expand Down
22 changes: 20 additions & 2 deletions test/gtest/ucp/test_ucp_rma.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -295,11 +295,23 @@ 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<send_func_t>(&test_ucp_rma::get_b));
test_mem_types(static_cast<send_func_t>(&test_ucp_rma::get_b), 1);
}

UCS_TEST_P(test_ucp_rma, get_nonblocking) {
test_mem_types(static_cast<send_func_t>(&test_ucp_rma::get_nbi));
test_mem_types(static_cast<send_func_t>(&test_ucp_rma::get_nbi), 1);
}

UCS_TEST_P(test_ucp_rma, get_bcopy_forced_success, "PROTOS=get/bcopy,reconfig")
{
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<send_func_t>(&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,
Expand Down Expand Up @@ -656,6 +668,12 @@ UCS_TEST_P(test_ucp_proto_emulation_enable, get_zcopy_forced_success,
test_forced_message_sizes(static_cast<send_func_t>(&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<send_func_t>(&test_ucp_rma::get_b));
}

UCS_TEST_P(test_ucp_proto_emulation_enable, put_zcopy_forced_success,
"PROTOS=put/offload/*,reconfig")
{
Expand Down
Loading