Skip to content

UCP/RMA: Sample multi-rail zcopy throughput #12002

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

Merged
evgeny-leksikov merged 8 commits into from
Oct 2, 2026
Merged

UCP/RMA: Sample multi-rail zcopy throughput #12002

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
136cb21
UCP/RMA: Sample multi-rail zcopy throughput
gleon99 Sep 30, 2026
94249a3
UCP/RMA: Make throughput sampling restart-safe and allocation-free
gleon99 Sep 30, 2026
1fee062
TEST/UCP: Remove some tests
gleon99 Sep 30, 2026
bcf1ac7
UCP/RMA: Simplify
gleon99 Sep 30, 2026
bd52808
UCP/RMA: CR fix
gleon99 Sep 30, 2026
7a755c6
UCP/RMA: CR fix
gleon99 Oct 1, 2026
39221b5
UCP/RMA: Support multi-lane sample
gleon99 Oct 1, 2026
92f46e6
UCP/RMA: Address measured sampling review findings
gleon99 Oct 1, 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
6 changes: 6 additions & 0 deletions src/ucp/core/ucp_context.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 two-rail PUT/GET zcopy operations. "

@evgeny-leksikov evgeny-leksikov Sep 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

maybe "multi-lane"?

@gleon99 gleon99 Oct 1, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed

@evgeny-leksikov evgeny-leksikov Oct 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

could you please point me where? I would suggest to unify everywhere if we are limited by 2 rails at this phase

@gleon99 gleon99 Oct 1, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed now

"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,
Expand Down
2 changes: 2 additions & 0 deletions src/ucp/core/ucp_context.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 */
Expand Down
15 changes: 10 additions & 5 deletions src/ucp/core/ucp_request.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment thread
gleon99 marked this conversation as resolved.
Resolved
};


Expand Down Expand Up @@ -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 */
void *bw_sample;
};
} rma;

struct {
Expand Down
4 changes: 4 additions & 0 deletions src/ucp/core/ucp_worker.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -314,6 +314,10 @@ 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 */
ucs_time_t rma_bw_next_sample; /* Next sample admission */
unsigned rma_bw_active; /* Active RMA samples */
uint64_t rma_bw_valid; /* Valid RMA samples */
Comment thread
gleon99 marked this conversation as resolved.
Resolved
uint64_t rma_bw_rejected; /* Rejected RMA samples */
ucp_tl_bitmap_t atomic_tls; /* Which resources can be used for atomics */

int inprogress;
Expand Down
57 changes: 47 additions & 10 deletions src/ucp/rma/get_offload.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#endif

#include "rma.inl"
#include "rma_bw.inl"

#include <ucp/core/ucp_request.inl>
#include <ucp/dt/datatype_iter.inl>
Expand Down Expand Up @@ -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,
Comment thread
gleon99 marked this conversation as resolved.
Resolved
lpriv->super.rkey_index);
size_t offset = req->send.state.dt_iter.offset;
const ucp_proto_multi_priv_t *mpriv;
uct_completion_t *comp = &req->send.state.uct_comp;
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),
Expand All @@ -176,22 +179,56 @@ 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);
}
Comment on lines +182 to +191

@evgeny-leksikov evgeny-leksikov Sep 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

why do we need 2 timestamps? should not doorbell of zcopy be constant?
I think we need it only if status == INPROGRESS
if status is err, then sampling probably does not make sense

@gleon99 gleon99 Oct 1, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

The return-time clock now runs only for sync UCS_OK. An async frag still needs the pre-post time and later callback time. Err contributes no bytes.


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);
const ucp_proto_multi_priv_t *mpriv = req->send.proto_config->priv;

if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) {
ucp_rma_bw_sample_start(req, mpriv);
}

/* 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,
(req->flags & UCP_REQUEST_FLAG_RMA_BW_SAMPLE) ?
ucp_proto_get_offload_zcopy_sampled_send_func :
ucp_proto_get_offload_zcopy_send_func,
ucp_request_invoke_uct_completion_success,
ucp_proto_request_zcopy_completion);
(req->flags & UCP_REQUEST_FLAG_RMA_BW_SAMPLE) ?
ucp_rma_bw_sample_complete : ucp_proto_request_zcopy_completion);
}

static void
Expand Down Expand Up @@ -258,6 +295,6 @@ ucp_proto_t ucp_get_offload_zcopy_proto = {
.probe = ucp_proto_get_offload_zcopy_probe,
.query = ucp_proto_multi_query,
.progress = {ucp_proto_get_offload_zcopy_progress},
.abort = ucp_proto_request_zcopy_abort,
.abort = ucp_rma_bw_abort,
.reset = ucp_proto_get_offload_zcopy_reset
};
52 changes: 43 additions & 9 deletions src/ucp/rma/put_offload.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

#include "rma.h"
#include "rma.inl"
#include "rma_bw.inl"

#include <ucp/core/ucp_request.inl>
#include <ucp/dt/datatype_iter.inl>
Expand Down Expand Up @@ -245,26 +246,34 @@ 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)
Comment thread
gleon99 marked this conversation as resolved.
Resolved
{
ucp_ep_h ep = req->send.ep;
uct_ep_h uct_ep = ucp_ep_get_lane(ep, lpriv->super.lane);
uint64_t address = req->send.rma.remote_addr +
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;

ucp_datatype_iter_next_iov(&req->send.state.dt_iter,
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);
Expand All @@ -273,18 +282,43 @@ 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);
const ucp_proto_multi_priv_t *mpriv = req->send.proto_config->priv;

if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) {
ucp_rma_bw_sample_start(req, mpriv);
}

/* 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,
(req->flags & UCP_REQUEST_FLAG_RMA_BW_SAMPLE) ?
Comment thread
gleon99 marked this conversation as resolved.
Resolved
ucp_proto_put_offload_zcopy_sampled_send_func :
ucp_proto_put_offload_zcopy_send_func,
ucp_request_invoke_uct_completion_success,
ucp_proto_request_zcopy_completion);
(req->flags & UCP_REQUEST_FLAG_RMA_BW_SAMPLE) ?
ucp_rma_bw_sample_complete : ucp_proto_request_zcopy_completion);
}

static void
Expand Down Expand Up @@ -354,7 +388,7 @@ ucp_proto_t ucp_put_offload_zcopy_proto = {
.probe = ucp_proto_put_offload_zcopy_probe,
.query = ucp_proto_multi_query,
.progress = {ucp_proto_put_offload_zcopy_progress},
.abort = ucp_proto_request_zcopy_abort,
.abort = ucp_rma_bw_abort,
.reset = ucp_proto_offload_zcopy_reset
};

Expand Down
Loading
Loading