Skip to content

UCP/RMA: Apply ep_based weak fence once per worker fence #12029

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
michael-beebe wants to merge 2 commits into
base: master
Choose a base branch
from
Open

UCP/RMA: Apply ep_based weak fence once per worker fence #12029

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8db1627
UCP/RMA: Apply ep_based weak fence once per worker fence
michael-beebe Oct 5, 2026
ca6fc76
UCP/RMA: Address review comments on ep-based fence
michael-beebe Oct 5, 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 AUTHORS
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@ Manu Shantharam <manu.shantharam@amd.com>
Marek Schimara <Marek.Schimara@bull.net>
Mark Allen <markalle@us.ibm.com>
Matthew Baker <bakermb@ornl.gov>
Michael Beebe <michael.n.beebe@outlook.com>
Michael Braverman <michaelbr@nvidia.com>
Michal Shalev <mshalev@nvidia.com>
Mike Dubman <miked@mellanox.com>
Expand Down
1 change: 1 addition & 0 deletions src/ucp/core/ucp_ep.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,7 @@ static ucp_ep_h ucp_ep_allocate(ucp_worker_h worker, const char *peer_name)
ep->ext->peer_mem = NULL;
ep->ext->unflushed_lanes = 0;
ep->ext->fence_seq = 0;
ep->ext->fenced_lanes = 0;
ep->ext->uct_eps = NULL;
ep->ext->flush_sys_dev_map = 0;

Expand Down
3 changes: 3 additions & 0 deletions src/ucp/core/ucp_ep.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -607,6 +607,9 @@ typedef struct ucp_ep_ext {

ucp_lane_map_t unflushed_lanes; /* Bitmap of lanes which have
unflushed operations */
ucp_lane_map_t fenced_lanes; /* Lanes ordered by a weak
fence since the last
strong fence */
uint64_t fence_seq; /* Sequence number for fence
detection */

Expand Down
1 change: 1 addition & 0 deletions src/ucp/rma/flush.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -905,6 +905,7 @@ ucs_status_t ucp_ep_fence_strong(ucp_ep_h ep)
}

