UCX Notified Communication by joe-explr · Pull Request #20 · devreal/ompi

joe-explr · 2026-03-31T20:54:10Z

Summary

This PR implements MPI-5.1 notified RMA in the osc/ucx component.

Design

Counters have their own UCX memory registration (notify_mem), kept
separate from the window data. The data region of an
MPI_WIN_FLAVOR_CREATE window belongs to the user and has no room for
them, and a dynamic window has no data region at all. As a result, the
counters work the same way for every window flavor.
Reserved vs. attached counters. notify_capacity counters are
registered per rank, and the value is the same on every rank. Only the
first notify_counts[rank] of them are attached, as set by
MPI_Win_set_num_notify. Notification indices are checked against the
target's attached count.
No reservation by default. mpi_assert_max_num_notify caps the growth
Counters grow on demand. MPI_Win_set_num_notify collectively
re-registers a larger region when a request exceeds the current capacity.
Ordering. Each notified operation issues its base operation, orders it
ahead of the counter update and then posts a UCX atomic add of 1 on the target's counter.
Request completion covers the notification. In the request-based
variants, the request is created after the counter atomic is posted.

Changes by file

`ompi/mca/osc/ucx/osc_ucx.h`

New fields in ompi_osc_ucx_module_t:

notify_counts: the number of counters attached at each rank (array of
size comm_size). It is kept the same on every rank by an allgather in
MPI_Win_set_num_notify, and is always <= notify_capacity.
notify_capacity: the number of counters reserved per rank (the same on
every rank). It starts at the largest mpi_assert_max_num_notify any rank
gave, or 0, and only grows.
notify_max_assert: this rank's own mpi_assert_max_num_notify, or 0.
MPI_Win_set_num_notify rejects requests from this rank above it.
notify_addrs: the remote base address of each rank's counter region
(NULL until counters are reserved).
notify_base: this rank's local counter array.
notify_mem: the UCX registration for the counters.

New prototypes for the notified operations, the counter-management
functions, and the internal helper ompi_osc_ucx_grow_notify_counters().

`ompi/mca/osc/ucx/osc_ucx_component.c`

Adds the 13 new notify entry points to the module template.
osc_ucx_reserved_notify_counters(): parses mpi_assert_max_num_notify.
If the key is absent or 0, there is no limit and no reservation. If the
value is malformed or negative, it returns MPI_ERR_INFO.
component_select():
- The existing creation-time allreduce grows from 4 to 5 values. The fifth
  value finds the largest assertion across ranks. A malformed value on any
  rank is sent as LONG_MIN, so every rank fails window creation together
  instead of some ranks hanging in a later collective.
- notify_counts starts at 0 everywhere, so no counters are attached at
  creation.
- If any rank gave an assertion, the agreed reservation is registered
  immediately.
ompi_osc_ucx_grow_notify_counters() (collective): allocates and zeroes a
new counter array, registers it (rkeys are exchanged with the whole group),
and allgathers the new base addresses. That allgather also acts as the
barrier that makes it safe to free the old registration.
ompi_osc_ucx_free(): releases all notification state.

`ompi/mca/osc/ucx/osc_ucx_comm.c`

Helpers:

osc_ucx_notify_counter_addr(): returns the remote address of a target's
counter notify.
CHECK_NOTIFY_IDX(): returns MPI_ERR_RMA_NOTIFICATION if the index is
outside the target's attached range.
osc_ucx_notify_target(): posts the UCP_ATOMIC_POST_OP_ADD of 1 through
notify_mem.

Counter management:

win_get_notify_value: progresses the UCX worker before reading. Without
this, a transport that emulates atomics in software never advances the
counter while the consumer spins on this call.
win_reset_notify_value: uses a UCX atomic swap on its own rank rather than
a CPU swap, so the reset is atomic with respect to network atomics arriving
from other ranks. It returns the previous value.
win_get_num_notify: a local query that returns
notify_counts[target_rank] (or MPI_ERR_RANK if the rank is out of
range).
win_set_num_notify (collective):
- Zeroes the local counters before the allgather, so no peer can wipe a
  notification after the call returns.
