Skip to content

UCP/RMA: Add AM-based software emulation protocol for SGL put #11943

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
michal-shalev wants to merge 8 commits into
base: master
Choose a base branch
from
Open

UCP/RMA: Add AM-based software emulation protocol for SGL put #11943

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
23de252
UCP/RMA: Add AM-based software emulation protocol for SGL put
michal-shalev Sep 13, 2026
a6b9b00
GTEST/UCP: Add coverage for the SGL put emulation protocol
michal-shalev Sep 13, 2026
f2dfc07
UCP/RMA: Use the shared SGL rkey config for the AM put header memory …
michal-shalev Sep 13, 2026
9d76430
GTEST/UCP: Run the SGL emulation tests without a zcopy lane
michal-shalev Sep 13, 2026
1dad743
UCP/RMA: Take the remote memory type from the request in the AM put h…
michal-shalev Sep 13, 2026
ddf20d2
GTEST/UCP: Clarify the in-place completion check in the SGL put test
michal-shalev Sep 13, 2026
fa7d955
UCP/RMA: Pass the SGL element length as the pack total length
michal-shalev Sep 13, 2026
fdd38cb
GTEST/UCP: Deny immediate completion in the SGL put callback tests
michal-shalev Sep 13, 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
1 change: 1 addition & 0 deletions src/ucp/proto/proto.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
_macro(ucp_get_offload_zcopy_proto) \
_macro(ucp_get_rndv_proto) \
_macro(ucp_put_am_bcopy_proto) \
_macro(ucp_put_sgl_am_bcopy_proto) \
_macro(ucp_put_offload_short_proto) \
_macro(ucp_put_offload_bcopy_proto) \
_macro(ucp_put_offload_zcopy_proto) \
Expand Down
114 changes: 106 additions & 8 deletions src/ucp/rma/put_am.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,16 +17,33 @@
#include <ucp/proto/proto_multi.inl>


typedef struct {
ucp_request_t *req;
const void *buffer;
uint64_t remote_addr;
size_t length;
size_t elem_length;
} ucp_proto_put_sgl_am_bcopy_pack_ctx_t;


static UCS_F_ALWAYS_INLINE void
Comment thread
michal-shalev marked this conversation as resolved.
ucp_proto_put_am_bcopy_pack_hdr(ucp_put_hdr_t *puth, ucp_request_t *req,
uint64_t remote_addr)
{
puth->address = remote_addr;
puth->ep_id = ucp_send_request_get_ep_remote_id(req);
puth->mem_type = req->send.rma.rkey->mem_type;
}

static size_t ucp_proto_put_am_bcopy_pack(void *dest, void *arg)
{
ucp_proto_multi_pack_ctx_t *pack_ctx = arg;
ucp_request_t *req = pack_ctx->req;
ucp_put_hdr_t *puth = dest;

puth->address = req->send.rma.remote_addr +
req->send.state.dt_iter.offset;
puth->ep_id = ucp_send_request_get_ep_remote_id(req);
puth->mem_type = req->send.rma.rkey->mem_type;
ucp_proto_put_am_bcopy_pack_hdr(puth, req,
req->send.rma.remote_addr +
req->send.state.dt_iter.offset);

return sizeof(*puth) + ucp_proto_multi_data_pack(pack_ctx, puth + 1);
}
Expand All @@ -48,7 +65,64 @@ ucp_proto_put_am_bcopy_send_func(ucp_request_t *req,
ucp_proto_put_am_bcopy_pack, &pack_ctx, NULL);
}

static ucs_status_t ucp_proto_put_am_bcopy_progress(uct_pending_req_t *self)
static size_t ucp_proto_put_sgl_am_bcopy_pack(void *dest, void *arg)
{
ucp_proto_put_sgl_am_bcopy_pack_ctx_t *pack_ctx = arg;
ucp_request_t *req = pack_ctx->req;
ucp_datatype_iter_t *dt_iter = &req->send.state.dt_iter;
ucp_put_hdr_t *puth = dest;

ucp_proto_put_am_bcopy_pack_hdr(puth, req, pack_ctx->remote_addr);

ucp_dt_contig_pack(req->send.ep->worker, puth + 1, pack_ctx->buffer,
pack_ctx->length,
(ucs_memory_type_t)dt_iter->mem_info.type,
pack_ctx->elem_length);

return sizeof(*puth) + pack_ctx->length;
}