ep->ext->unflushed_lanes = 0;
ep->ext->fenced_lanes = 0;
ep->ext->fence_seq = ep->worker->fence_seq;
return UCS_OK;
}
Expand Down
16 changes: 15 additions & 1 deletion src/ucp/rma/rma.inl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,13 +169,27 @@ ucp_ep_rma_handle_fence(ucp_ep_h ep, ucp_request_t *req,
/* Apply a fence if EP's sequence is behind worker's */
if (ucs_unlikely(req->flags & UCP_REQUEST_FLAG_FENCE_REQUIRED)) {
if (ucs_unlikely(ep->ext->unflushed_lanes == 0)) {
status = UCS_OK;
/* No earlier operation to order */

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 5, 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.

In the unflushed_lanes == 0 branch: fenced_lanes is always a subset of unflushed_lanes, so this store is redundant. maybe replace it with an assert to document the invariant that ucp_ep_fence_strong() and the new branch rely on?

Suggested change
/* No earlier operation to order */
/* No earlier operation to order */
ucs_assert(ep->ext->fenced_lanes == 0);
ep->ext->fence_seq = ep->worker->fence_seq;
status = UCS_OK;

@michael-beebe michael-beebe Oct 5, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Done in ca6fc76. The store is replaced with ucs_assert(ep->ext->fenced_lane == UCP_NULL_LANE).

ep->ext->fence_seq = ep->worker->fence_seq;
ep->ext->fenced_lanes = 0;
status = UCS_OK;
} else if (ucs_likely(
ucs_is_pow2_or_zero(ep->ext->unflushed_lanes | lane_map))) {
status = ucp_ep_fence_weak(ep);
if (status == UCS_OK) {

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 5, 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.

In the weak fence branch: why set fence_seq to the current worker->fence_seq and not to the value that was pending when the request was created? if the user calls ucp_worker_fence() again after this request was submitted but before it is progressed, the EP is marked as satisfying the newer fence too, and a later operation on the same lane will not be fenced against this one. (same assumption already exists in ucp_ep_fence_strong(), so maybe worth a follow-up rather than this PR.)

@michael-beebe michael-beebe Oct 5, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Agreed, but I would keep this out of this PR. The weak branch makes the same assumption that ucp_ep_fence_strong() already makes, so this change is consistent with the current code. Recording the sequence number pending at request creation would have to change both paths, plus the request state, so it seems better as a separate follow-up.

/* The fence orders all earlier operations on this single
* lane, so later operations on it need no further fence */
ep->ext->fence_seq = ep->worker->fence_seq;
ep->ext->fenced_lanes = ep->ext->unflushed_lanes;
}
} else {
status = ucp_ep_fence_strong(ep);
}
} else if (ucs_unlikely(lane_map & ~ep->ext->fenced_lanes) &&

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 5, 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.

pls test fenced_lanes != 0 first: in the common case it is 0, so lane_map & ~fenced_lanes is non-zero and the ucs_unlikely hint is on the wrong sub-expression.

Suggested change
} else if (ucs_unlikely(lane_map & ~ep->ext->fenced_lanes) &&
} else if (ucs_unlikely(ep->ext->fenced_lanes != 0) &&
(lane_map & ~ep->ext->fenced_lanes)) {

@michael-beebe michael-beebe Oct 5, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Done in ca6fc76. The NULL check now comes first, inside the new ucp_ep_rma_is_strong_fence_required() helper, so the ucs_unlikely hint applies to the rare case. The field is now a single fenced_lane index, so the check is ucs_unlikely(fenced_lane != UCP_NULL_LANE) && (lane_map & ~UCS_BIT(fenced_lane)). Testing for NULL first also avoids computing UCS_BIT(UCP_NULL_LANE).

(ep->ext->fenced_lanes != 0)) {
/* A weak fence does not order earlier operations on the fenced lane
* with operations on another lane */
status = ucp_ep_fence_strong(ep);
} else {
status = UCS_OK;
}
Expand Down
4 changes: 3 additions & 1 deletion src/ucp/rma/rma_send.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,9 @@ ucp_put_send_short(ucp_ep_h ep, const void *buffer, size_t length,

tl_rkey = ucp_rkey_get_tl_rkey(rkey, rkey_config->put_short.rkey_index);

if (ucs_unlikely(ucp_ep_rma_is_fence_required(ep))) {
if (ucs_unlikely(ucp_ep_rma_is_fence_required(ep) ||

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 5, 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: can we hide this in an inline helper in rma.inl next to ucp_ep_rma_is_fence_required()? the fence bookkeeping is otherwise fully contained there, and put_send_short now open-codes part of it.

@michael-beebe michael-beebe Oct 5, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Done in ca6fc76. Added ucp_ep_rma_is_strong_fence_required(ep, lane_map) next to ucp_ep_rma_is_fence_required() in rma.inl. Both ucp_ep_rma_handle_fence() and ucp_put_send_short() use it now.

(ep->ext->fenced_lanes &
~UCS_BIT(rkey_config->put_short.lane)))) {
/* TODO: check support for fence in fast path short */
return UCS_ERR_NO_RESOURCE;
}
Expand Down
40 changes: 40 additions & 0 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 @@ -955,6 +955,25 @@ class test_ucp_ep_based_fence : public test_ucp_rma {
* iterations.
*/
sender().ep()->ext->unflushed_lanes = 0;
sender().ep()->ext->fenced_lanes = 0;
}

void do_rma_op_with_fence_applied_once(op_type_t op, void *sbuf,

@svc-nvidia-pr-review svc-nvidia-pr-review Oct 5, 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.

the second EXPECT_EQ(worker_fence_seq(), ep_fence_seq()) cannot fail — re-applying a fence does not change ep_fence_seq, so the test does not actually verify that the fence is applied only once. pls also cover the new path: weak fence on one lane followed by an operation on another lane (should trigger the strong fence), e.g. small op + fence + small op, then a large multi-lane op.

@michael-beebe michael-beebe Oct 5, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Done in ca6fc76.

  • The *_once tests now mock ep_fence on each lane's iface. They check that the fence count does not change on the next op after the fence, which is a real check. I confirmed they fail without the fence_seq update.
  • Added test_ep_based_fence_other_lane_{put,get}. Each does op + fence + op, then a 1-byte put (which can take the short path) and a large multi-lane put, and expects a strong fence (fenced_lane reset) whenever an op lands on a lane other than the fenced one.
  • I left out an atomic variant. When AMOs are emulated in software, the strong fence blocks until the receiver makes progress, which a single-threaded test cannot provide.
  • On a single-lane config (e.g. local cma) the other-lane check is trivially satisfied. It takes effect with MAX_RMA_LANES=2 on multi-lane transports such as rc.

size_t size, void *target,
ucp_rkey_h rkey)
{
do_rma_op_with_fence(op, sbuf, size, target, rkey);
flush_worker(sender());

/* Applying the fence satisfies it, so later operations on the
* endpoint must not apply it again */
EXPECT_EQ(worker_fence_seq(), ep_fence_seq());

do_rma_op(op, sbuf, size, target, rkey);
flush_worker(sender());
EXPECT_EQ(worker_fence_seq(), ep_fence_seq());

flush_sender_cleanup();
}

void do_rma_op_with_fence_before(op_type_t op, void *sbuf, size_t size,
Expand Down Expand Up @@ -999,6 +1018,27 @@ UCS_TEST_P(test_ucp_ep_based_fence, test_ep_based_fence_before_atomic) {
OP_ATOMIC, &test_ucp_ep_based_fence::do_rma_op_with_fence_before);
}

UCS_TEST_P(test_ucp_ep_based_fence, test_ep_based_fence_applied_once_put)
{
test_ep_based_fence_common(
OP_PUT,
&test_ucp_ep_based_fence::do_rma_op_with_fence_applied_once);
}

UCS_TEST_P(test_ucp_ep_based_fence, test_ep_based_fence_applied_once_get)
{
test_ep_based_fence_common(
OP_GET,
&test_ucp_ep_based_fence::do_rma_op_with_fence_applied_once);
}

UCS_TEST_P(test_ucp_ep_based_fence, test_ep_based_fence_applied_once_atomic)
{
test_ep_based_fence_common(
OP_ATOMIC,
&test_ucp_ep_based_fence::do_rma_op_with_fence_applied_once);
}

UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_ep_based_fence, all, "all")

class test_ucp_rma_sgl : public test_ucp_rma {
Expand Down
Loading