- Allgathers the requested counts. A negative request, or one above this
  rank's assertion, is sent as -1 rather than returning early, so every
  rank returns MPI_ERR_ARG together and the attached counts are left
  unchanged.
- Grows the capacity to the largest request if needed, then records the
  requested counts. Gathering the requested values directly means
  MPI_Win_get_num_notify returns exactly what was set, even when the count
  is lowered.
win_get_notify_bounds: NUM_SB is the reservation made at creation (or
INT_MAX if there was none), NUM_UB is INT_MAX (counters grow on
demand), and VALUE_UB is INT64_MAX (the value is returned as a signed
MPI_Count).

Notified operations:

Function	Base operation	Ordering before the counter increment
`put_notify`	put	fence
`get_notify`	get	blocking EP flush
`accumulate_notify`	accumulate	fence
`get_accumulate_notify`	get_accumulate	blocking EP flush
`rput_notify`	put	fence, then request over flush
`rget_notify`	get	blocking flush, then request over flush
`raccumulate_notify`	blocking `accumulate_req`	fence, then request over flush
`rget_accumulate_notify`	blocking `get_accumulate_req`	fence, then request over flush

devreal

Looks good, just a few comments.

devreal · 2026-04-01T12:27:54Z

+           + (uint64_t)notify * sizeof(uint64_t);
+}
+
+#define CHECK_NOTIFY_IDX(notify)                                              \


I would suggest passing in the module and checking against a value stored in there. OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS is just a crutch until we have a better solution

devreal · 2026-04-01T12:29:05Z

+
+    /* Counters are local memory — just read with a barrier to ensure
+     * any preceding remote writes to this counter are visible. */
+    opal_atomic_rmb();


Not sure this actually does anything. Shouldn't the read barrier come after the counter read to prevent subsequent reads from being reordered?

devreal · 2026-04-01T12:35:14Z

Oh, I guess we'll need to get #9 in first?

devreal · 2026-05-20T13:37:36Z

Is this ready to go in?

github-actions · 2026-07-24T20:01:05Z

Hello! The Git Commit Checker CI bot found a few problems with this PR:

ad1271b: Get and set notify changes

check_signed_off: does not contain a valid Signed-off-by line

Please fix these problems and, if necessary, force-push new commits back up to the PR branch. Thanks!

github-actions · 2026-09-11T18:42:08Z

Hello! The Git Commit Checker CI bot found a few problems with this PR:

d477d4f: Fix MPI_WIN_SET_NUM_NOTIFY hang on an invalid argu...

check_cherry_pick: contains a cherry pick message that refers to a commit that exists, but is in an as-yet unmerged pull request: fe59a33

Please fix these problems and, if necessary, force-push new commits back up to the PR branch. Thanks!

github-actions · 2026-09-11T18:53:38Z

Hello! The Git Commit Checker CI bot found a few problems with this PR:

d477d4f: Fix MPI_WIN_SET_NUM_NOTIFY hang on an invalid argu...

check_cherry_pick: contains a cherry pick message that refers to a commit that exists, but is in an as-yet unmerged pull request: fe59a33

Please fix these problems and, if necessary, force-push new commits back up to the PR branch. Thanks!

MPI_WIN_SET_NUM_NOTIFY is a blocking, synchronizing collective, but its num_notifications argument is local: MPI-5.1 section 12.6.1 states that the number of notification counters "can be different for different MPI processes". Both the C binding and osc/sm validated that argument and returned early, before the osc module's internal allgather. A single rank passing a bad value therefore returned an error while every other rank stayed blocked in that allgather forever, turning an erroneous argument into a hang. The binding rejected negative counts, and osc/sm additionally rejected counts above an mpi_assert_max_num_notify assertion -- a case the binding never covered, so osc/sm could hang even before this change. Drop the range check from the binding and carry each rank's verdict through the collective instead. osc/sm gathers ULONG_MAX as a sentinel that no legal count can collide with, since valid counts come from an int and never exceed INT_MAX; a single-process window has nobody to agree with and still answers immediately. All ranks then see the same gathered array and fail identically, so the window cannot end up half-reconfigured. osc/ucx already carried its verdict through the allgather and needed no change. Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