static UCS_F_ALWAYS_INLINE ucs_status_t
ucp_proto_put_sgl_am_bcopy_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_datatype_iter_t *dt_iter = &req->send.state.dt_iter;
void *buffer = NULL;
size_t length = 0;
uint64_t remote_addr = 0;
size_t elem_index = 0;
ucp_proto_put_sgl_am_bcopy_pack_ctx_t pack_ctx;
size_t max_payload;

ucs_assertv(lpriv->max_frag > sizeof(ucp_put_hdr_t), "max_frag=%zu",
lpriv->max_frag);
max_payload = lpriv->max_frag - sizeof(ucp_put_hdr_t);

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 13, 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 we add a short comment on why ucp_proto_multi_max_payload() is not used here? for SGL dt_iter->length is the element count, so the weight scaling inside the helper would cap the payload by the number of elements — without a note someone is likely to "simplify" this back to the helper.


if (ucp_datatype_iter_next_sgl_frags(dt_iter,
req->send.rma.sgl.remote_addrs, 1,
max_payload, next_iter, &buffer,
&length, &remote_addr,
&elem_index) == 0) {
return UCS_OK;
}

pack_ctx.req = req;
pack_ctx.buffer = buffer;
pack_ctx.remote_addr = remote_addr;
pack_ctx.length = length;
pack_ctx.elem_length = dt_iter->type.sgl.lengths[elem_index];

return ucp_rma_sw_do_am_bcopy(req, UCP_AM_ID_PUT, lpriv->super.lane,
ucp_proto_put_sgl_am_bcopy_pack, &pack_ctx,
NULL);
}

static UCS_F_ALWAYS_INLINE ucs_status_t
ucp_proto_put_am_bcopy_common_progress(uct_pending_req_t *self,
ucp_proto_send_multi_cb_t send_func,
unsigned dt_mask)
{
ucp_request_t *req = ucs_container_of(self, ucp_request_t,
send.uct);
Expand All @@ -74,10 +148,22 @@ static ucs_status_t ucp_proto_put_am_bcopy_progress(uct_pending_req_t *self)
req->flags |= UCP_REQUEST_FLAG_PROTO_INITIALIZED;
}

return ucp_proto_multi_progress(req, mpriv,
ucp_proto_put_am_bcopy_send_func,
return ucp_proto_multi_progress(req, mpriv, send_func,
ucp_proto_request_bcopy_complete_success,
UCP_DT_MASK_CONTIG_IOV);
dt_mask);
}

static ucs_status_t ucp_proto_put_am_bcopy_progress(uct_pending_req_t *self)
{
return ucp_proto_put_am_bcopy_common_progress(
self, ucp_proto_put_am_bcopy_send_func, UCP_DT_MASK_CONTIG_IOV);
}

static ucs_status_t ucp_proto_put_sgl_am_bcopy_progress(uct_pending_req_t *self)
{
return ucp_proto_put_am_bcopy_common_progress(
self, ucp_proto_put_sgl_am_bcopy_send_func,
UCS_BIT(UCP_DATATYPE_SGL));
}

static void
Expand Down Expand Up @@ -137,3 +223,15 @@ ucp_proto_t ucp_put_am_bcopy_proto = {
.abort = ucp_proto_request_bcopy_abort,
.reset = ucp_proto_request_bcopy_reset
};

ucp_proto_t ucp_put_sgl_am_bcopy_proto = {
.name = "put/sgl/am/bcopy",
.desc = "sgl " UCP_PROTO_RMA_EMULATION_DESC,
.flags = 0,
.dt_mask = UCS_BIT(UCP_DATATYPE_SGL),
.probe = ucp_proto_put_am_bcopy_probe,
Comment thread
michal-shalev marked this conversation as resolved.
.query = ucp_proto_multi_query,
.progress = {ucp_proto_put_sgl_am_bcopy_progress},
.abort = ucp_proto_request_bcopy_abort,
.reset = ucp_proto_request_bcopy_reset
};
2 changes: 1 addition & 1 deletion src/ucp/rma/put_offload.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -624,7 +624,7 @@ ucp_proto_put_sgl_offload_sw_progress(uct_pending_req_t *self)

ucp_proto_t ucp_put_sgl_offload_sw_proto = {
.name = "put/sgl/offload_sw",
.desc = "sgl " UCP_PROTO_RMA_EMULATION_DESC,
.desc = "sgl per-element zcopy",

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 13, 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: use the existing macro, like the other protos in this file:

Suggested change
.desc = "sgl per-element zcopy",
.desc = "sgl per-element " UCP_PROTO_ZCOPY_DESC,

.flags = 0,
.dt_mask = UCS_BIT(UCP_DATATYPE_SGL),
.probe = ucp_proto_put_sgl_offload_sw_probe,
Expand Down
83 changes: 63 additions & 20 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 @@ -1001,6 +1001,8 @@ UCS_TEST_P(test_ucp_ep_based_fence, test_ep_based_fence_before_atomic) {

UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_ep_based_fence, all, "all")

#define UCP_SGL_EMULATION_PROTOS "PROTOS=put/sgl/am/*,reconfig"

class test_ucp_rma_sgl : public test_ucp_rma {
public:
static void get_base_variants(std::vector<ucp_test_variant>& variants) {
Expand All @@ -1014,11 +1016,6 @@ class test_ucp_rma_sgl : public test_ucp_rma {
}

virtual void init() override {
/* FIXME: sporadic failure on CUDA memory type. re-enable once fixed */
if (mem_type() == UCS_MEMORY_TYPE_CUDA) {
UCS_TEST_SKIP_R("sporadic failure on CUDA memory type");
}

modify_config("MAX_RMA_RAILS", "2");
Comment thread
michal-shalev marked this conversation as resolved.
Comment thread
michal-shalev marked this conversation as resolved.
Comment thread
michal-shalev marked this conversation as resolved.
test_ucp_rma::init();
}
Expand All @@ -1033,6 +1030,8 @@ class test_ucp_rma_sgl : public test_ucp_rma {
SGL_OP_GET
};

bool m_require_zcopy_lane = true;

struct sgl_ctx {
std::vector<mapped_buffer> src;
std::vector<mapped_buffer> dst;
Expand Down Expand Up @@ -1194,7 +1193,7 @@ class test_ucp_rma_sgl : public test_ucp_rma {

uint64_t zcopy_cap = (op == SGL_OP_PUT) ? UCT_IFACE_FLAG_PUT_ZCOPY :
UCT_IFACE_FLAG_GET_ZCOPY;
if (!sender().has_lane_with_caps(zcopy_cap)) {
if (m_require_zcopy_lane && !sender().has_lane_with_caps(zcopy_cap)) {
UCS_TEST_SKIP_R("zcopy is not supported");
}

Expand Down Expand Up @@ -1222,7 +1221,8 @@ class test_ucp_rma_sgl : public test_ucp_rma {

if (use_callback) {
param.op_attr_mask |= UCP_OP_ATTR_FIELD_CALLBACK |
UCP_OP_ATTR_FIELD_USER_DATA;
UCP_OP_ATTR_FIELD_USER_DATA |
UCP_OP_ATTR_FLAG_NO_IMM_CMPL;
param.cb.send = [](void *request, ucs_status_t status,
void *user_data) {
cb_state *s = static_cast<cb_state*>(user_data);
Expand All @@ -1242,8 +1242,6 @@ class test_ucp_rma_sgl : public test_ucp_rma {
return;
}

ASSERT_TRUE(UCS_PTR_IS_PTR(sptr));

auto verify_sgl_buffers = [&]() {
ucs_memory_type_t mtype = mem_type();
for (size_t i = 0; i < num; i++) {
Expand All @@ -1260,24 +1258,41 @@ class test_ucp_rma_sgl : public test_ucp_rma {
}
};

if (use_callback) {
while (!cb.completed) {
ucp_worker_progress(sender().worker());
ucp_worker_progress(receiver().worker());
}
EXPECT_UCS_OK(cb.status);
if (!UCS_PTR_IS_PTR(sptr)) {
ASSERT_UCS_OK(UCS_PTR_STATUS(sptr));
} else {
while (!ucp_request_is_completed(sptr)) {
ucp_worker_progress(sender().worker());
ucp_worker_progress(receiver().worker());
if (use_callback) {
while (!cb.completed) {
ucp_worker_progress(sender().worker());
ucp_worker_progress(receiver().worker());
}
EXPECT_UCS_OK(cb.status);
} else {
while (!ucp_request_is_completed(sptr)) {
ucp_worker_progress(sender().worker());
ucp_worker_progress(receiver().worker());
}
}

ucp_request_release(sptr);
}

ucp_request_release(sptr);
flush_ep(sender());
verify_sgl_buffers();
}

static bool offload_proto_selected(ucs_status_ptr_t sptr) {
if (!UCS_PTR_IS_PTR(sptr)) {
/* Only the emulation protocol can complete the put in-place, by
copying the data to a bounce buffer */
return false;
}

const ucp_request_t *req = (const ucp_request_t*)sptr - 1;
return strstr(req->send.proto_config->proto->name,
"put/sgl/offload") != nullptr;
}

void test_put_sgl(const std::vector<size_t> &elem_sizes,
bool use_memhs = true, bool use_callback = false,
bool set_remote_count = true,
Expand Down Expand Up @@ -1415,6 +1430,29 @@ UCS_TEST_P(test_ucp_rma_sgl, put_no_remote_count) {
test_put_sgl(4, 2 * UCS_KBYTE, true, false, false);
Comment thread
michal-shalev marked this conversation as resolved.
}

UCS_TEST_P(test_ucp_rma_sgl, put_emulation, UCP_SGL_EMULATION_PROTOS) {
Comment thread
michal-shalev marked this conversation as resolved.
Comment thread
michal-shalev marked this conversation as resolved.
m_require_zcopy_lane = false;
test_put_sgl({64, 256, UCS_KBYTE, 4 * UCS_KBYTE, 512});
}

UCS_TEST_P(test_ucp_rma_sgl, put_emulation_with_callback,
UCP_SGL_EMULATION_PROTOS) {
m_require_zcopy_lane = false;
test_put_sgl(10, UCS_KBYTE, true, true);
}

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 13, 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: with the emulation proto the put normally completes in place, so the callback is never invoked and this test adds nothing over put_emulation. can we set UCP_OP_ATTR_FLAG_NO_IMM_CMPL here to actually cover the callback path?

With the emulation protocol every fragment is sent synchronously by uct_ep_am_bcopy(), so ucp_request_send() drains the whole SGL in one call and ucp_proto_request_send_op_common() returns via ucp_request_imm_cmpl_param() — a status, not a request. In that case UCX does not invoke the send callback, and the new !UCS_PTR_IS_PTR(sptr) branch in test_sgl() skips the cb.completed/cb.status checks entirely. For 10×1KB this is the normal outcome.


UCS_TEST_P(test_ucp_rma_sgl, put_emulation_no_memhs,
UCP_SGL_EMULATION_PROTOS) {
m_require_zcopy_lane = false;
test_put_sgl(4, 2 * UCS_KBYTE, false);
}

UCS_TEST_SKIP_COND_P(test_ucp_rma_sgl, put_emulation_fragmented,
RUNNING_ON_VALGRIND, UCP_SGL_EMULATION_PROTOS) {
m_require_zcopy_lane = false;
test_put_sgl(4, 256 * UCS_KBYTE);
}

UCS_TEST_P(test_ucp_rma_sgl, put_split_between_lanes) {
static constexpr size_t NUM_ELEMS = 16;

Expand All @@ -1432,7 +1470,11 @@ UCS_TEST_P(test_ucp_rma_sgl, put_split_between_lanes) {
ucs_status_ptr_t sptr = sgl_op_nbx(SGL_OP_PUT, &local, NUM_ELEMS,
UCP_REMOTE_ADDR_INVALID,
UCP_RKEY_INVALID, &param);
ASSERT_TRUE(UCS_PTR_IS_PTR(sptr));
if (!offload_proto_selected(sptr)) {
ASSERT_UCS_OK(request_wait(sptr));
flush_ep(sender());
UCS_TEST_SKIP_R("SGL offload protocol was not selected");
}

/* All the elements fit into a single post, so at least one outstanding post
per lane of the selected protocol means they were split between them */
Expand Down Expand Up @@ -1698,3 +1740,4 @@ UCS_TEST_SKIP_COND_P(test_ucp_rma_sgl, put_without_proto,
}

UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_rma_sgl, all, "all")
UCP_INSTANTIATE_TEST_CASE_TLS_GPU_AWARE(test_ucp_rma_sgl, tcp, "tcp")