The previous fix computed "bad" but never sent the ULONG_MAX sentinel through the allgather, leaving the agree: label unused. Only -1 hit the sentinel, by accident of the cast; -5 became a huge request and crashed in the grow path, and a request above mpi_assert_max_num_notify grew instead of failing. Also correct the comment that called the assertion not a limit. Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

Test the broadcast descriptor with the shmem framework's validity flag instead of peeking at seg_name. segment_create sets the flag only on success, and the zeroed descriptor rank 0 sends when it cannot create the segment has it clear, so every rank still fails together. Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

Removing unecesaary comments Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

The tests expects the counter to grow past the set max_assert_num_notify value. This commit fixes the comments and the code to cap the growth at the flag value. Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

The base branch renamed the notified-communication error class from MPI_ERR_NOTIFY_IDX to MPI_ERR_RMA_NOTIFICATION; osc/ucx still referenced the old, now-undefined name. Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

Adds the four accumulate-flavored notified operations, the notification window attributes, and a configurable counter reservation, and corrects MPI_Win_set_num_notify. MPI_Accumulate_notify, MPI_Get_accumulate_notify, MPI_Raccumulate_notify and MPI_Rget_accumulate_notify follow the pattern already used by the put/get variants: issue the base operation, order it ahead of the counter with a fence (or a flush where a result must also be locally valid), then post the atomic increment. The five copies of that increment are now one helper. osc_win_get_notify_bounds was left NULL, so ompi_win_init cached zero for MPI_WIN_NOTIFICATION_NUM_SB, MPI_WIN_NOTIFICATION_NUM_UB and MPI_WIN_NOTIFICATION_VALUE_UB on every UCX window. Zero is how a component says it supports no counters at all, so a program that checked the attributes before calling MPI_Win_set_num_notify would skip notified communication even though UCX implements it. MPI_Win_set_num_notify only ever raised the attached count, but §12.6.1 says MPI_WIN_GET_NUM_NOTIFY returns the value given to MPI_WIN_SET_NUM_NOTIFY, so lowering it has to take effect. Gathering the requested value directly gives that, and also stops an origin from addressing counters the target has since detached. A rank whose argument is out of range no longer returns before the allgather either; that left the rest of the group blocked in a synchronizing collective. The reservation was a hard-coded 16. It is now the osc_ucx_num_notify_counters MCA parameter, overridable per window with the mpi_assert_max_num_notify info key, matching osc/sm. Because the counters share the window's memory registration, which cannot grow, the reservation is a real upper bound and is reported as such in MPI_WIN_NOTIFICATION_NUM_UB; osc/sm can reallocate its segment and so advertises INT_MAX. Ranks agree on one value so the shared-memory layout stays uniform. MPI_Win_notify_threshold (§12.6.3) is still unimplemented. Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The counters were appended to the window data and registered as part of the same region. For MPI_WIN_FLAVOR_CREATE the data region is the caller's buffer, sized for the window data alone, so this registered and then wrote past the end of memory the MPI library does not own -- at window creation, and again on every remote notification. Making the reservation configurable turned that from a fixed 128-byte overrun into one the user can scale. It went unnoticed because the notified communication tests only ever call MPI_Win_allocate. The counters now get their own registered region, alongside the window state rather than inside the window data. That removes the overrun, and since the region no longer depends on there being a data region at all, dynamic windows can support notified communication instead of being refused with MPI_ERR_RMA_FLAVOR. The shared-memory segment layout goes back to what it was before notified communication was added. Also from review of the previous commit: MPI_Win_set_num_notify validated its argument on each rank before the allgather, so a rank whose count was out of range returned while the rest of the group stayed blocked in the collective. The check now rides the collective. A refused call also no longer republishes the attached counts, so the group is not left half-reconfigured; the counters are still reset first, which is what makes the standard's "will not return until all processes have adjusted" hold. Window creation had the same defect: a malformed mpi_assert_max_num_notify value returned before the reservation allreduce. The failure is carried through that collective now. The reservation is also range-checked before being narrowed to int for the exchange. Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Notification counters are incremented by remote origins with a UCX atomic. Depending on the transport and the atomic mode in use, UCX may emulate that atomic in software on the target's worker rather than offloading it to the NIC. When it does, the counter only advances while the local worker is progressed. MPI_WIN_GET_NOTIFY_VALUE read the counter through a volatile pointer without progressing the worker, so a consumer polling it in a loop -- the natural way to wait for a notification, and the pattern the notified-communication interface exists to support -- could spin forever on a value that can never change. The operation completed only if the application happened to call some other MPI routine that progressed the worker as a side effect. Progress the worker before reading the counter, as every other spin-wait in this component already does. Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

The notification counters were reserved once at window creation and that reservation was a hard cap: MPI_WIN_SET_NUM_NOTIFY rejected any larger count, and MPI_WIN_NOTIFICATION_NUM_UB reported the reservation. MPI-5.1 section 12.2 defines the mpi_assert_max_num_notify info key as an assertion by the caller that it will not request more counters than the given value, and states that when the key is absent (zero) "the implementation does not assume any limit on the number of notification counters". Capping an unasserted window at the default reservation contradicts that, and left osc/ucx unable to satisfy programs that osc/sm -- which grows its counters -- accepts. Grow the counters instead. MPI_WIN_SET_NUM_NOTIFY already gathers every rank's requested count, so all ranks reach the same decision from the same array without extra communication, and grow to the largest request. Growing re-registers the region and re-exchanges base addresses, which is safe precisely because this procedure is a blocking, synchronizing collective: it resets every counter to zero, so a freshly allocated region is already the required contents, and it is erroneous to call it while an access epoch is open, so no remote atomic can be in flight against the region being replaced. The address allgather doubles as the barrier that lets the old region be released. The reservation is never shrunk, so a rank that lowers its count keeps its space and only genuine growth costs a re-registration. A window keeps a hard cap only when *every* rank asserted a bound; a rank that passed no key made no promise. NUM_UB now reports INT_MAX for an unasserted window and the asserted value otherwise, and NUM_SB follows the reservation, so both attributes stay consistent with what MPI_WIN_SET_NUM_NOTIFY will actually accept. Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

The request-based notified operations built their request from the underlying rput/rget/accumulate and only then issued the counter update, so the request described the data movement alone. Two consequences. MPI-5.1 section 12.6.4 advises that "completion at the origin entails that the notification counter update has been sent to the target and thus notifications do not rely on progress of decoupled MPI activities at the origin". Because the counter update is a non-fetching atomic issued after the request-bearing flush, MPI_WAIT could return with that update still queued locally, leaving a polling target waiting until the origin happened to re-enter MPI. The failure paths were also malformed: once the underlying operation had succeeded, *request held a live request, yet a failing fence or atomic returned an error. A caller following the usual convention -- on error, do not touch the request -- would then never complete it. The request could not simply be released there either, since UCX already holds it with a completion callback. Issue the data movement and the counter update first, and build the request afterwards. The request is attached to ucp_worker_flush_nb, which covers every operation already issued on the worker regardless of which registration it used, so the flush now covers the counter update as well. Allocating the request last also means every failure point precedes it and those paths return with *request untouched. The accumulate variants get the same treatment. Their underlying accumulate is synchronous -- it ends with a blocking flush and completes its request immediately -- so passing a NULL request runs the accumulate to completion at the target and lets the notification that follows be ordered after it, with a fresh request covering both. Factor the request construction shared by all six operations into osc_ucx_request_over_flush(). rput and rget keep their previous behaviour. Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

MPI_NOTIFICATION_NUM_SB now is set to INT_MAX. The default counters set at window creation is either 0 or the maximum of what max_assert_num_notify info key has. Removed changes that allocated default number of counters. Set up an all reduce call that finds the maximum counter size required to allocate in the state memory. ompi_osc_ucx_grow_notify_counters collectively reserve new_capacity notification counters per rank, replacing the current registration if there is one. Signed-off-by: Joseph Antony <joseph.antony@stonybrook.edu>

Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

devreal

LGTM. We'll iterate over some missing pieces (info handling, thresholds) but for now we'll move on.

devreal requested changes Apr 1, 2026

View reviewed changes

joe-explr force-pushed the notified-rma-ucx branch from 93195db to ad1271b Compare July 24, 2026 20:00

joe-explr force-pushed the notified-rma-ucx branch from ad1271b to 2ac49ff Compare July 24, 2026 20:14

joe-explr force-pushed the notified-rma-ucx branch from 2ac49ff to 8081698 Compare August 12, 2026 04:59

joe-explr requested a review from devreal August 13, 2026 14:31

devreal force-pushed the notified-rma branch from 17ba85d to bc02964 Compare August 17, 2026 17:28

joe-explr force-pushed the notified-rma-ucx branch from c5c7b00 to b140aa3 Compare September 4, 2026 20:15

joe-explr commented Sep 10, 2026

View reviewed changes

Comment thread ompi/mca/osc/ucx/osc_ucx.h Outdated

joe-explr commented Sep 10, 2026

View reviewed changes

Comment thread ompi/mca/osc/ucx/osc_ucx_comm.c Outdated

joe-explr commented Sep 10, 2026

View reviewed changes

Comment thread ompi/mca/osc/ucx/osc_ucx_comm.c Outdated

joe-explr force-pushed the notified-rma-ucx branch from b327326 to e93bc96 Compare September 11, 2026 18:41

joe-explr force-pushed the notified-rma-ucx branch from e93bc96 to 44b3202 Compare September 11, 2026 18:53

joe-explr added 4 commits September 11, 2026 14:55

osc/sm: Fixing comments

e88846a

Removing unecesaary comments Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

joe-explr force-pushed the notified-rma-ucx branch from 44b3202 to 4a510d7 Compare September 11, 2026 19:08

joe-explr and others added 8 commits September 11, 2026 16:55

osc/sm: Fixing failng test that caps counter growth

cece024

The tests expects the counter to grow past the set max_assert_num_notify value. This commit fixes the comments and the code to cap the growth at the flag value. Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

UCX Notified Communication

35f399e

Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

UCX Notified Communication changes

0666639

Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

Get and set Notify changes

a8ea3e6

Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

joe-explr added 4 commits September 11, 2026 17:16

osc/ucx: Comment change

ee29836

Signed-off-by: Joseph Antony <jajoseph.antony18@gmail.com>

joe-explr force-pushed the notified-rma-ucx branch from 05dd2a4 to ee29836 Compare September 11, 2026 21:16

devreal approved these changes Sep 11, 2026

View reviewed changes

Navigation Menu

UCX Notified Communication - #20

UCX Notified Communication#20
joe-explr wants to merge 16 commits into
devreal:notified-rmafrom
joe-explr:notified-rma-ucx

joe-explr commented Mar 31, 2026 •

edited

Loading

Uh oh!

devreal left a comment

Uh oh!

devreal Apr 1, 2026

Uh oh!

devreal Apr 1, 2026

Uh oh!

Uh oh!

Uh oh!

devreal commented Apr 1, 2026

Uh oh!

devreal commented May 20, 2026

Uh oh!

github-actions Bot commented Jul 24, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Sep 11, 2026

Uh oh!

github-actions Bot commented Sep 11, 2026

Uh oh!

devreal left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

joe-explr commented Mar 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Design

Changes by file

ompi/mca/osc/ucx/osc_ucx.h

ompi/mca/osc/ucx/osc_ucx_component.c

ompi/mca/osc/ucx/osc_ucx_comm.c

Uh oh!

devreal left a comment

Choose a reason for hiding this comment

Uh oh!

devreal Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

devreal Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

devreal commented Apr 1, 2026

Uh oh!

devreal commented May 20, 2026

Uh oh!

github-actions Bot commented Jul 24, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Sep 11, 2026

Uh oh!

github-actions Bot commented Sep 11, 2026

Uh oh!

devreal left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

joe-explr commented Mar 31, 2026 •

edited

Loading

`ompi/mca/osc/ucx/osc_ucx.h`

`ompi/mca/osc/ucx/osc_ucx_component.c`

`ompi/mca/osc/ucx/osc_ucx_comm.c`