From 2c880b84156950e05f66eee1ee59168485f03954 Mon Sep 17 00:00:00 2001 From: Joseph Antony Date: Tue, 31 Mar 2026 16:46:27 -0400 Subject: [PATCH 01/10] UCX Notified Communication Signed-off-by: Joseph Antony --- ompi/mca/osc/ucx/osc_ucx.h | 28 ++++ ompi/mca/osc/ucx/osc_ucx_comm.c | 192 +++++++++++++++++++++++++++ ompi/mca/osc/ucx/osc_ucx_component.c | 34 ++++- 3 files changed, 250 insertions(+), 4 deletions(-) diff --git a/ompi/mca/osc/ucx/osc_ucx.h b/ompi/mca/osc/ucx/osc_ucx.h index bc3dc8a91b3..0266f930ec3 100644 --- a/ompi/mca/osc/ucx/osc_ucx.h +++ b/ompi/mca/osc/ucx/osc_ucx.h @@ -27,6 +27,7 @@ #define OMPI_OSC_UCX_POST_PEER_MAX 32 #define OMPI_OSC_UCX_ATTACH_MAX 48 #define OMPI_OSC_UCX_MEM_ADDR_MAX_LEN 1024 +#define OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS 16 typedef struct ompi_osc_ucx_component { @@ -277,6 +278,33 @@ int ompi_osc_find_attached_region_position(ompi_osc_dynamic_win_info_t *dynamic_ int ompi_osc_ucx_dynamic_lock(ompi_osc_ucx_module_t *module, int target); int ompi_osc_ucx_dynamic_unlock(ompi_osc_ucx_module_t *module, int target); +int ompi_osc_ucx_put_notify(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + int notify, struct ompi_win_t *win); +int ompi_osc_ucx_get_notify(void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + int notify, struct ompi_win_t *win); +int ompi_osc_ucx_rput_notify(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + int notify, struct ompi_win_t *win, + struct ompi_request_t **request); +int ompi_osc_ucx_rget_notify(void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + int notify, struct ompi_win_t *win, + struct ompi_request_t **request); +int ompi_osc_ucx_win_get_notify_value(struct ompi_win_t *win, int notify, + OMPI_MPI_COUNT_TYPE *value); +int ompi_osc_ucx_win_reset_notify_value(struct ompi_win_t *win, int notify, + OMPI_MPI_COUNT_TYPE *value); + /* returns the size at the peer */ static inline size_t ompi_osc_ucx_get_size(ompi_osc_ucx_module_t *module, int rank) { diff --git a/ompi/mca/osc/ucx/osc_ucx_comm.c b/ompi/mca/osc/ucx/osc_ucx_comm.c index 0354edb71c0..53d1cf23cfd 100644 --- a/ompi/mca/osc/ucx/osc_ucx_comm.c +++ b/ompi/mca/osc/ucx/osc_ucx_comm.c @@ -603,6 +603,198 @@ int ompi_osc_ucx_get(void *origin_addr, size_t origin_count, } } +/* Returns the remote address of notify counter[notify] for the given target. + * Counters are appended directly after the target's window data in the same + * registered memory region (module->mem), so the rkey that covers window data + * also covers the counters. */ +static inline uint64_t +osc_ucx_notify_counter_addr(ompi_osc_ucx_module_t *module, int target, int notify) +{ + return module->addrs[target] + + ompi_osc_ucx_get_size(module, target) + + (uint64_t)notify * sizeof(uint64_t); +} + +#define CHECK_NOTIFY_IDX(notify) \ + if ((notify) < 0 || (notify) >= OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS) { \ + return MPI_ERR_NOTIFY_IDX; \ + } + +int ompi_osc_ucx_win_get_notify_value(struct ompi_win_t *win, int notify, + OMPI_MPI_COUNT_TYPE *value) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + int my_rank = ompi_comm_rank(module->comm); + + CHECK_NOTIFY_IDX(notify); + + /* Counters are local memory — just read with a barrier to ensure + * any preceding remote writes to this counter are visible. */ + opal_atomic_rmb(); + volatile uint64_t *counter = + (volatile uint64_t *)(module->addrs[my_rank] + module->size) + notify; + *value = (OMPI_MPI_COUNT_TYPE)*counter; + return OMPI_SUCCESS; +} + +int ompi_osc_ucx_win_reset_notify_value(struct ompi_win_t *win, int notify, + OMPI_MPI_COUNT_TYPE *value) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + int my_rank = ompi_comm_rank(module->comm); + + CHECK_NOTIFY_IDX(notify); + + volatile uint64_t *counter = + (volatile uint64_t *)(module->addrs[my_rank] + module->size) + notify; + *value = (OMPI_MPI_COUNT_TYPE)opal_atomic_swap_64((volatile int64_t *)counter, 0); + return OMPI_SUCCESS; +} + +int ompi_osc_ucx_put_notify(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + int notify, struct ompi_win_t *win) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + ucp_ep_h *ep; + int ret; + + CHECK_NOTIFY_IDX(notify); + + OSC_UCX_GET_DEFAULT_EP(ep, module, target); + + ret = ompi_osc_ucx_put(origin_addr, origin_count, origin_dt, + target, target_disp, target_count, target_dt, win); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* Flush to ensure the PUT is visible at the target before the counter + * increment arrives. */ + ret = opal_common_ucx_wpmem_fence(module->mem); + if (OPAL_SUCCESS != ret) { + return OMPI_ERROR; + } + + /* Atomically increment the target's notify counter in-place using the + * same mem handle as the window data. */ + ret = opal_common_ucx_wpmem_post(module->mem, + UCP_ATOMIC_POST_OP_ADD, 1, + target, sizeof(uint64_t), + osc_ucx_notify_counter_addr(module, target, notify), + ep); + return (OPAL_SUCCESS == ret) ? OMPI_SUCCESS : OMPI_ERROR; +} + +int ompi_osc_ucx_get_notify(void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + int notify, struct ompi_win_t *win) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + ucp_ep_h *ep; + int ret; + + CHECK_NOTIFY_IDX(notify); + + OSC_UCX_GET_DEFAULT_EP(ep, module, target); + + ret = ompi_osc_ucx_get(origin_addr, origin_count, origin_dt, + target, target_disp, target_count, target_dt, win); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* Flush to ensure the GET data is locally available before issuing the + * counter increment back to the target. */ + ret = opal_common_ucx_ctx_flush(module->ctx, OPAL_COMMON_UCX_SCOPE_EP, target); + if (OPAL_SUCCESS != ret) { + return OMPI_ERROR; + } + + ret = opal_common_ucx_wpmem_post(module->mem, + UCP_ATOMIC_POST_OP_ADD, 1, + target, sizeof(uint64_t), + osc_ucx_notify_counter_addr(module, target, notify), + ep); + return (OPAL_SUCCESS == ret) ? OMPI_SUCCESS : OMPI_ERROR; +} + +int ompi_osc_ucx_rput_notify(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + int notify, struct ompi_win_t *win, + struct ompi_request_t **request) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + ucp_ep_h *ep; + int ret; + + CHECK_NOTIFY_IDX(notify); + + OSC_UCX_GET_DEFAULT_EP(ep, module, target); + + ret = ompi_osc_ucx_rput(origin_addr, origin_count, origin_dt, + target, target_disp, target_count, target_dt, + win, request); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* Fence to order the PUT before the counter increment. */ + ret = opal_common_ucx_wpmem_fence(module->mem); + if (OPAL_SUCCESS != ret) { + return OMPI_ERROR; + } + + ret = opal_common_ucx_wpmem_post(module->mem, + UCP_ATOMIC_POST_OP_ADD, 1, + target, sizeof(uint64_t), + osc_ucx_notify_counter_addr(module, target, notify), + ep); + return (OPAL_SUCCESS == ret) ? OMPI_SUCCESS : OMPI_ERROR; +} + +int ompi_osc_ucx_rget_notify(void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + int notify, struct ompi_win_t *win, + struct ompi_request_t **request) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + ucp_ep_h *ep; + int ret; + + CHECK_NOTIFY_IDX(notify); + + OSC_UCX_GET_DEFAULT_EP(ep, module, target); + + ret = ompi_osc_ucx_rget(origin_addr, origin_count, origin_dt, + target, target_disp, target_count, target_dt, + win, request); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* Flush to ensure GET data is locally available before notifying target. */ + ret = opal_common_ucx_ctx_flush(module->ctx, OPAL_COMMON_UCX_SCOPE_EP, target); + if (OPAL_SUCCESS != ret) { + return OMPI_ERROR; + } + + ret = opal_common_ucx_wpmem_post(module->mem, + UCP_ATOMIC_POST_OP_ADD, 1, + target, sizeof(uint64_t), + osc_ucx_notify_counter_addr(module, target, notify), + ep); + return (OPAL_SUCCESS == ret) ? OMPI_SUCCESS : OMPI_ERROR; +} + static inline bool ompi_osc_need_acc_lock(ompi_osc_ucx_module_t *module, int target) { ompi_osc_ucx_lock_t *lock = NULL; diff --git a/ompi/mca/osc/ucx/osc_ucx_component.c b/ompi/mca/osc/ucx/osc_ucx_component.c index 635a53a3e0f..fcf9d23aee4 100644 --- a/ompi/mca/osc/ucx/osc_ucx_component.c +++ b/ompi/mca/osc/ucx/osc_ucx_component.c @@ -102,6 +102,13 @@ ompi_osc_ucx_module_t ompi_osc_ucx_module_template = { .osc_fetch_and_op = ompi_osc_ucx_fetch_and_op, .osc_get_accumulate = ompi_osc_ucx_get_accumulate, + .osc_put_notify = ompi_osc_ucx_put_notify, + .osc_get_notify = ompi_osc_ucx_get_notify, + .osc_rput_notify = ompi_osc_ucx_rput_notify, + .osc_rget_notify = ompi_osc_ucx_rget_notify, + .osc_win_get_notify_value = ompi_osc_ucx_win_get_notify_value, + .osc_win_reset_notify_value = ompi_osc_ucx_win_reset_notify_value, + .osc_rput = ompi_osc_ucx_rput, .osc_rget = ompi_osc_ucx_rget, .osc_raccumulate = ompi_osc_ucx_raccumulate, @@ -785,8 +792,10 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt /* create the segment */ size_t total = 0; + size_t notify_size = OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t); for (i = 0 ; i < comm_size ; ++i) { - total += ompi_osc_ucx_get_size(module, i); + /* each rank's slot holds its window data plus its notify counters */ + total += ompi_osc_ucx_get_size(module, i) + notify_size; } module->segment_base = NULL; @@ -849,14 +858,16 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt goto error; } - + /* Each rank's window slot is (peer_size + notify_size) bytes; the + * notify counters for rank i are at shmem_addrs[i] + peer_size. */ for (i = 0, total = 0; i < comm_size ; ++i) { size_t peer_size = ompi_osc_ucx_get_size(module, i); if (peer_size || !module->noncontig_shared_win) { module->shmem_addrs[i] = ((uint64_t) module->segment_base) + total; - total += peer_size; + total += peer_size + notify_size; } else { module->shmem_addrs[i] = (uint64_t)NULL; + total += notify_size; } } @@ -884,7 +895,16 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt ret = OMPI_ERR_BAD_PARAM; goto error; } - ret = opal_common_ucx_wpmem_create(module->ctx, mem_base, module->size, + /* Append notify counters after the window data in the same registered + * memory region. For ALLOCATE flavor the UCX allocator will hand back + * a buffer of this extended size; for CREATE/SHARED the user buffer is + * large enough to hold only the window data, but we still register the + * extra bytes so that remote atomic operations on the counters can use + * the same rkey as the window data. */ + size_t notify_reg_size = (flavor == MPI_WIN_FLAVOR_DYNAMIC) ? 0 : + OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t); + ret = opal_common_ucx_wpmem_create(module->ctx, mem_base, + module->size + notify_reg_size, mem_type, &exchange_len_info, OPAL_COMMON_UCX_WPMEM_ADDR_EXCHANGE_FULL, (void *)module->comm, @@ -957,6 +977,12 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt module->state.acc_lock = TARGET_LOCK_UNLOCKED; module->state.dynamic_lock = TARGET_LOCK_UNLOCKED; module->state.dynamic_win_count = 0; + + /* initialize notify counters to zero; they live at base + size */ + if (flavor != MPI_WIN_FLAVOR_DYNAMIC && *base != NULL) { + memset((char *)*base + module->size, 0, + OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t)); + } for (i = 0; i < OMPI_OSC_UCX_ATTACH_MAX; i++) { module->local_dynamic_win_info[i].refcnt = 0; } From 37ecf0a57c003f72f07c77f2bc980026158aeafd Mon Sep 17 00:00:00 2001 From: Joseph Antony Date: Mon, 20 Jul 2026 23:16:52 -0400 Subject: [PATCH 02/10] UCX Notified Communication changes Signed-off-by: Joseph Antony --- ompi/mca/osc/ucx/osc_ucx.h | 1 + ompi/mca/osc/ucx/osc_ucx_comm.c | 44 +++++++++++++++++++--------- ompi/mca/osc/ucx/osc_ucx_component.c | 11 +++++-- 3 files changed, 39 insertions(+), 17 deletions(-) diff --git a/ompi/mca/osc/ucx/osc_ucx.h b/ompi/mca/osc/ucx/osc_ucx.h index 0266f930ec3..021e7095a10 100644 --- a/ompi/mca/osc/ucx/osc_ucx.h +++ b/ompi/mca/osc/ucx/osc_ucx.h @@ -123,6 +123,7 @@ typedef struct ompi_osc_ucx_module { struct ompi_communicator_t *comm; int flavor; size_t size; + int num_notify; /* number of notify counters allocated per rank in this window */ size_t *sizes; /* used if not every process has the same size */ uint64_t *addrs; uint64_t *state_addrs; diff --git a/ompi/mca/osc/ucx/osc_ucx_comm.c b/ompi/mca/osc/ucx/osc_ucx_comm.c index 53d1cf23cfd..8b1022a78bf 100644 --- a/ompi/mca/osc/ucx/osc_ucx_comm.c +++ b/ompi/mca/osc/ucx/osc_ucx_comm.c @@ -615,8 +615,8 @@ osc_ucx_notify_counter_addr(ompi_osc_ucx_module_t *module, int target, int notif + (uint64_t)notify * sizeof(uint64_t); } -#define CHECK_NOTIFY_IDX(notify) \ - if ((notify) < 0 || (notify) >= OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS) { \ +#define CHECK_NOTIFY_IDX(module, notify) \ + if ((notify) < 0 || (notify) >= (module)->num_notify) { \ return MPI_ERR_NOTIFY_IDX; \ } @@ -626,14 +626,12 @@ int ompi_osc_ucx_win_get_notify_value(struct ompi_win_t *win, int notify, ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; int my_rank = ompi_comm_rank(module->comm); - CHECK_NOTIFY_IDX(notify); + CHECK_NOTIFY_IDX(module, notify); - /* Counters are local memory — just read with a barrier to ensure - * any preceding remote writes to this counter are visible. */ - opal_atomic_rmb(); volatile uint64_t *counter = (volatile uint64_t *)(module->addrs[my_rank] + module->size) + notify; *value = (OMPI_MPI_COUNT_TYPE)*counter; + opal_atomic_rmb(); return OMPI_SUCCESS; } @@ -642,12 +640,30 @@ int ompi_osc_ucx_win_reset_notify_value(struct ompi_win_t *win, int notify, { ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; int my_rank = ompi_comm_rank(module->comm); + uint64_t result_value = 0; + ucp_ep_h *ep; + int ret; - CHECK_NOTIFY_IDX(notify); + CHECK_NOTIFY_IDX(module, notify); - volatile uint64_t *counter = - (volatile uint64_t *)(module->addrs[my_rank] + module->size) + notify; - *value = (OMPI_MPI_COUNT_TYPE)opal_atomic_swap_64((volatile int64_t *)counter, 0); + OSC_UCX_GET_DEFAULT_EP(ep, module, my_rank); + + /* The counter is incremented by remote origins through UCX network atomic + * operations, so reset it with a UCX atomic swap (targeting our own rank) + * rather than a CPU atomic. That keeps the read-and-zero atomic with + * respect to those concurrent network atomics — a plain CPU swap is not + * ordered against them. The fetch returns the counter's previous value. */ + ret = opal_common_ucx_wpmem_fetch(module->mem, + UCP_ATOMIC_FETCH_OP_SWAP, 0, + my_rank, &result_value, sizeof(result_value), + osc_ucx_notify_counter_addr(module, my_rank, notify), + ep); + if (OPAL_SUCCESS != ret) { + OSC_UCX_VERBOSE(1, "opal_common_ucx_wpmem_fetch failed: %d", ret); + return OMPI_ERROR; + } + + *value = (OMPI_MPI_COUNT_TYPE)result_value; return OMPI_SUCCESS; } @@ -661,7 +677,7 @@ int ompi_osc_ucx_put_notify(const void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_IDX(notify); + CHECK_NOTIFY_IDX(module, notify); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -698,7 +714,7 @@ int ompi_osc_ucx_get_notify(void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_IDX(notify); + CHECK_NOTIFY_IDX(module, notify); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -734,7 +750,7 @@ int ompi_osc_ucx_rput_notify(const void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_IDX(notify); + CHECK_NOTIFY_IDX(module, notify); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -770,7 +786,7 @@ int ompi_osc_ucx_rget_notify(void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_IDX(notify); + CHECK_NOTIFY_IDX(module, notify); OSC_UCX_GET_DEFAULT_EP(ep, module, target); diff --git a/ompi/mca/osc/ucx/osc_ucx_component.c b/ompi/mca/osc/ucx/osc_ucx_component.c index fcf9d23aee4..a3cab7b7b68 100644 --- a/ompi/mca/osc/ucx/osc_ucx_component.c +++ b/ompi/mca/osc/ucx/osc_ucx_component.c @@ -679,6 +679,11 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt module->flavor = flavor; module->size = size; + /* Number of notify counters allocated per rank in this window. Stored in + * the module so every user references a single value instead of the + * OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS compile-time constant; this is the one + * place that can later be driven from the info object. */ + module->num_notify = OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS; module->no_locks = check_config_value_bool ("no_locks", info); module->acc_single_intrinsic = check_config_value_bool ("acc_single_intrinsic", info); module->skip_sync_check = false; @@ -792,7 +797,7 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt /* create the segment */ size_t total = 0; - size_t notify_size = OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t); + size_t notify_size = module->num_notify * sizeof(uint64_t); for (i = 0 ; i < comm_size ; ++i) { /* each rank's slot holds its window data plus its notify counters */ total += ompi_osc_ucx_get_size(module, i) + notify_size; @@ -902,7 +907,7 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt * extra bytes so that remote atomic operations on the counters can use * the same rkey as the window data. */ size_t notify_reg_size = (flavor == MPI_WIN_FLAVOR_DYNAMIC) ? 0 : - OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t); + module->num_notify * sizeof(uint64_t); ret = opal_common_ucx_wpmem_create(module->ctx, mem_base, module->size + notify_reg_size, mem_type, &exchange_len_info, @@ -981,7 +986,7 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt /* initialize notify counters to zero; they live at base + size */ if (flavor != MPI_WIN_FLAVOR_DYNAMIC && *base != NULL) { memset((char *)*base + module->size, 0, - OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t)); + module->num_notify * sizeof(uint64_t)); } for (i = 0; i < OMPI_OSC_UCX_ATTACH_MAX; i++) { module->local_dynamic_win_info[i].refcnt = 0; From 21395bcc464f018b7091bc404d0e411422d0fd65 Mon Sep 17 00:00:00 2001 From: Joseph Antony Date: Fri, 24 Jul 2026 15:38:02 -0400 Subject: [PATCH 03/10] Get and set Notify changes Signed-off-by: Joseph Antony --- ompi/mca/osc/ucx/osc_ucx.h | 11 ++- ompi/mca/osc/ucx/osc_ucx_comm.c | 109 +++++++++++++++++++++++++-- ompi/mca/osc/ucx/osc_ucx_component.c | 21 +++--- 3 files changed, 123 insertions(+), 18 deletions(-) diff --git a/ompi/mca/osc/ucx/osc_ucx.h b/ompi/mca/osc/ucx/osc_ucx.h index 021e7095a10..830b65d6a3d 100644 --- a/ompi/mca/osc/ucx/osc_ucx.h +++ b/ompi/mca/osc/ucx/osc_ucx.h @@ -123,7 +123,12 @@ typedef struct ompi_osc_ucx_module { struct ompi_communicator_t *comm; int flavor; size_t size; - int num_notify; /* number of notify counters allocated per rank in this window */ + int *notify_counts; /* per-rank number of notification counters attached at each + * rank (size comm_size). A fixed region of + * OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS counters is registered per + * rank at window creation; this tracks how many of them are + * currently attached, as set by MPI_WIN_SET_NUM_NOTIFY and + * kept consistent across the group by an allgather. */ size_t *sizes; /* used if not every process has the same size */ uint64_t *addrs; uint64_t *state_addrs; @@ -305,6 +310,10 @@ int ompi_osc_ucx_win_get_notify_value(struct ompi_win_t *win, int notify, OMPI_MPI_COUNT_TYPE *value); int ompi_osc_ucx_win_reset_notify_value(struct ompi_win_t *win, int notify, OMPI_MPI_COUNT_TYPE *value); +int ompi_osc_ucx_win_set_num_notify(struct ompi_win_t *win, struct opal_info_t *info, + int num_notifications); +int ompi_osc_ucx_win_get_num_notify(struct ompi_win_t *win, int target_rank, + int *num_notifications); /* returns the size at the peer */ static inline size_t ompi_osc_ucx_get_size(ompi_osc_ucx_module_t *module, int rank) diff --git a/ompi/mca/osc/ucx/osc_ucx_comm.c b/ompi/mca/osc/ucx/osc_ucx_comm.c index 8b1022a78bf..125d8e28c64 100644 --- a/ompi/mca/osc/ucx/osc_ucx_comm.c +++ b/ompi/mca/osc/ucx/osc_ucx_comm.c @@ -615,18 +615,33 @@ osc_ucx_notify_counter_addr(ompi_osc_ucx_module_t *module, int target, int notif + (uint64_t)notify * sizeof(uint64_t); } -#define CHECK_NOTIFY_IDX(module, notify) \ - if ((notify) < 0 || (notify) >= (module)->num_notify) { \ +/* A fixed region of OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS notification counters is + * registered per rank at window creation (see osc_ucx_component.c), but only + * the first notify_counts[rank] of them are considered *attached* by + * MPI_WIN_SET_NUM_NOTIFY. Per the MPI Standard it is erroneous to reference a + * counter that is out of range at the target, so validate against the target + * rank's attached count. */ +#define CHECK_NOTIFY_IDX(module, notify, rank) \ + if ((notify) < 0 || (notify) >= (module)->notify_counts[rank]) { \ return MPI_ERR_NOTIFY_IDX; \ } +/* Notification counters live in the window's registered memory region, which + * is not allocated for dynamic windows. Reject notified operations on them + * rather than issuing remote atomics against unregistered memory. */ +#define CHECK_NOTIFY_FLAVOR(module) \ + if (MPI_WIN_FLAVOR_DYNAMIC == (module)->flavor) { \ + return MPI_ERR_RMA_FLAVOR; \ + } + int ompi_osc_ucx_win_get_notify_value(struct ompi_win_t *win, int notify, OMPI_MPI_COUNT_TYPE *value) { ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; int my_rank = ompi_comm_rank(module->comm); - CHECK_NOTIFY_IDX(module, notify); + CHECK_NOTIFY_FLAVOR(module); + CHECK_NOTIFY_IDX(module, notify, my_rank); volatile uint64_t *counter = (volatile uint64_t *)(module->addrs[my_rank] + module->size) + notify; @@ -644,7 +659,8 @@ int ompi_osc_ucx_win_reset_notify_value(struct ompi_win_t *win, int notify, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_IDX(module, notify); + CHECK_NOTIFY_FLAVOR(module); + CHECK_NOTIFY_IDX(module, notify, my_rank); OSC_UCX_GET_DEFAULT_EP(ep, module, my_rank); @@ -667,6 +683,79 @@ int ompi_osc_ucx_win_reset_notify_value(struct ompi_win_t *win, int notify, return OMPI_SUCCESS; } +int ompi_osc_ucx_win_get_num_notify(struct ompi_win_t *win, int target_rank, + int *num_notifications) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + + if (target_rank < 0 || target_rank >= ompi_comm_size(module->comm)) { + return MPI_ERR_RANK; + } + + /* Local query (MPI_WIN_GET_NUM_NOTIFY, §12.6.1): return the number of + * notification counters currently attached at target_rank, as last + * published by MPI_WIN_SET_NUM_NOTIFY. */ + *num_notifications = module->notify_counts[target_rank]; + return OMPI_SUCCESS; +} + +int ompi_osc_ucx_win_set_num_notify(struct ompi_win_t *win, struct opal_info_t *info, + int num_notifications) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + int my_rank = ompi_comm_rank(module->comm); + int my_count; + int ret; + + (void) info; /* "mpi_assert_same_num_notifications" is an optimization hint only */ + + if (num_notifications < 0) { + return MPI_ERR_ARG; + } + + /* Notification counters live in the window's registered memory region, + * which is not allocated for dynamic windows. */ + CHECK_NOTIFY_FLAVOR(module); + + /* A fixed region of OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS counters is registered + * once at window creation (the effective MPI_WIN_NOTIFICATION_NUM_UB), so + * we can attach up to that many without re-registering memory. Requesting + * more than the capacity is out of range. */ + if (num_notifications > OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS) { + return MPI_ERR_ARG; + } + + /* The number of attached notification counters is never decreased + * (§12.6.1). */ + if (num_notifications > module->notify_counts[my_rank]) { + module->notify_counts[my_rank] = num_notifications; + } + + /* All notification counters (existing and newly attached) are reset to zero + * by this call. It is erroneous to call MPI_WIN_SET_NUM_NOTIFY while an + * access epoch is open, so no concurrent network atomics touch the counters + * and a plain local reset is sufficient. */ + if (0 != module->addrs[my_rank]) { + memset((void *)(module->addrs[my_rank] + module->size), 0, + OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t)); + } + opal_atomic_wmb(); + + /* Publish every rank's attached count to the whole group so that origins + * can validate notification indices against the target's count. This is + * the blocking, synchronizing collective required by the standard. */ + my_count = module->notify_counts[my_rank]; + ret = module->comm->c_coll->coll_allgather(&my_count, 1, MPI_INT, + module->notify_counts, 1, MPI_INT, + module->comm, + module->comm->c_coll->coll_allgather_module); + if (OMPI_SUCCESS != ret) { + return ret; + } + + return OMPI_SUCCESS; +} + int ompi_osc_ucx_put_notify(const void *origin_addr, size_t origin_count, struct ompi_datatype_t *origin_dt, int target, ptrdiff_t target_disp, size_t target_count, @@ -677,7 +766,8 @@ int ompi_osc_ucx_put_notify(const void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_IDX(module, notify); + CHECK_NOTIFY_FLAVOR(module); + CHECK_NOTIFY_IDX(module, notify, target); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -714,7 +804,8 @@ int ompi_osc_ucx_get_notify(void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_IDX(module, notify); + CHECK_NOTIFY_FLAVOR(module); + CHECK_NOTIFY_IDX(module, notify, target); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -750,7 +841,8 @@ int ompi_osc_ucx_rput_notify(const void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_IDX(module, notify); + CHECK_NOTIFY_FLAVOR(module); + CHECK_NOTIFY_IDX(module, notify, target); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -786,7 +878,8 @@ int ompi_osc_ucx_rget_notify(void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_IDX(module, notify); + CHECK_NOTIFY_FLAVOR(module); + CHECK_NOTIFY_IDX(module, notify, target); OSC_UCX_GET_DEFAULT_EP(ep, module, target); diff --git a/ompi/mca/osc/ucx/osc_ucx_component.c b/ompi/mca/osc/ucx/osc_ucx_component.c index a3cab7b7b68..f24899d2215 100644 --- a/ompi/mca/osc/ucx/osc_ucx_component.c +++ b/ompi/mca/osc/ucx/osc_ucx_component.c @@ -108,6 +108,8 @@ ompi_osc_ucx_module_t ompi_osc_ucx_module_template = { .osc_rget_notify = ompi_osc_ucx_rget_notify, .osc_win_get_notify_value = ompi_osc_ucx_win_get_notify_value, .osc_win_reset_notify_value = ompi_osc_ucx_win_reset_notify_value, + .osc_win_set_num_notify = ompi_osc_ucx_win_set_num_notify, + .osc_win_get_num_notify = ompi_osc_ucx_win_get_num_notify, .osc_rput = ompi_osc_ucx_rput, .osc_rget = ompi_osc_ucx_rget, @@ -679,11 +681,6 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt module->flavor = flavor; module->size = size; - /* Number of notify counters allocated per rank in this window. Stored in - * the module so every user references a single value instead of the - * OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS compile-time constant; this is the one - * place that can later be driven from the info object. */ - module->num_notify = OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS; module->no_locks = check_config_value_bool ("no_locks", info); module->acc_single_intrinsic = check_config_value_bool ("acc_single_intrinsic", info); module->skip_sync_check = false; @@ -797,7 +794,7 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt /* create the segment */ size_t total = 0; - size_t notify_size = module->num_notify * sizeof(uint64_t); + size_t notify_size = OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t); for (i = 0 ; i < comm_size ; ++i) { /* each rank's slot holds its window data plus its notify counters */ total += ompi_osc_ucx_get_size(module, i) + notify_size; @@ -907,7 +904,7 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt * extra bytes so that remote atomic operations on the counters can use * the same rkey as the window data. */ size_t notify_reg_size = (flavor == MPI_WIN_FLAVOR_DYNAMIC) ? 0 : - module->num_notify * sizeof(uint64_t); + OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t); ret = opal_common_ucx_wpmem_create(module->ctx, mem_base, module->size + notify_reg_size, mem_type, &exchange_len_info, @@ -966,6 +963,10 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt module->addrs = calloc(comm_size, sizeof(uint64_t)); module->state_addrs = calloc(comm_size, sizeof(uint64_t)); module->comm_world_ranks = calloc(comm_size, sizeof(uint64_t)); + /* Number of notification counters attached at each rank; starts at zero + * everywhere (consistent without communication) and is updated by + * MPI_WIN_SET_NUM_NOTIFY. Counters must be attached before use. */ + module->notify_counts = calloc(comm_size, sizeof(int)); for (i = 0; i < comm_size; i++) { memcpy(&(module->addrs[i]), recv_buf + i * 3 * sizeof(uint64_t), sizeof(uint64_t)); memcpy(&(module->state_addrs[i]), recv_buf + i * 3 * sizeof(uint64_t) + sizeof(uint64_t), sizeof(uint64_t)); @@ -983,10 +984,11 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt module->state.dynamic_lock = TARGET_LOCK_UNLOCKED; module->state.dynamic_win_count = 0; - /* initialize notify counters to zero; they live at base + size */ + /* initialize the fixed set of notify counters to zero; they live at + * base + size, immediately after the window data */ if (flavor != MPI_WIN_FLAVOR_DYNAMIC && *base != NULL) { memset((char *)*base + module->size, 0, - module->num_notify * sizeof(uint64_t)); + OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t)); } for (i = 0; i < OMPI_OSC_UCX_ATTACH_MAX; i++) { module->local_dynamic_win_info[i].refcnt = 0; @@ -1276,6 +1278,7 @@ int ompi_osc_ucx_free(struct ompi_win_t *win) { free(module->addrs); free(module->state_addrs); free(module->comm_world_ranks); + free(module->notify_counts); opal_common_ucx_wpmem_free(module->state_mem); if (NULL != module->mem) { From d944d8a6d515e89116feb1a519d45ec29516ec71 Mon Sep 17 00:00:00 2001 From: Joseph Antony Date: Wed, 12 Aug 2026 00:43:48 -0400 Subject: [PATCH 04/10] osc/ucx: use MPI_ERR_RMA_NOTIFICATION for invalid notification index 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 --- ompi/mca/osc/ucx/osc_ucx_comm.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ompi/mca/osc/ucx/osc_ucx_comm.c b/ompi/mca/osc/ucx/osc_ucx_comm.c index 125d8e28c64..c13062a1eec 100644 --- a/ompi/mca/osc/ucx/osc_ucx_comm.c +++ b/ompi/mca/osc/ucx/osc_ucx_comm.c @@ -623,7 +623,7 @@ osc_ucx_notify_counter_addr(ompi_osc_ucx_module_t *module, int target, int notif * rank's attached count. */ #define CHECK_NOTIFY_IDX(module, notify, rank) \ if ((notify) < 0 || (notify) >= (module)->notify_counts[rank]) { \ - return MPI_ERR_NOTIFY_IDX; \ + return MPI_ERR_RMA_NOTIFICATION; \ } /* Notification counters live in the window's registered memory region, which From 26bcd0daac53711fde0e32f9644472f694e241a3 Mon Sep 17 00:00:00 2001 From: Joseph Antony Date: Wed, 12 Aug 2026 01:27:16 -0400 Subject: [PATCH 05/10] osc/ucx: complete notified communication except thresholds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Co-Authored-By: Claude Opus 5 --- ompi/mca/osc/ucx/osc_ucx.h | 55 ++++- ompi/mca/osc/ucx/osc_ucx_comm.c | 288 +++++++++++++++++++++------ ompi/mca/osc/ucx/osc_ucx_component.c | 91 ++++++++- 3 files changed, 367 insertions(+), 67 deletions(-) diff --git a/ompi/mca/osc/ucx/osc_ucx.h b/ompi/mca/osc/ucx/osc_ucx.h index 830b65d6a3d..0e133c2fbd5 100644 --- a/ompi/mca/osc/ucx/osc_ucx.h +++ b/ompi/mca/osc/ucx/osc_ucx.h @@ -27,7 +27,11 @@ #define OMPI_OSC_UCX_POST_PEER_MAX 32 #define OMPI_OSC_UCX_ATTACH_MAX 48 #define OMPI_OSC_UCX_MEM_ADDR_MAX_LEN 1024 -#define OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS 16 +/* Default number of RMA notification counters reserved per MPI process in each + * window's registered memory region. Overridden per job by the + * osc_ucx_num_notify_counters MCA parameter and per window by the + * "mpi_assert_max_num_notify" info key. */ +#define OMPI_OSC_UCX_DEFAULT_NOTIFY_COUNTERS 16 typedef struct ompi_osc_ucx_component { @@ -44,6 +48,9 @@ typedef struct ompi_osc_ucx_component { bool no_locks; /* Default value of the no_locks info key for new windows */ bool acc_single_intrinsic; unsigned int priority; + /* Number of notification counters reserved per MPI process in each window, + * unless the window's info gives "mpi_assert_max_num_notify". */ + unsigned int num_notify_counters; /* directory where to place backing files */ char *backing_directory; } ompi_osc_ucx_component_t; @@ -123,12 +130,14 @@ typedef struct ompi_osc_ucx_module { struct ompi_communicator_t *comm; int flavor; size_t size; - int *notify_counts; /* per-rank number of notification counters attached at each - * rank (size comm_size). A fixed region of - * OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS counters is registered per - * rank at window creation; this tracks how many of them are - * currently attached, as set by MPI_WIN_SET_NUM_NOTIFY and - * kept consistent across the group by an allgather. */ + int *notify_counts; /* per-rank number of notification counters *attached* at each + * rank (size comm_size), as set by MPI_WIN_SET_NUM_NOTIFY and + * kept consistent across the group by an allgather. Always + * <= notify_capacity. */ + unsigned int notify_capacity; /* notification counters reserved per rank in the + * registered region at window creation. Agreed on + * across the group, and a hard upper bound: the + * registration cannot grow afterwards. */ size_t *sizes; /* used if not every process has the same size */ uint64_t *addrs; uint64_t *state_addrs; @@ -306,8 +315,40 @@ int ompi_osc_ucx_rget_notify(void *origin_addr, size_t origin_count, struct ompi_datatype_t *target_dt, int notify, struct ompi_win_t *win, struct ompi_request_t **request); +int ompi_osc_ucx_accumulate_notify(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + struct ompi_op_t *op, int notify, + struct ompi_win_t *win); +int ompi_osc_ucx_get_accumulate_notify(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + void *result_addr, size_t result_count, + struct ompi_datatype_t *result_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + struct ompi_op_t *op, int notify, + struct ompi_win_t *win); +int ompi_osc_ucx_raccumulate_notify(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + struct ompi_op_t *op, int notify, + struct ompi_win_t *win, + struct ompi_request_t **request); +int ompi_osc_ucx_rget_accumulate_notify(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + void *result_addr, size_t result_count, + struct ompi_datatype_t *result_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + struct ompi_op_t *op, int notify, + struct ompi_win_t *win, + struct ompi_request_t **request); int ompi_osc_ucx_win_get_notify_value(struct ompi_win_t *win, int notify, OMPI_MPI_COUNT_TYPE *value); +int ompi_osc_ucx_win_get_notify_bounds(struct ompi_win_t *win, int *num_sb, int *num_ub, + OMPI_MPI_COUNT_TYPE *value_ub); int ompi_osc_ucx_win_reset_notify_value(struct ompi_win_t *win, int notify, OMPI_MPI_COUNT_TYPE *value); int ompi_osc_ucx_win_set_num_notify(struct ompi_win_t *win, struct opal_info_t *info, diff --git a/ompi/mca/osc/ucx/osc_ucx_comm.c b/ompi/mca/osc/ucx/osc_ucx_comm.c index c13062a1eec..e6d8fb40ab4 100644 --- a/ompi/mca/osc/ucx/osc_ucx_comm.c +++ b/ompi/mca/osc/ucx/osc_ucx_comm.c @@ -17,6 +17,8 @@ #include "ompi/mca/osc/base/osc_base_obj_convert.h" #include "opal/mca/common/ucx/common_ucx.h" +#include + #include "osc_ucx.h" #include "osc_ucx_request.h" @@ -615,9 +617,9 @@ osc_ucx_notify_counter_addr(ompi_osc_ucx_module_t *module, int target, int notif + (uint64_t)notify * sizeof(uint64_t); } -/* A fixed region of OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS notification counters is - * registered per rank at window creation (see osc_ucx_component.c), but only - * the first notify_counts[rank] of them are considered *attached* by +/* A region of module->notify_capacity notification counters is registered per + * rank at window creation (see osc_ucx_component.c), but only the first + * notify_counts[rank] of them are considered *attached* by * MPI_WIN_SET_NUM_NOTIFY. Per the MPI Standard it is erroneous to reference a * counter that is out of range at the target, so validate against the target * rank's attached count. */ @@ -634,6 +636,22 @@ osc_ucx_notify_counter_addr(ompi_osc_ucx_module_t *module, int target, int notif return MPI_ERR_RMA_FLAVOR; \ } +/* Increments the target's notification counter once the preceding data + * operation has been ordered ahead of it. Shared by every notified + * operation; they differ only in which base operation they issue first and + * in whether a fence or a flush is needed to order it. */ +static inline int +osc_ucx_notify_target(ompi_osc_ucx_module_t *module, int target, int notify, + ucp_ep_h *ep) +{ + int ret = opal_common_ucx_wpmem_post(module->mem, + UCP_ATOMIC_POST_OP_ADD, 1, + target, sizeof(uint64_t), + osc_ucx_notify_counter_addr(module, target, notify), + ep); + return (OPAL_SUCCESS == ret) ? OMPI_SUCCESS : OMPI_ERROR; +} + int ompi_osc_ucx_win_get_notify_value(struct ompi_win_t *win, int notify, OMPI_MPI_COUNT_TYPE *value) { @@ -644,7 +662,7 @@ int ompi_osc_ucx_win_get_notify_value(struct ompi_win_t *win, int notify, CHECK_NOTIFY_IDX(module, notify, my_rank); volatile uint64_t *counter = - (volatile uint64_t *)(module->addrs[my_rank] + module->size) + notify; + (volatile uint64_t *)osc_ucx_notify_counter_addr(module, my_rank, notify); *value = (OMPI_MPI_COUNT_TYPE)*counter; opal_atomic_rmb(); return OMPI_SUCCESS; @@ -704,48 +722,45 @@ int ompi_osc_ucx_win_set_num_notify(struct ompi_win_t *win, struct opal_info_t * { ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; int my_rank = ompi_comm_rank(module->comm); - int my_count; - int ret; + int comm_size = ompi_comm_size(module->comm); + int requested = num_notifications; + int ret, i; (void) info; /* "mpi_assert_same_num_notifications" is an optimization hint only */ - if (num_notifications < 0) { - return MPI_ERR_ARG; - } - - /* Notification counters live in the window's registered memory region, - * which is not allocated for dynamic windows. */ + /* Notification counters live in the window's registered memory region, which + * is not allocated for dynamic windows. The flavor is uniform across the + * group, so returning here cannot desynchronize the collective below. */ CHECK_NOTIFY_FLAVOR(module); - /* A fixed region of OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS counters is registered - * once at window creation (the effective MPI_WIN_NOTIFICATION_NUM_UB), so - * we can attach up to that many without re-registering memory. Requesting - * more than the capacity is out of range. */ - if (num_notifications > OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS) { - return MPI_ERR_ARG; - } - - /* The number of attached notification counters is never decreased - * (§12.6.1). */ - if (num_notifications > module->notify_counts[my_rank]) { - module->notify_counts[my_rank] = num_notifications; + /* The counters are registered once at window creation and that registration + * cannot grow, so notify_capacity is a hard upper bound (reported as + * MPI_WIN_NOTIFICATION_NUM_UB). This is a synchronizing collective, so a + * rank with a bad argument must not return before the allgather below -- + * that would leave the rest of the group blocked in it. Mark the request + * instead and let every rank discover the error from the gathered values. */ + if (requested < 0 || (unsigned int) requested > module->notify_capacity) { + requested = -1; } /* All notification counters (existing and newly attached) are reset to zero * by this call. It is erroneous to call MPI_WIN_SET_NUM_NOTIFY while an * access epoch is open, so no concurrent network atomics touch the counters - * and a plain local reset is sufficient. */ - if (0 != module->addrs[my_rank]) { - memset((void *)(module->addrs[my_rank] + module->size), 0, - OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t)); + * and a plain local reset is sufficient. Resetting before the allgather + * means no rank can leave the collective and then observe a peer's stale + * counter. */ + if (0 <= requested && 0 != module->addrs[my_rank]) { + memset((void *)osc_ucx_notify_counter_addr(module, my_rank, 0), 0, + module->notify_capacity * sizeof(uint64_t)); } opal_atomic_wmb(); - /* Publish every rank's attached count to the whole group so that origins - * can validate notification indices against the target's count. This is - * the blocking, synchronizing collective required by the standard. */ - my_count = module->notify_counts[my_rank]; - ret = module->comm->c_coll->coll_allgather(&my_count, 1, MPI_INT, + /* Publish every rank's attached count to the whole group so that origins can + * validate notification indices against the target's count. This is the + * blocking, synchronizing collective required by the standard. Gathering + * the requested value directly is what makes MPI_WIN_GET_NUM_NOTIFY return + * the value given here, including when it lowers the count. */ + ret = module->comm->c_coll->coll_allgather(&requested, 1, MPI_INT, module->notify_counts, 1, MPI_INT, module->comm, module->comm->c_coll->coll_allgather_module); @@ -753,6 +768,45 @@ int ompi_osc_ucx_win_set_num_notify(struct ompi_win_t *win, struct opal_info_t * return ret; } + for (i = 0; i < comm_size; i++) { + if (0 > module->notify_counts[i]) { + /* Some rank asked for a count outside [0, notify_capacity]. Report + * the error on every rank rather than leaving part of the group + * believing the window was reconfigured. */ + module->notify_counts[i] = 0; + ret = MPI_ERR_ARG; + } + } + + return ret; +} + +int ompi_osc_ucx_win_get_notify_bounds(struct ompi_win_t *win, int *num_sb, int *num_ub, + OMPI_MPI_COUNT_TYPE *value_ub) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + + /* Dynamic windows have no registered data region to append counters to, so + * they support none; reporting zero is what tells a conforming program not + * to attempt notified communication on them. */ + if (MPI_WIN_FLAVOR_DYNAMIC == module->flavor) { + *num_sb = 0; + *num_ub = 0; + *value_ub = 0; + return OMPI_SUCCESS; + } + + /* The counters share the window's memory registration, which cannot grow + * after window creation, so the reservation is both the efficiently + * supported count and the hard upper bound. */ + *num_sb = (int) module->notify_capacity; + *num_ub = (int) module->notify_capacity; + + /* Counters are uint64_t and only ever incremented by one per notified + * operation, but they are returned to the user as a signed MPI_Count, so + * the representable maximum of that type is the real bound. */ + *value_ub = (OMPI_MPI_COUNT_TYPE) INT64_MAX; + return OMPI_SUCCESS; } @@ -786,12 +840,7 @@ int ompi_osc_ucx_put_notify(const void *origin_addr, size_t origin_count, /* Atomically increment the target's notify counter in-place using the * same mem handle as the window data. */ - ret = opal_common_ucx_wpmem_post(module->mem, - UCP_ATOMIC_POST_OP_ADD, 1, - target, sizeof(uint64_t), - osc_ucx_notify_counter_addr(module, target, notify), - ep); - return (OPAL_SUCCESS == ret) ? OMPI_SUCCESS : OMPI_ERROR; + return osc_ucx_notify_target(module, target, notify, ep); } int ompi_osc_ucx_get_notify(void *origin_addr, size_t origin_count, @@ -822,12 +871,7 @@ int ompi_osc_ucx_get_notify(void *origin_addr, size_t origin_count, return OMPI_ERROR; } - ret = opal_common_ucx_wpmem_post(module->mem, - UCP_ATOMIC_POST_OP_ADD, 1, - target, sizeof(uint64_t), - osc_ucx_notify_counter_addr(module, target, notify), - ep); - return (OPAL_SUCCESS == ret) ? OMPI_SUCCESS : OMPI_ERROR; + return osc_ucx_notify_target(module, target, notify, ep); } int ompi_osc_ucx_rput_notify(const void *origin_addr, size_t origin_count, @@ -859,12 +903,7 @@ int ompi_osc_ucx_rput_notify(const void *origin_addr, size_t origin_count, return OMPI_ERROR; } - ret = opal_common_ucx_wpmem_post(module->mem, - UCP_ATOMIC_POST_OP_ADD, 1, - target, sizeof(uint64_t), - osc_ucx_notify_counter_addr(module, target, notify), - ep); - return (OPAL_SUCCESS == ret) ? OMPI_SUCCESS : OMPI_ERROR; + return osc_ucx_notify_target(module, target, notify, ep); } int ompi_osc_ucx_rget_notify(void *origin_addr, size_t origin_count, @@ -896,12 +935,147 @@ int ompi_osc_ucx_rget_notify(void *origin_addr, size_t origin_count, return OMPI_ERROR; } - ret = opal_common_ucx_wpmem_post(module->mem, - UCP_ATOMIC_POST_OP_ADD, 1, - target, sizeof(uint64_t), - osc_ucx_notify_counter_addr(module, target, notify), - ep); - return (OPAL_SUCCESS == ret) ? OMPI_SUCCESS : OMPI_ERROR; + return osc_ucx_notify_target(module, target, notify, ep); +} + +int ompi_osc_ucx_accumulate_notify(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + struct ompi_op_t *op, int notify, + struct ompi_win_t *win) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + ucp_ep_h *ep; + int ret; + + CHECK_NOTIFY_FLAVOR(module); + CHECK_NOTIFY_IDX(module, notify, target); + + OSC_UCX_GET_DEFAULT_EP(ep, module, target); + + ret = ompi_osc_ucx_accumulate(origin_addr, origin_count, origin_dt, + target, target_disp, target_count, target_dt, + op, win); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* Fence so that the accumulate is applied at the target before the counter + * increment, as §12.6.4 requires. */ + ret = opal_common_ucx_wpmem_fence(module->mem); + if (OPAL_SUCCESS != ret) { + return OMPI_ERROR; + } + + return osc_ucx_notify_target(module, target, notify, ep); +} + +int ompi_osc_ucx_get_accumulate_notify(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + void *result_addr, size_t result_count, + struct ompi_datatype_t *result_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + struct ompi_op_t *op, int notify, + struct ompi_win_t *win) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + ucp_ep_h *ep; + int ret; + + CHECK_NOTIFY_FLAVOR(module); + CHECK_NOTIFY_IDX(module, notify, target); + + OSC_UCX_GET_DEFAULT_EP(ep, module, target); + + ret = ompi_osc_ucx_get_accumulate(origin_addr, origin_count, origin_dt, + result_addr, result_count, result_dt, + target, target_disp, target_count, target_dt, + op, win); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* Flush rather than fence: the result buffer must be locally valid, and the + * update must have been applied at the target, before the target may + * observe the notification. */ + ret = opal_common_ucx_ctx_flush(module->ctx, OPAL_COMMON_UCX_SCOPE_EP, target); + if (OPAL_SUCCESS != ret) { + return OMPI_ERROR; + } + + return osc_ucx_notify_target(module, target, notify, ep); +} + +int ompi_osc_ucx_raccumulate_notify(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + struct ompi_op_t *op, int notify, + struct ompi_win_t *win, + struct ompi_request_t **request) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + ucp_ep_h *ep; + int ret; + + CHECK_NOTIFY_FLAVOR(module); + CHECK_NOTIFY_IDX(module, notify, target); + + OSC_UCX_GET_DEFAULT_EP(ep, module, target); + + ret = ompi_osc_ucx_raccumulate(origin_addr, origin_count, origin_dt, + target, target_disp, target_count, target_dt, + op, win, request); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* Fence to order the accumulate ahead of the counter increment. */ + ret = opal_common_ucx_wpmem_fence(module->mem); + if (OPAL_SUCCESS != ret) { + return OMPI_ERROR; + } + + return osc_ucx_notify_target(module, target, notify, ep); +} + +int ompi_osc_ucx_rget_accumulate_notify(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + void *result_addr, size_t result_count, + struct ompi_datatype_t *result_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + struct ompi_op_t *op, int notify, + struct ompi_win_t *win, + struct ompi_request_t **request) +{ + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; + ucp_ep_h *ep; + int ret; + + CHECK_NOTIFY_FLAVOR(module); + CHECK_NOTIFY_IDX(module, notify, target); + + OSC_UCX_GET_DEFAULT_EP(ep, module, target); + + ret = ompi_osc_ucx_rget_accumulate(origin_addr, origin_count, origin_dt, + result_addr, result_count, result_dt, + target, target_disp, target_count, target_dt, + op, win, request); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* Flush so the fetched result is locally valid and the update has landed at + * the target before it can observe the notification. */ + ret = opal_common_ucx_ctx_flush(module->ctx, OPAL_COMMON_UCX_SCOPE_EP, target); + if (OPAL_SUCCESS != ret) { + return OMPI_ERROR; + } + + return osc_ucx_notify_target(module, target, notify, ep); } static inline bool ompi_osc_need_acc_lock(ompi_osc_ucx_module_t *module, int target) diff --git a/ompi/mca/osc/ucx/osc_ucx_component.c b/ompi/mca/osc/ucx/osc_ucx_component.c index f24899d2215..0ca2763de48 100644 --- a/ompi/mca/osc/ucx/osc_ucx_component.c +++ b/ompi/mca/osc/ucx/osc_ucx_component.c @@ -25,6 +25,8 @@ #include "osc_ucx.h" #include "osc_ucx_request.h" #include "opal/util/sys_limits.h" +#include "opal/util/info.h" +#include "opal/class/opal_cstring.h" #define memcpy_off(_dst, _src, _len, _off) \ memcpy(((char*)(_dst)) + (_off), _src, _len); \ @@ -106,10 +108,15 @@ ompi_osc_ucx_module_t ompi_osc_ucx_module_template = { .osc_get_notify = ompi_osc_ucx_get_notify, .osc_rput_notify = ompi_osc_ucx_rput_notify, .osc_rget_notify = ompi_osc_ucx_rget_notify, + .osc_accumulate_notify = ompi_osc_ucx_accumulate_notify, + .osc_get_accumulate_notify = ompi_osc_ucx_get_accumulate_notify, + .osc_raccumulate_notify = ompi_osc_ucx_raccumulate_notify, + .osc_rget_accumulate_notify = ompi_osc_ucx_rget_accumulate_notify, .osc_win_get_notify_value = ompi_osc_ucx_win_get_notify_value, .osc_win_reset_notify_value = ompi_osc_ucx_win_reset_notify_value, .osc_win_set_num_notify = ompi_osc_ucx_win_set_num_notify, .osc_win_get_num_notify = ompi_osc_ucx_win_get_num_notify, + .osc_win_get_notify_bounds = ompi_osc_ucx_win_get_notify_bounds, .osc_rput = ompi_osc_ucx_rput, .osc_rget = ompi_osc_ucx_rget, @@ -159,6 +166,46 @@ static bool check_config_value_bool (char *key, opal_info_t *info) return flag_value[0]; } +/* Read the mpi_assert_max_num_notify info key (MPI-5.1 section 12.2) to decide + * how many notification counters to reserve per MPI process. The counters are + * registered as part of the window's memory region and that registration cannot + * grow later, so the reservation is a hard upper bound on + * MPI_WIN_SET_NUM_NOTIFY for the lifetime of the window. */ +static int osc_ucx_reserved_notify_counters(opal_info_t *info, unsigned int *reserved) +{ + opal_cstring_t *value_string; + int flag = 0, value = 0; + + *reserved = mca_osc_ucx_component.num_notify_counters; + + if (NULL == info) { + return OMPI_SUCCESS; + } + + if (OMPI_SUCCESS != opal_info_get(info, "mpi_assert_max_num_notify", + &value_string, &flag) || !flag) { + return OMPI_SUCCESS; + } + + if (OPAL_SUCCESS != opal_cstring_to_int(value_string, &value)) { + OBJ_RELEASE(value_string); + return MPI_ERR_INFO; + } + OBJ_RELEASE(value_string); + + /* A negative value is a malformed key rather than "no assertion"; only 0 + * carries the "assume nothing" meaning. */ + if (value < 0) { + return MPI_ERR_INFO; + } + + if (0 != value) { + *reserved = (unsigned int) value; + } + + return OMPI_SUCCESS; +} + static int component_open(void) { opal_common_ucx_mca_register(); @@ -229,6 +276,23 @@ static int component_register(void) { MCA_BASE_VAR_SCOPE_GROUP, &ompi_osc_ucx_outstanding_ops_flush_threshold); free(description_str); + mca_osc_ucx_component.num_notify_counters = OMPI_OSC_UCX_DEFAULT_NOTIFY_COUNTERS; + + opal_asprintf(&description_str, + "Number of RMA notification counters reserved per MPI process " + "in the registered memory region of each window. Windows whose " + "info gives an mpi_assert_max_num_notify value use that instead. " + "This is a hard upper bound: the counters share the window's " + "registration, so MPI_Win_set_num_notify cannot exceed it " + "(default: %u)", + mca_osc_ucx_component.num_notify_counters); + (void) mca_base_component_var_register(&mca_osc_ucx_component.super.osc_version, + "num_notify_counters", description_str, + MCA_BASE_VAR_TYPE_UNSIGNED_INT, NULL, 0, 0, + OPAL_INFO_LVL_3, MCA_BASE_VAR_SCOPE_GROUP, + &mca_osc_ucx_component.num_notify_counters); + free(description_str); + opal_common_ucx_mca_var_register(&mca_osc_ucx_component.super.osc_version); if (0 == access ("/dev/shm", W_OK)) { @@ -681,6 +745,27 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt module->flavor = flavor; module->size = size; + + /* How many notification counters to reserve per MPI process. Read before + * the window memory is sized, since the reservation is part of its layout. */ + unsigned int notify_reserved = 0; + ret = osc_ucx_reserved_notify_counters(info, ¬ify_reserved); + if (OMPI_SUCCESS != ret) { + goto error; + } + /* info is allowed to differ between MPI processes, but the SHARED-flavor + * segment is laid out from a single per-rank stride, so agree on one + * reservation for the whole window. Taking the maximum keeps every rank's + * own assertion satisfiable. */ + int notify_reserved_int = (int) notify_reserved; + ret = module->comm->c_coll->coll_allreduce(MPI_IN_PLACE, ¬ify_reserved_int, 1, + MPI_INT, MPI_MAX, module->comm, + module->comm->c_coll->coll_allreduce_module); + if (OMPI_SUCCESS != ret) { + goto error; + } + module->notify_capacity = (unsigned int) notify_reserved_int; + module->no_locks = check_config_value_bool ("no_locks", info); module->acc_single_intrinsic = check_config_value_bool ("acc_single_intrinsic", info); module->skip_sync_check = false; @@ -794,7 +879,7 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt /* create the segment */ size_t total = 0; - size_t notify_size = OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t); + size_t notify_size = module->notify_capacity * sizeof(uint64_t); for (i = 0 ; i < comm_size ; ++i) { /* each rank's slot holds its window data plus its notify counters */ total += ompi_osc_ucx_get_size(module, i) + notify_size; @@ -904,7 +989,7 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt * extra bytes so that remote atomic operations on the counters can use * the same rkey as the window data. */ size_t notify_reg_size = (flavor == MPI_WIN_FLAVOR_DYNAMIC) ? 0 : - OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t); + module->notify_capacity * sizeof(uint64_t); ret = opal_common_ucx_wpmem_create(module->ctx, mem_base, module->size + notify_reg_size, mem_type, &exchange_len_info, @@ -988,7 +1073,7 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt * base + size, immediately after the window data */ if (flavor != MPI_WIN_FLAVOR_DYNAMIC && *base != NULL) { memset((char *)*base + module->size, 0, - OMPI_OSC_UCX_MAX_NOTIFY_COUNTERS * sizeof(uint64_t)); + module->notify_capacity * sizeof(uint64_t)); } for (i = 0; i < OMPI_OSC_UCX_ATTACH_MAX; i++) { module->local_dynamic_win_info[i].refcnt = 0; From caaa0bc1d2581bac02f1ef6d8d998dc8aba3326c Mon Sep 17 00:00:00 2001 From: Joseph Antony Date: Wed, 12 Aug 2026 02:27:48 -0400 Subject: [PATCH 06/10] osc/ucx: give notification counters their own registration 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 Co-Authored-By: Claude Opus 5 --- ompi/mca/osc/ucx/osc_ucx.h | 16 +++- ompi/mca/osc/ucx/osc_ucx_comm.c | 108 +++++++++++--------------- ompi/mca/osc/ucx/osc_ucx_component.c | 109 ++++++++++++++++++--------- 3 files changed, 128 insertions(+), 105 deletions(-) diff --git a/ompi/mca/osc/ucx/osc_ucx.h b/ompi/mca/osc/ucx/osc_ucx.h index 0e133c2fbd5..c02b0984010 100644 --- a/ompi/mca/osc/ucx/osc_ucx.h +++ b/ompi/mca/osc/ucx/osc_ucx.h @@ -134,10 +134,13 @@ typedef struct ompi_osc_ucx_module { * rank (size comm_size), as set by MPI_WIN_SET_NUM_NOTIFY and * kept consistent across the group by an allgather. Always * <= notify_capacity. */ - unsigned int notify_capacity; /* notification counters reserved per rank in the - * registered region at window creation. Agreed on - * across the group, and a hard upper bound: the - * registration cannot grow afterwards. */ + unsigned int notify_capacity; /* notification counters reserved per rank at window + * creation. Agreed on across the group, and a hard + * upper bound: the registration cannot grow + * afterwards. */ + uint64_t *notify_addrs; /* per-rank base address of the notification counters + * (size comm_size) */ + void *notify_base; /* this rank's counters; notify_capacity uint64_t */ size_t *sizes; /* used if not every process has the same size */ uint64_t *addrs; uint64_t *state_addrs; @@ -165,6 +168,11 @@ typedef struct ompi_osc_ucx_module { opal_common_ucx_ctx_t *ctx; opal_common_ucx_wpmem_t *mem; opal_common_ucx_wpmem_t *state_mem; + /* Notification counters get their own registration rather than being + * appended to the window data: the data region for MPI_WIN_FLAVOR_CREATE + * belongs to the user and has no room for them, and a dynamic window has + * no data region at all. */ + opal_common_ucx_wpmem_t *notify_mem; ompi_osc_ucx_mem_ranges_t *epoc_outstanding_ops_mems; bool skip_sync_check; bool noncontig_shared_win; diff --git a/ompi/mca/osc/ucx/osc_ucx_comm.c b/ompi/mca/osc/ucx/osc_ucx_comm.c index e6d8fb40ab4..648307f24bb 100644 --- a/ompi/mca/osc/ucx/osc_ucx_comm.c +++ b/ompi/mca/osc/ucx/osc_ucx_comm.c @@ -606,15 +606,12 @@ int ompi_osc_ucx_get(void *origin_addr, size_t origin_count, } /* Returns the remote address of notify counter[notify] for the given target. - * Counters are appended directly after the target's window data in the same - * registered memory region (module->mem), so the rkey that covers window data - * also covers the counters. */ + * The counters have their own registered region (module->notify_mem), separate + * from the window data, so this is independent of the window's flavor and size. */ static inline uint64_t osc_ucx_notify_counter_addr(ompi_osc_ucx_module_t *module, int target, int notify) { - return module->addrs[target] - + ompi_osc_ucx_get_size(module, target) - + (uint64_t)notify * sizeof(uint64_t); + return module->notify_addrs[target] + (uint64_t)notify * sizeof(uint64_t); } /* A region of module->notify_capacity notification counters is registered per @@ -628,23 +625,20 @@ osc_ucx_notify_counter_addr(ompi_osc_ucx_module_t *module, int target, int notif return MPI_ERR_RMA_NOTIFICATION; \ } -/* Notification counters live in the window's registered memory region, which - * is not allocated for dynamic windows. Reject notified operations on them - * rather than issuing remote atomics against unregistered memory. */ -#define CHECK_NOTIFY_FLAVOR(module) \ - if (MPI_WIN_FLAVOR_DYNAMIC == (module)->flavor) { \ - return MPI_ERR_RMA_FLAVOR; \ - } - /* Increments the target's notification counter once the preceding data * operation has been ordered ahead of it. Shared by every notified * operation; they differ only in which base operation they issue first and - * in whether a fence or a flush is needed to order it. */ + * in whether a fence or a flush is needed to order it. + * + * Note that for the request-based variants the data operation has already been + * issued and *request already handed back by the time this can fail, so an + * error return leaves that request outstanding and the caller still has to + * complete it. A transport failure here is not recoverable in any case. */ static inline int osc_ucx_notify_target(ompi_osc_ucx_module_t *module, int target, int notify, ucp_ep_h *ep) { - int ret = opal_common_ucx_wpmem_post(module->mem, + int ret = opal_common_ucx_wpmem_post(module->notify_mem, UCP_ATOMIC_POST_OP_ADD, 1, target, sizeof(uint64_t), osc_ucx_notify_counter_addr(module, target, notify), @@ -658,7 +652,6 @@ int ompi_osc_ucx_win_get_notify_value(struct ompi_win_t *win, int notify, ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; int my_rank = ompi_comm_rank(module->comm); - CHECK_NOTIFY_FLAVOR(module); CHECK_NOTIFY_IDX(module, notify, my_rank); volatile uint64_t *counter = @@ -677,7 +670,6 @@ int ompi_osc_ucx_win_reset_notify_value(struct ompi_win_t *win, int notify, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_FLAVOR(module); CHECK_NOTIFY_IDX(module, notify, my_rank); OSC_UCX_GET_DEFAULT_EP(ep, module, my_rank); @@ -687,7 +679,7 @@ int ompi_osc_ucx_win_reset_notify_value(struct ompi_win_t *win, int notify, * rather than a CPU atomic. That keeps the read-and-zero atomic with * respect to those concurrent network atomics — a plain CPU swap is not * ordered against them. The fetch returns the counter's previous value. */ - ret = opal_common_ucx_wpmem_fetch(module->mem, + ret = opal_common_ucx_wpmem_fetch(module->notify_mem, UCP_ATOMIC_FETCH_OP_SWAP, 0, my_rank, &result_value, sizeof(result_value), osc_ucx_notify_counter_addr(module, my_rank, notify), @@ -721,18 +713,13 @@ int ompi_osc_ucx_win_set_num_notify(struct ompi_win_t *win, struct opal_info_t * int num_notifications) { ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; - int my_rank = ompi_comm_rank(module->comm); int comm_size = ompi_comm_size(module->comm); int requested = num_notifications; + int *requested_counts; int ret, i; (void) info; /* "mpi_assert_same_num_notifications" is an optimization hint only */ - /* Notification counters live in the window's registered memory region, which - * is not allocated for dynamic windows. The flavor is uniform across the - * group, so returning here cannot desynchronize the collective below. */ - CHECK_NOTIFY_FLAVOR(module); - /* The counters are registered once at window creation and that registration * cannot grow, so notify_capacity is a hard upper bound (reported as * MPI_WIN_NOTIFICATION_NUM_UB). This is a synchronizing collective, so a @@ -743,42 +730,53 @@ int ompi_osc_ucx_win_set_num_notify(struct ompi_win_t *win, struct opal_info_t * requested = -1; } + requested_counts = malloc(comm_size * sizeof(int)); + if (NULL == requested_counts) { + return OMPI_ERR_TEMP_OUT_OF_RESOURCE; + } + /* All notification counters (existing and newly attached) are reset to zero - * by this call. It is erroneous to call MPI_WIN_SET_NUM_NOTIFY while an + * by this call. Resetting before the allgather is what makes the standard's + * "will not return until ... all processes have adjusted the number of + * notification counters" hold: completing the collective implies every rank + * has already reset, so no rank can return and then have a peer wipe the + * notification it just delivered. It is erroneous to call this while an * access epoch is open, so no concurrent network atomics touch the counters - * and a plain local reset is sufficient. Resetting before the allgather - * means no rank can leave the collective and then observe a peer's stale - * counter. */ - if (0 <= requested && 0 != module->addrs[my_rank]) { - memset((void *)osc_ucx_notify_counter_addr(module, my_rank, 0), 0, - module->notify_capacity * sizeof(uint64_t)); + * and a plain local reset is sufficient. */ + if (NULL != module->notify_base) { + memset(module->notify_base, 0, module->notify_capacity * sizeof(uint64_t)); } opal_atomic_wmb(); - /* Publish every rank's attached count to the whole group so that origins can - * validate notification indices against the target's count. This is the - * blocking, synchronizing collective required by the standard. Gathering + /* Publish every rank's requested count to the whole group so that origins + * can validate notification indices against the target's count. Gathering * the requested value directly is what makes MPI_WIN_GET_NUM_NOTIFY return * the value given here, including when it lowers the count. */ ret = module->comm->c_coll->coll_allgather(&requested, 1, MPI_INT, - module->notify_counts, 1, MPI_INT, + requested_counts, 1, MPI_INT, module->comm, module->comm->c_coll->coll_allgather_module); if (OMPI_SUCCESS != ret) { + free(requested_counts); return ret; } for (i = 0; i < comm_size; i++) { - if (0 > module->notify_counts[i]) { - /* Some rank asked for a count outside [0, notify_capacity]. Report - * the error on every rank rather than leaving part of the group - * believing the window was reconfigured. */ - module->notify_counts[i] = 0; - ret = MPI_ERR_ARG; + if (0 > requested_counts[i]) { + /* Some rank asked for a count outside [0, notify_capacity]. Every + * rank sees the same gathered array and bails identically, so the + * attached counts stay as they were rather than the group ending up + * half-reconfigured. The counters have been zeroed, which is + * harmless for a call that is erroneous anyway. */ + free(requested_counts); + return MPI_ERR_ARG; } } - return ret; + memcpy(module->notify_counts, requested_counts, comm_size * sizeof(int)); + free(requested_counts); + + return OMPI_SUCCESS; } int ompi_osc_ucx_win_get_notify_bounds(struct ompi_win_t *win, int *num_sb, int *num_ub, @@ -786,19 +784,9 @@ int ompi_osc_ucx_win_get_notify_bounds(struct ompi_win_t *win, int *num_sb, int { ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; - /* Dynamic windows have no registered data region to append counters to, so - * they support none; reporting zero is what tells a conforming program not - * to attempt notified communication on them. */ - if (MPI_WIN_FLAVOR_DYNAMIC == module->flavor) { - *num_sb = 0; - *num_ub = 0; - *value_ub = 0; - return OMPI_SUCCESS; - } - - /* The counters share the window's memory registration, which cannot grow - * after window creation, so the reservation is both the efficiently - * supported count and the hard upper bound. */ + /* The counters have their own registration, which cannot grow after window + * creation, so the reservation is both the efficiently supported count and + * the hard upper bound. It does not depend on the window's flavor. */ *num_sb = (int) module->notify_capacity; *num_ub = (int) module->notify_capacity; @@ -820,7 +808,6 @@ int ompi_osc_ucx_put_notify(const void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_FLAVOR(module); CHECK_NOTIFY_IDX(module, notify, target); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -853,7 +840,6 @@ int ompi_osc_ucx_get_notify(void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_FLAVOR(module); CHECK_NOTIFY_IDX(module, notify, target); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -885,7 +871,6 @@ int ompi_osc_ucx_rput_notify(const void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_FLAVOR(module); CHECK_NOTIFY_IDX(module, notify, target); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -917,7 +902,6 @@ int ompi_osc_ucx_rget_notify(void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_FLAVOR(module); CHECK_NOTIFY_IDX(module, notify, target); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -949,7 +933,6 @@ int ompi_osc_ucx_accumulate_notify(const void *origin_addr, size_t origin_count, ucp_ep_h *ep; int ret; - CHECK_NOTIFY_FLAVOR(module); CHECK_NOTIFY_IDX(module, notify, target); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -984,7 +967,6 @@ int ompi_osc_ucx_get_accumulate_notify(const void *origin_addr, size_t origin_co ucp_ep_h *ep; int ret; - CHECK_NOTIFY_FLAVOR(module); CHECK_NOTIFY_IDX(module, notify, target); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -1020,7 +1002,6 @@ int ompi_osc_ucx_raccumulate_notify(const void *origin_addr, size_t origin_count ucp_ep_h *ep; int ret; - CHECK_NOTIFY_FLAVOR(module); CHECK_NOTIFY_IDX(module, notify, target); OSC_UCX_GET_DEFAULT_EP(ep, module, target); @@ -1055,7 +1036,6 @@ int ompi_osc_ucx_rget_accumulate_notify(const void *origin_addr, size_t origin_c ucp_ep_h *ep; int ret; - CHECK_NOTIFY_FLAVOR(module); CHECK_NOTIFY_IDX(module, notify, target); OSC_UCX_GET_DEFAULT_EP(ep, module, target); diff --git a/ompi/mca/osc/ucx/osc_ucx_component.c b/ompi/mca/osc/ucx/osc_ucx_component.c index 0ca2763de48..4776f287a1a 100644 --- a/ompi/mca/osc/ucx/osc_ucx_component.c +++ b/ompi/mca/osc/ucx/osc_ucx_component.c @@ -28,6 +28,8 @@ #include "opal/util/info.h" #include "opal/class/opal_cstring.h" +#include + #define memcpy_off(_dst, _src, _len, _off) \ memcpy(((char*)(_dst)) + (_off), _src, _len); \ (_off) += (_len); @@ -632,7 +634,8 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt opal_common_ucx_mem_type_t mem_type; char *my_mem_addr; int my_mem_addr_size; - uint64_t my_info[3] = {0}; + uint64_t my_info[4] = {0}; + void *notify_base = NULL; char *recv_buf = NULL; void *dynamic_base = NULL; unsigned long adjusted_size = size; @@ -747,24 +750,35 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt module->size = size; /* How many notification counters to reserve per MPI process. Read before - * the window memory is sized, since the reservation is part of its layout. */ + * the counter region is allocated below. A malformed info value must not + * make this rank skip the allreduce that follows -- that would leave the + * rest of the group blocked in window creation -- so the failure is carried + * through the collective as a negative reservation instead. */ unsigned int notify_reserved = 0; - ret = osc_ucx_reserved_notify_counters(info, ¬ify_reserved); - if (OMPI_SUCCESS != ret) { - goto error; - } - /* info is allowed to differ between MPI processes, but the SHARED-flavor - * segment is laid out from a single per-rank stride, so agree on one + int notify_values[2]; + /* The reservation is exchanged as an int and is used to size an allocation, + * so a value that does not fit is a bad configuration rather than a request + * to honor. Both failures ride the flag, not the value: MPI_MAX would hide + * a sentinel value behind some other rank's larger reservation. */ + bool notify_bad = (OMPI_SUCCESS != osc_ucx_reserved_notify_counters(info, ¬ify_reserved)) + || notify_reserved > (unsigned int) INT_MAX; + notify_values[0] = notify_bad ? 1 : 0; + notify_values[1] = notify_bad ? 0 : (int) notify_reserved; + + /* info is allowed to differ between MPI processes, so agree on one * reservation for the whole window. Taking the maximum keeps every rank's - * own assertion satisfiable. */ - int notify_reserved_int = (int) notify_reserved; - ret = module->comm->c_coll->coll_allreduce(MPI_IN_PLACE, ¬ify_reserved_int, 1, + * own assertion satisfiable, and propagates any rank's failure flag. */ + ret = module->comm->c_coll->coll_allreduce(MPI_IN_PLACE, notify_values, 2, MPI_INT, MPI_MAX, module->comm, module->comm->c_coll->coll_allreduce_module); if (OMPI_SUCCESS != ret) { goto error; } - module->notify_capacity = (unsigned int) notify_reserved_int; + if (0 != notify_values[0]) { + ret = MPI_ERR_INFO; + goto error; + } + module->notify_capacity = (unsigned int) notify_values[1]; module->no_locks = check_config_value_bool ("no_locks", info); module->acc_single_intrinsic = check_config_value_bool ("acc_single_intrinsic", info); @@ -879,10 +893,8 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt /* create the segment */ size_t total = 0; - size_t notify_size = module->notify_capacity * sizeof(uint64_t); for (i = 0 ; i < comm_size ; ++i) { - /* each rank's slot holds its window data plus its notify counters */ - total += ompi_osc_ucx_get_size(module, i) + notify_size; + total += ompi_osc_ucx_get_size(module, i); } module->segment_base = NULL; @@ -945,16 +957,13 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt goto error; } - /* Each rank's window slot is (peer_size + notify_size) bytes; the - * notify counters for rank i are at shmem_addrs[i] + peer_size. */ for (i = 0, total = 0; i < comm_size ; ++i) { size_t peer_size = ompi_osc_ucx_get_size(module, i); if (peer_size || !module->noncontig_shared_win) { module->shmem_addrs[i] = ((uint64_t) module->segment_base) + total; - total += peer_size + notify_size; + total += peer_size; } else { module->shmem_addrs[i] = (uint64_t)NULL; - total += notify_size; } } @@ -982,16 +991,8 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt ret = OMPI_ERR_BAD_PARAM; goto error; } - /* Append notify counters after the window data in the same registered - * memory region. For ALLOCATE flavor the UCX allocator will hand back - * a buffer of this extended size; for CREATE/SHARED the user buffer is - * large enough to hold only the window data, but we still register the - * extra bytes so that remote atomic operations on the counters can use - * the same rkey as the window data. */ - size_t notify_reg_size = (flavor == MPI_WIN_FLAVOR_DYNAMIC) ? 0 : - module->notify_capacity * sizeof(uint64_t); ret = opal_common_ucx_wpmem_create(module->ctx, mem_base, - module->size + notify_reg_size, + module->size, mem_type, &exchange_len_info, OPAL_COMMON_UCX_WPMEM_ADDR_EXCHANGE_FULL, (void *)module->comm, @@ -1006,6 +1007,37 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt ucp_rkey_buffer_release(my_mem_addr); } + /* Notification counters live in their own registered region, like the + * window state does. Appending them to the window data would mean writing + * past the end of the user's buffer for MPI_WIN_FLAVOR_CREATE, and would + * leave dynamic windows -- which have no data region -- with nowhere to put + * them. */ + if (0 != module->notify_capacity) { + module->notify_base = calloc(module->notify_capacity, sizeof(uint64_t)); + if (NULL == module->notify_base) { + ret = OMPI_ERR_TEMP_OUT_OF_RESOURCE; + goto error; + } + + notify_base = module->notify_base; + ret = opal_common_ucx_wpmem_create(module->ctx, ¬ify_base, + module->notify_capacity * sizeof(uint64_t), + OPAL_COMMON_UCX_MEM_MAP, + &exchange_len_info, + OPAL_COMMON_UCX_WPMEM_ADDR_EXCHANGE_FULL, + (void *)module->comm, + &my_mem_addr, &my_mem_addr_size, + &module->notify_mem); + if (ret != OMPI_SUCCESS) { + goto error; + } + + if (my_mem_addr_size != 0) { + /* rkey object is already distributed among comm processes */ + ucp_rkey_buffer_release(my_mem_addr); + } + } + state_base = (void *)&(module->state); ret = opal_common_ucx_wpmem_create(module->ctx, &state_base, sizeof(ompi_osc_ucx_state_t), @@ -1033,6 +1065,7 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt } my_info[1] = (uint64_t)state_base; my_info[2] = ompi_comm_rank(&ompi_mpi_comm_world.comm); + my_info[3] = (uint64_t)module->notify_base; recv_buf = (char *)calloc(comm_size, sizeof(my_info)); ret = comm->c_coll->coll_allgather((void *)my_info, sizeof(my_info), @@ -1052,10 +1085,13 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt * everywhere (consistent without communication) and is updated by * MPI_WIN_SET_NUM_NOTIFY. Counters must be attached before use. */ module->notify_counts = calloc(comm_size, sizeof(int)); + module->notify_addrs = calloc(comm_size, sizeof(uint64_t)); for (i = 0; i < comm_size; i++) { - memcpy(&(module->addrs[i]), recv_buf + i * 3 * sizeof(uint64_t), sizeof(uint64_t)); - memcpy(&(module->state_addrs[i]), recv_buf + i * 3 * sizeof(uint64_t) + sizeof(uint64_t), sizeof(uint64_t)); - memcpy(&(module->comm_world_ranks[i]), recv_buf + i * 3 * sizeof(uint64_t) + 2 * sizeof(uint64_t), sizeof(uint64_t)); + const char *entry = recv_buf + i * sizeof(my_info); + memcpy(&(module->addrs[i]), entry, sizeof(uint64_t)); + memcpy(&(module->state_addrs[i]), entry + sizeof(uint64_t), sizeof(uint64_t)); + memcpy(&(module->comm_world_ranks[i]), entry + 2 * sizeof(uint64_t), sizeof(uint64_t)); + memcpy(&(module->notify_addrs[i]), entry + 3 * sizeof(uint64_t), sizeof(uint64_t)); } free(recv_buf); @@ -1069,12 +1105,6 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt module->state.dynamic_lock = TARGET_LOCK_UNLOCKED; module->state.dynamic_win_count = 0; - /* initialize the fixed set of notify counters to zero; they live at - * base + size, immediately after the window data */ - if (flavor != MPI_WIN_FLAVOR_DYNAMIC && *base != NULL) { - memset((char *)*base + module->size, 0, - module->notify_capacity * sizeof(uint64_t)); - } for (i = 0; i < OMPI_OSC_UCX_ATTACH_MAX; i++) { module->local_dynamic_win_info[i].refcnt = 0; } @@ -1364,11 +1394,16 @@ int ompi_osc_ucx_free(struct ompi_win_t *win) { free(module->state_addrs); free(module->comm_world_ranks); free(module->notify_counts); + free(module->notify_addrs); opal_common_ucx_wpmem_free(module->state_mem); if (NULL != module->mem) { opal_common_ucx_wpmem_free(module->mem); } + if (NULL != module->notify_mem) { + opal_common_ucx_wpmem_free(module->notify_mem); + } + free(module->notify_base); opal_common_ucx_wpctx_release(module->ctx); From 74e978a6984b8b4e08b35e5c02feaaa9d1ecb606 Mon Sep 17 00:00:00 2001 From: Joseph Antony Date: Wed, 12 Aug 2026 18:39:56 -0400 Subject: [PATCH 07/10] osc/ucx: progress the worker in MPI_WIN_GET_NOTIFY_VALUE 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 --- ompi/mca/osc/ucx/osc_ucx_comm.c | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/ompi/mca/osc/ucx/osc_ucx_comm.c b/ompi/mca/osc/ucx/osc_ucx_comm.c index 648307f24bb..b7d421180d2 100644 --- a/ompi/mca/osc/ucx/osc_ucx_comm.c +++ b/ompi/mca/osc/ucx/osc_ucx_comm.c @@ -654,6 +654,15 @@ int ompi_osc_ucx_win_get_notify_value(struct ompi_win_t *win, int notify, CHECK_NOTIFY_IDX(module, notify, my_rank); + /* Origins increment this counter with a UCX atomic, which the transport may + * emulate in software on the local worker rather than offload to the NIC. + * In that case the counter only advances while the worker is progressed, so + * a consumer spinning on MPI_WIN_GET_NOTIFY_VALUE -- the natural way to wait + * for a notification -- would never observe the update. Progress the worker + * here so that such a loop makes forward progress on its own, as every other + * spin-wait in this component does. */ + opal_common_ucx_wpool_progress(mca_osc_ucx_component.wpool); + volatile uint64_t *counter = (volatile uint64_t *)osc_ucx_notify_counter_addr(module, my_rank, notify); *value = (OMPI_MPI_COUNT_TYPE)*counter; From fe59a335fb759fbd6c89eabb45ab22aef40ad0e7 Mon Sep 17 00:00:00 2001 From: Joseph Antony Date: Thu, 13 Aug 2026 10:23:16 -0400 Subject: [PATCH 08/10] Fix MPI_WIN_SET_NUM_NOTIFY hang on an invalid argument 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 --- ompi/mca/osc/sm/osc_sm_comm.c | 26 ++++++++++++++++++++++++-- ompi/mpi/c/win_set_num_notify.c.in | 11 +++++++++-- 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/ompi/mca/osc/sm/osc_sm_comm.c b/ompi/mca/osc/sm/osc_sm_comm.c index f20c00c2517..f553aa528d2 100644 --- a/ompi/mca/osc/sm/osc_sm_comm.c +++ b/ompi/mca/osc/sm/osc_sm_comm.c @@ -223,7 +223,7 @@ ompi_osc_sm_win_set_num_notify(struct ompi_win_t *win, int rank = ompi_comm_rank(module->comm); unsigned long requested = (unsigned long) num_notifications; unsigned long *new_caps; - bool grow = false; + bool grow = false, bad; int ret, i; /* "mpi_assert_same_num_notifications" would let us skip the allgather below @@ -232,7 +232,18 @@ ompi_osc_sm_win_set_num_notify(struct ompi_win_t *win, * synchronizing and collective. */ (void) info; - if (num_notifications < 0) { + /* num_notifications is a local argument -- MPI-5.1 12.6.1 allows it to + * differ between MPI processes -- but this is a synchronizing collective. + * A rank that rejected its own value and returned here would leave every + * other rank blocked in the allgather below, turning an erroneous argument + * into a hang. So the validity rides through the collective as a sentinel + * and all ranks fail together. A multi-process window defers the decision; + * a single-process one has nobody to agree with and can answer now. */ + bad = (num_notifications < 0) + || (0 != module->notify_max_assert && + requested > (unsigned long) module->notify_max_assert); + + if (bad && 1 == comm_size) { return MPI_ERR_ARG; } @@ -267,6 +278,7 @@ ompi_osc_sm_win_set_num_notify(struct ompi_win_t *win, return OMPI_SUCCESS; } +agree: new_caps = malloc(sizeof(*new_caps) * comm_size); if (NULL == new_caps) { return OMPI_ERR_TEMP_OUT_OF_RESOURCE; @@ -281,6 +293,16 @@ ompi_osc_sm_win_set_num_notify(struct ompi_win_t *win, return ret; } + for (i = 0 ; i < comm_size ; ++i) { + if (ULONG_MAX == new_caps[i]) { + /* Some rank supplied an invalid count. Every rank sees the same + * gathered array, so they all report the same error and none of + * them reconfigures. */ + free(new_caps); + return MPI_ERR_ARG; + } + } + for (i = 0 ; i < comm_size ; ++i) { if (new_caps[i] > module->node_states[i].notify_counter_capacity) { grow = true; diff --git a/ompi/mpi/c/win_set_num_notify.c.in b/ompi/mpi/c/win_set_num_notify.c.in index cc1d39a9e77..f85bf1c8fd8 100644 --- a/ompi/mpi/c/win_set_num_notify.c.in +++ b/ompi/mpi/c/win_set_num_notify.c.in @@ -30,10 +30,17 @@ PROTOTYPE ERROR_CLASS win_set_num_notify(WIN win, INFO info, INT num_notificatio return OMPI_ERRHANDLER_NOHANDLE_INVOKE(MPI_ERR_WIN, FUNC_NAME); } else if (NULL != info && MPI_INFO_NULL != info && ompi_info_is_freed(info)) { rc = MPI_ERR_INFO; - } else if (num_notifications < 0) { - rc = MPI_ERR_ARG; } + /* num_notifications is deliberately *not* range-checked here. + MPI_WIN_SET_NUM_NOTIFY is a synchronizing collective whose count + argument is local -- MPI-5.1 12.6.1 allows it to differ between MPI + processes -- so a rank that rejected its own value here would return + while every other rank stayed blocked in the osc module's internal + collective, turning an erroneous argument into a hang. The module + carries each rank's validity through that collective instead, so all + ranks agree to fail together. */ + OMPI_ERRHANDLER_CHECK(rc, win, rc, FUNC_NAME); } From 6fa28addc5cac6bea1da1c116ef4f1d82f4dc39d Mon Sep 17 00:00:00 2001 From: Joseph Antony Date: Thu, 13 Aug 2026 10:24:11 -0400 Subject: [PATCH 09/10] osc/ucx: grow notification counters on demand 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 --- ompi/mca/osc/ucx/osc_ucx.h | 21 ++++- ompi/mca/osc/ucx/osc_ucx_comm.c | 63 ++++++++++++-- ompi/mca/osc/ucx/osc_ucx_component.c | 119 +++++++++++++++++++++++++-- 3 files changed, 183 insertions(+), 20 deletions(-) diff --git a/ompi/mca/osc/ucx/osc_ucx.h b/ompi/mca/osc/ucx/osc_ucx.h index c02b0984010..f79db5d1070 100644 --- a/ompi/mca/osc/ucx/osc_ucx.h +++ b/ompi/mca/osc/ucx/osc_ucx.h @@ -134,10 +134,17 @@ typedef struct ompi_osc_ucx_module { * rank (size comm_size), as set by MPI_WIN_SET_NUM_NOTIFY and * kept consistent across the group by an allgather. Always * <= notify_capacity. */ - unsigned int notify_capacity; /* notification counters reserved per rank at window - * creation. Agreed on across the group, and a hard - * upper bound: the registration cannot grow - * afterwards. */ + unsigned int notify_capacity; /* notification counters currently reserved per rank. + * Agreed on across the group and uniform. Grown on + * demand by MPI_WIN_SET_NUM_NOTIFY unless + * notify_max_assert caps it. */ + unsigned int notify_max_assert; /* non-zero only if *every* rank passed + * "mpi_assert_max_num_notify" at window creation. + * Then the agreed reservation is a hard cap and the + * counters never grow (MPI-5.1 12.2: the assertion + * lets the implementation optimize the allocation). + * Zero means no rank asserted a bound, so the + * standard's "does not assume any limit" applies. */ uint64_t *notify_addrs; /* per-rank base address of the notification counters * (size comm_size) */ void *notify_base; /* this rank's counters; notify_capacity uint64_t */ @@ -361,6 +368,12 @@ int ompi_osc_ucx_win_reset_notify_value(struct ompi_win_t *win, int notify, OMPI_MPI_COUNT_TYPE *value); int ompi_osc_ucx_win_set_num_notify(struct ompi_win_t *win, struct opal_info_t *info, int num_notifications); +/* Collectively re-reserve new_capacity notification counters per rank, replacing + * the current registration. Defined in osc_ucx_component.c because it needs the + * component's address-exchange helper. */ +int ompi_osc_ucx_grow_notify_counters(ompi_osc_ucx_module_t *module, + unsigned int new_capacity); + int ompi_osc_ucx_win_get_num_notify(struct ompi_win_t *win, int target_rank, int *num_notifications); diff --git a/ompi/mca/osc/ucx/osc_ucx_comm.c b/ompi/mca/osc/ucx/osc_ucx_comm.c index b7d421180d2..243810c8f5d 100644 --- a/ompi/mca/osc/ucx/osc_ucx_comm.c +++ b/ompi/mca/osc/ucx/osc_ucx_comm.c @@ -22,6 +22,10 @@ #include "osc_ucx.h" #include "osc_ucx_request.h" +#include + +#include "ompi/attribute/attribute.h" + #define CHECK_VALID_RKEY(_module, _target, _count) \ if (!((_module)->win_info_array[_target]).rkey_init && ((_count) > 0)) { \ @@ -725,17 +729,23 @@ int ompi_osc_ucx_win_set_num_notify(struct ompi_win_t *win, struct opal_info_t * int comm_size = ompi_comm_size(module->comm); int requested = num_notifications; int *requested_counts; + unsigned int needed; int ret, i; (void) info; /* "mpi_assert_same_num_notifications" is an optimization hint only */ - /* The counters are registered once at window creation and that registration - * cannot grow, so notify_capacity is a hard upper bound (reported as - * MPI_WIN_NOTIFICATION_NUM_UB). This is a synchronizing collective, so a + /* When every rank asserted "mpi_assert_max_num_notify" at window creation, + * that value is a hard upper bound and asking for more is erroneous. + * Otherwise MPI-5.1 12.2 says no limit is assumed, so a request above the + * current reservation grows the counters below rather than failing. + * + * This is a synchronizing collective, so a * rank with a bad argument must not return before the allgather below -- * that would leave the rest of the group blocked in it. Mark the request * instead and let every rank discover the error from the gathered values. */ - if (requested < 0 || (unsigned int) requested > module->notify_capacity) { + if (requested < 0 || + (0 != module->notify_max_assert && + (unsigned int) requested > module->notify_max_assert)) { requested = -1; } @@ -782,6 +792,39 @@ int ompi_osc_ucx_win_set_num_notify(struct ompi_win_t *win, struct opal_info_t * } } + /* Every rank sees the same gathered array, so they all reach the same + * decision about whether to grow and to what size, without extra + * communication. The reservation is uniform across the window (window + * creation agrees it with an allreduce), so it grows to the largest request + * anyone made. Never shrink: a rank that lowered its count keeps the space + * it already has, so only genuine growth costs a re-registration and + * alternating high/low requests do not thrash the NIC. */ + needed = module->notify_capacity; + for (i = 0; i < comm_size; i++) { + if ((unsigned int) requested_counts[i] > needed) { + needed = (unsigned int) requested_counts[i]; + } + } + + if (needed > module->notify_capacity) { + ret = ompi_osc_ucx_grow_notify_counters(module, needed); + if (OMPI_SUCCESS != ret) { + free(requested_counts); + return ret; + } + + /* MPI_WIN_NOTIFICATION_NUM_SB is the count supported without paying for + * a re-registration, so it has to follow the reservation rather than + * stay at whatever was cached when the window was created. */ + ret = ompi_attr_set_int(WIN_ATTR, win, &win->w_keyhash, + MPI_WIN_NOTIFICATION_NUM_SB, + (int) module->notify_capacity, true); + if (OMPI_SUCCESS != ret) { + free(requested_counts); + return ret; + } + } + memcpy(module->notify_counts, requested_counts, comm_size * sizeof(int)); free(requested_counts); @@ -793,11 +836,15 @@ int ompi_osc_ucx_win_get_notify_bounds(struct ompi_win_t *win, int *num_sb, int { ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t *)win->w_osc_module; - /* The counters have their own registration, which cannot grow after window - * creation, so the reservation is both the efficiently supported count and - * the hard upper bound. It does not depend on the window's flavor. */ + /* The current reservation is what is supported without paying for a + * re-registration, so it is the suggested bound. The hard bound is only + * real when every rank asserted "mpi_assert_max_num_notify" at window + * creation; otherwise MPI_WIN_SET_NUM_NOTIFY grows the counters on demand + * and the only limit is what can be allocated. Neither depends on the + * window's flavor. */ *num_sb = (int) module->notify_capacity; - *num_ub = (int) module->notify_capacity; + *num_ub = (0 != module->notify_max_assert) ? (int) module->notify_max_assert + : INT_MAX; /* Counters are uint64_t and only ever incremented by one per notified * operation, but they are returned to the user as a signed MPI_Count, so diff --git a/ompi/mca/osc/ucx/osc_ucx_component.c b/ompi/mca/osc/ucx/osc_ucx_component.c index 4776f287a1a..4c33b366543 100644 --- a/ompi/mca/osc/ucx/osc_ucx_component.c +++ b/ompi/mca/osc/ucx/osc_ucx_component.c @@ -169,16 +169,22 @@ static bool check_config_value_bool (char *key, opal_info_t *info) } /* Read the mpi_assert_max_num_notify info key (MPI-5.1 section 12.2) to decide - * how many notification counters to reserve per MPI process. The counters are - * registered as part of the window's memory region and that registration cannot - * grow later, so the reservation is a hard upper bound on - * MPI_WIN_SET_NUM_NOTIFY for the lifetime of the window. */ -static int osc_ucx_reserved_notify_counters(opal_info_t *info, unsigned int *reserved) + * how many notification counters to reserve per MPI process initially, and + * report in *asserted whether the key was actually given. + * + * The key is an assertion by the caller that it will not ask + * MPI_WIN_SET_NUM_NOTIFY for more than this, which lets us size the + * registration once and treat it as a hard upper bound. Without it the + * standard is explicit that "the implementation does not assume any limit", so + * the reservation is only a starting size and the counters grow on demand. */ +static int osc_ucx_reserved_notify_counters(opal_info_t *info, unsigned int *reserved, + bool *asserted) { opal_cstring_t *value_string; int flag = 0, value = 0; *reserved = mca_osc_ucx_component.num_notify_counters; + *asserted = false; if (NULL == info) { return OMPI_SUCCESS; @@ -203,6 +209,7 @@ static int osc_ucx_reserved_notify_counters(opal_info_t *info, unsigned int *res if (0 != value) { *reserved = (unsigned int) value; + *asserted = true; } return OMPI_SUCCESS; @@ -755,20 +762,25 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt * rest of the group blocked in window creation -- so the failure is carried * through the collective as a negative reservation instead. */ unsigned int notify_reserved = 0; - int notify_values[2]; + bool notify_asserted = false; + int notify_values[3]; /* The reservation is exchanged as an int and is used to size an allocation, * so a value that does not fit is a bad configuration rather than a request * to honor. Both failures ride the flag, not the value: MPI_MAX would hide * a sentinel value behind some other rank's larger reservation. */ - bool notify_bad = (OMPI_SUCCESS != osc_ucx_reserved_notify_counters(info, ¬ify_reserved)) + bool notify_bad = (OMPI_SUCCESS != osc_ucx_reserved_notify_counters(info, ¬ify_reserved, + ¬ify_asserted)) || notify_reserved > (unsigned int) INT_MAX; notify_values[0] = notify_bad ? 1 : 0; notify_values[1] = notify_bad ? 0 : (int) notify_reserved; + /* Carried as "some rank did NOT assert" so that it combines under MPI_MAX + * along with the other two values. */ + notify_values[2] = notify_asserted ? 0 : 1; /* info is allowed to differ between MPI processes, so agree on one * reservation for the whole window. Taking the maximum keeps every rank's * own assertion satisfiable, and propagates any rank's failure flag. */ - ret = module->comm->c_coll->coll_allreduce(MPI_IN_PLACE, notify_values, 2, + ret = module->comm->c_coll->coll_allreduce(MPI_IN_PLACE, notify_values, 3, MPI_INT, MPI_MAX, module->comm, module->comm->c_coll->coll_allreduce_module); if (OMPI_SUCCESS != ret) { @@ -779,6 +791,11 @@ static int component_select(struct ompi_win_t *win, void **base, size_t size, pt goto error; } module->notify_capacity = (unsigned int) notify_values[1]; + /* Only treat the reservation as a hard cap when *every* rank asserted a + * bound. A rank that gave no key made no promise, so the window has to + * stay growable for it -- MPI-5.1 12.2 says an absent (zero) key means the + * implementation assumes no limit. */ + module->notify_max_assert = notify_values[2] ? 0 : (unsigned int) notify_values[1]; module->no_locks = check_config_value_bool ("no_locks", info); module->acc_single_intrinsic = check_config_value_bool ("acc_single_intrinsic", info); @@ -1239,6 +1256,92 @@ int ompi_osc_ucx_dynamic_unlock(ompi_osc_ucx_module_t *module, int target) { return OMPI_SUCCESS; } +/* Collectively replace the notification-counter registration with a larger one. + * + * Called from MPI_WIN_SET_NUM_NOTIFY, which the standard defines as a blocking, + * synchronizing collective procedure -- that is what makes this safe. Every + * rank has to take part even if its own request fits, because registering the + * memory exchanges rkeys with the whole group. + * + * Two properties of MPI_WIN_SET_NUM_NOTIFY keep this simple. It resets every + * counter to zero, so a freshly calloc'd region is already the required + * contents and no value has to be carried across. And it is erroneous to call + * it while an access epoch is open or with an active notification-threshold + * request, so no remote atomic can be in flight against the old region while it + * is being replaced. + * + * The allgather of the new base addresses doubles as the barrier that lets the + * old region be released: once it completes, every rank has published its new + * address and no rank can issue a notified operation until it returns from the + * enclosing collective. */ +int ompi_osc_ucx_grow_notify_counters(ompi_osc_ucx_module_t *module, + unsigned int new_capacity) +{ + int comm_size = ompi_comm_size(module->comm); + opal_common_ucx_wpmem_t *new_mem = NULL; + void *new_base = NULL, *reg_base; + char *my_mem_addr = NULL; + uint64_t my_addr, *new_addrs = NULL; + int my_mem_addr_size = 0; + int ret; + + new_base = calloc(new_capacity, sizeof(uint64_t)); + if (NULL == new_base) { + return MPI_ERR_NO_MEM; + } + + new_addrs = calloc(comm_size, sizeof(uint64_t)); + if (NULL == new_addrs) { + free(new_base); + return MPI_ERR_NO_MEM; + } + + reg_base = new_base; + ret = opal_common_ucx_wpmem_create(module->ctx, ®_base, + new_capacity * sizeof(uint64_t), + OPAL_COMMON_UCX_MEM_MAP, + &exchange_len_info, + OPAL_COMMON_UCX_WPMEM_ADDR_EXCHANGE_FULL, + (void *)module->comm, + &my_mem_addr, &my_mem_addr_size, + &new_mem); + if (OMPI_SUCCESS != ret) { + free(new_addrs); + free(new_base); + return ret; + } + + if (0 != my_mem_addr_size) { + /* rkey object is already distributed among comm processes */ + ucp_rkey_buffer_release(my_mem_addr); + } + + my_addr = (uint64_t) new_base; + ret = module->comm->c_coll->coll_allgather(&my_addr, sizeof(uint64_t), MPI_BYTE, + new_addrs, sizeof(uint64_t), MPI_BYTE, + module->comm, + module->comm->c_coll->coll_allgather_module); + if (OMPI_SUCCESS != ret) { + opal_common_ucx_wpmem_free(new_mem); + free(new_addrs); + free(new_base); + return ret; + } + + if (NULL != module->notify_mem) { + opal_common_ucx_wpmem_free(module->notify_mem); + } + free(module->notify_base); + free(module->notify_addrs); + + module->notify_mem = new_mem; + module->notify_base = new_base; + module->notify_addrs = new_addrs; + module->notify_capacity = new_capacity; + + return OMPI_SUCCESS; +} + int ompi_osc_ucx_win_attach(struct ompi_win_t *win, void *base, size_t len) { ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t*) win->w_osc_module; int insert_index = -1, contain_index; From b140aa3d9e74655367a45f115e4c80564a969a19 Mon Sep 17 00:00:00 2001 From: Joseph Antony Date: Thu, 13 Aug 2026 10:24:46 -0400 Subject: [PATCH 10/10] osc/ucx: make notified request completion cover the notification 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 --- ompi/mca/osc/ucx/osc_ucx_comm.c | 219 +++++++++++++++++++++----------- 1 file changed, 145 insertions(+), 74 deletions(-) diff --git a/ompi/mca/osc/ucx/osc_ucx_comm.c b/ompi/mca/osc/ucx/osc_ucx_comm.c index 243810c8f5d..1f400f85fb9 100644 --- a/ompi/mca/osc/ucx/osc_ucx_comm.c +++ b/ompi/mca/osc/ucx/osc_ucx_comm.c @@ -609,6 +609,27 @@ int ompi_osc_ucx_get(void *origin_addr, size_t origin_count, } } +static int osc_ucx_request_over_flush(ompi_osc_ucx_module_t *module, + struct ompi_win_t *win, int target, + ucp_ep_h *ep, enum req_type req_type, + struct ompi_request_t **request); + +static int accumulate_req(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + struct ompi_op_t *op, struct ompi_win_t *win, + ompi_osc_ucx_accumulate_request_t *ucx_req); + +static int get_accumulate_req(const void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + void *result_addr, size_t result_count, + struct ompi_datatype_t *result_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, + struct ompi_op_t *op, struct ompi_win_t *win, + ompi_osc_ucx_accumulate_request_t *ucx_req); + /* Returns the remote address of notify counter[notify] for the given target. * The counters have their own registered region (module->notify_mem), separate * from the window data, so this is independent of the window's flavor and size. */ @@ -739,10 +760,10 @@ int ompi_osc_ucx_win_set_num_notify(struct ompi_win_t *win, struct opal_info_t * * Otherwise MPI-5.1 12.2 says no limit is assumed, so a request above the * current reservation grows the counters below rather than failing. * - * This is a synchronizing collective, so a - * rank with a bad argument must not return before the allgather below -- - * that would leave the rest of the group blocked in it. Mark the request - * instead and let every rank discover the error from the gathered values. */ + * This is a synchronizing collective, so a rank with a bad argument must not + * return before the allgather below -- that would leave the rest of the + * group blocked in it. Mark the request instead and let every rank discover + * the error from the gathered values. */ if (requested < 0 || (0 != module->notify_max_assert && (unsigned int) requested > module->notify_max_assert)) { @@ -931,9 +952,17 @@ int ompi_osc_ucx_rput_notify(const void *origin_addr, size_t origin_count, OSC_UCX_GET_DEFAULT_EP(ep, module, target); - ret = ompi_osc_ucx_rput(origin_addr, origin_count, origin_dt, - target, target_disp, target_count, target_dt, - win, request); + ret = check_sync_state(module, target, true); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* Issue the data movement and the notification first, and only then build + * the request over a flush -- that way the flush covers both, so completing + * the request means the counter update has been pushed out and not merely + * queued behind the origin's next MPI call. */ + ret = ompi_osc_ucx_put(origin_addr, origin_count, origin_dt, target, + target_disp, target_count, target_dt, win); if (OMPI_SUCCESS != ret) { return ret; } @@ -944,7 +973,12 @@ int ompi_osc_ucx_rput_notify(const void *origin_addr, size_t origin_count, return OMPI_ERROR; } - return osc_ucx_notify_target(module, target, notify, ep); + ret = osc_ucx_notify_target(module, target, notify, ep); + if (OMPI_SUCCESS != ret) { + return ret; + } + + return osc_ucx_request_over_flush(module, win, target, ep, RPUT_REQ, request); } int ompi_osc_ucx_rget_notify(void *origin_addr, size_t origin_count, @@ -962,20 +996,33 @@ int ompi_osc_ucx_rget_notify(void *origin_addr, size_t origin_count, OSC_UCX_GET_DEFAULT_EP(ep, module, target); - ret = ompi_osc_ucx_rget(origin_addr, origin_count, origin_dt, - target, target_disp, target_count, target_dt, - win, request); + ret = check_sync_state(module, target, true); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* As in rput_notify, the data movement and the notification both precede + * the request-bearing flush so that the request covers both. */ + ret = ompi_osc_ucx_get(origin_addr, origin_count, origin_dt, target, + target_disp, target_count, target_dt, win); if (OMPI_SUCCESS != ret) { return ret; } - /* Flush to ensure GET data is locally available before notifying target. */ + /* Blocking flush: the notification tells the target its window has been + * read, so the GET must have completed before the counter is incremented + * (MPI-5.1 12.6.4 requires that order). */ ret = opal_common_ucx_ctx_flush(module->ctx, OPAL_COMMON_UCX_SCOPE_EP, target); if (OPAL_SUCCESS != ret) { return OMPI_ERROR; } - return osc_ucx_notify_target(module, target, notify, ep); + ret = osc_ucx_notify_target(module, target, notify, ep); + if (OMPI_SUCCESS != ret) { + return ret; + } + + return osc_ucx_request_over_flush(module, win, target, ep, RGET_REQ, request); } int ompi_osc_ucx_accumulate_notify(const void *origin_addr, size_t origin_count, @@ -1062,9 +1109,20 @@ int ompi_osc_ucx_raccumulate_notify(const void *origin_addr, size_t origin_count OSC_UCX_GET_DEFAULT_EP(ep, module, target); - ret = ompi_osc_ucx_raccumulate(origin_addr, origin_count, origin_dt, - target, target_disp, target_count, target_dt, - op, win, request); + ret = check_sync_state(module, target, true); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* Passing a NULL request runs accumulate_req as the blocking accumulate: it + * ends with a blocking flush, so the read-modify-write has completed at the + * target before we return here. The notification issued below is therefore + * correctly ordered after the window access (MPI-5.1 12.6.4). Building our + * own request afterwards -- rather than taking the pre-completed one from + * ompi_osc_ucx_raccumulate -- is what lets that request also cover the + * notification. */ + ret = accumulate_req(origin_addr, origin_count, origin_dt, target, target_disp, + target_count, target_dt, op, win, NULL); if (OMPI_SUCCESS != ret) { return ret; } @@ -1075,7 +1133,15 @@ int ompi_osc_ucx_raccumulate_notify(const void *origin_addr, size_t origin_count return OMPI_ERROR; } - return osc_ucx_notify_target(module, target, notify, ep); + ret = osc_ucx_notify_target(module, target, notify, ep); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* The accumulate is already complete, so the request only has to represent + * the notification reaching the wire; RPUT_REQ selects the plain + * flush-completion behaviour rather than the accumulate state machine. */ + return osc_ucx_request_over_flush(module, win, target, ep, RPUT_REQ, request); } int ompi_osc_ucx_rget_accumulate_notify(const void *origin_addr, size_t origin_count, @@ -1096,22 +1162,34 @@ int ompi_osc_ucx_rget_accumulate_notify(const void *origin_addr, size_t origin_c OSC_UCX_GET_DEFAULT_EP(ep, module, target); - ret = ompi_osc_ucx_rget_accumulate(origin_addr, origin_count, origin_dt, - result_addr, result_count, result_dt, - target, target_disp, target_count, target_dt, - op, win, request); + ret = check_sync_state(module, target, true); if (OMPI_SUCCESS != ret) { return ret; } - /* Flush so the fetched result is locally valid and the update has landed at - * the target before it can observe the notification. */ - ret = opal_common_ucx_ctx_flush(module->ctx, OPAL_COMMON_UCX_SCOPE_EP, target); + /* As in raccumulate_notify: a NULL request makes this the blocking form, + * which ends with a blocking flush, so both the fetched result and the + * update at the target are complete before the notification is issued. */ + ret = get_accumulate_req(origin_addr, origin_count, origin_dt, + result_addr, result_count, result_dt, + target, target_disp, target_count, target_dt, + op, win, NULL); + if (OMPI_SUCCESS != ret) { + return ret; + } + + /* Fence to order the accumulate ahead of the counter increment. */ + ret = opal_common_ucx_wpmem_fence(module->mem); if (OPAL_SUCCESS != ret) { return OMPI_ERROR; } - return osc_ucx_notify_target(module, target, notify, ep); + ret = osc_ucx_notify_target(module, target, notify, ep); + if (OMPI_SUCCESS != ret) { + return ret; + } + + return osc_ucx_request_over_flush(module, win, target, ep, RGET_REQ, request); } static inline bool ompi_osc_need_acc_lock(ompi_osc_ucx_module_t *module, int target) @@ -2001,31 +2079,30 @@ int ompi_osc_ucx_get_accumulate_nb(const void *origin_addr, size_t origin_count, target_count, target_dt, op, win, GET_ACCUMULATE); } -int ompi_osc_ucx_rput(const void *origin_addr, size_t origin_count, - struct ompi_datatype_t *origin_dt, - int target, ptrdiff_t target_disp, size_t target_count, - struct ompi_datatype_t *target_dt, - struct ompi_win_t *win, struct ompi_request_t **request) { - ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t*) win->w_osc_module; - ucp_ep_h *ep; - OSC_UCX_GET_DEFAULT_EP(ep, module, target); +/* Attach an MPI request to a nonblocking worker flush. + * + * The request completes when ucp_worker_flush_nb completes, and that flush + * covers every operation already issued on this worker -- whichever memory + * registration they used. So the request waits for exactly what the caller + * issued *before* getting here. The notified variants rely on this: they issue + * their counter atomic first, so completing the request implies the + * notification has been pushed to the target rather than left queued locally + * (MPI-5.1 12.6.4, advice to implementors). + * + * Allocating the request last also keeps the failure contract clean: anything + * that can fail has already run, so callers return their errors with *request + * untouched, as MPI expects of a call that reports an error. */ +static int osc_ucx_request_over_flush(ompi_osc_ucx_module_t *module, + struct ompi_win_t *win, int target, + ucp_ep_h *ep, enum req_type req_type, + struct ompi_request_t **request) +{ opal_common_ucx_wpmem_t *mem = module->mem; uint64_t remote_addr = (module->state_addrs[target]) + OSC_UCX_STATE_REQ_FLAG_OFFSET; ompi_osc_ucx_generic_request_t *ucx_req = NULL; int ret = OMPI_SUCCESS; - ret = check_sync_state(module, target, true); - if (ret != OMPI_SUCCESS) { - return ret; - } - - ret = ompi_osc_ucx_put(origin_addr, origin_count, origin_dt, target, target_disp, - target_count, target_dt, win); - if (ret != OMPI_SUCCESS) { - return ret; - } - - OMPI_OSC_UCX_GENERIC_REQUEST_ALLOC(win, ucx_req, RPUT_REQ); + OMPI_OSC_UCX_GENERIC_REQUEST_ALLOC(win, ucx_req, req_type); ucx_req->super.module = module; OSC_UCX_INCREMENT_OUTSTANDING_NB_OPS(module); @@ -2055,17 +2132,14 @@ int ompi_osc_ucx_rput(const void *origin_addr, size_t origin_count, return ret; } -int ompi_osc_ucx_rget(void *origin_addr, size_t origin_count, +int ompi_osc_ucx_rput(const void *origin_addr, size_t origin_count, struct ompi_datatype_t *origin_dt, int target, ptrdiff_t target_disp, size_t target_count, - struct ompi_datatype_t *target_dt, struct ompi_win_t *win, - struct ompi_request_t **request) { + struct ompi_datatype_t *target_dt, + struct ompi_win_t *win, struct ompi_request_t **request) { ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t*) win->w_osc_module; ucp_ep_h *ep; OSC_UCX_GET_DEFAULT_EP(ep, module, target); - opal_common_ucx_wpmem_t *mem = module->mem; - uint64_t remote_addr = (module->state_addrs[target]) + OSC_UCX_STATE_REQ_FLAG_OFFSET; - ompi_osc_ucx_generic_request_t *ucx_req = NULL; int ret = OMPI_SUCCESS; ret = check_sync_state(module, target, true); @@ -2073,40 +2147,37 @@ int ompi_osc_ucx_rget(void *origin_addr, size_t origin_count, return ret; } - ret = ompi_osc_ucx_get(origin_addr, origin_count, origin_dt, target, target_disp, + ret = ompi_osc_ucx_put(origin_addr, origin_count, origin_dt, target, target_disp, target_count, target_dt, win); if (ret != OMPI_SUCCESS) { return ret; } - OMPI_OSC_UCX_GENERIC_REQUEST_ALLOC(win, ucx_req, RGET_REQ); - ucx_req->super.module = module; + return osc_ucx_request_over_flush(module, win, target, ep, RPUT_REQ, request); +} - OSC_UCX_INCREMENT_OUTSTANDING_NB_OPS(module); - ret = opal_common_ucx_wpmem_flush_ep_nb(mem, target, ompi_osc_ucx_req_completion, ucx_req, ep); +int ompi_osc_ucx_rget(void *origin_addr, size_t origin_count, + struct ompi_datatype_t *origin_dt, + int target, ptrdiff_t target_disp, size_t target_count, + struct ompi_datatype_t *target_dt, struct ompi_win_t *win, + struct ompi_request_t **request) { + ompi_osc_ucx_module_t *module = (ompi_osc_ucx_module_t*) win->w_osc_module; + ucp_ep_h *ep; + OSC_UCX_GET_DEFAULT_EP(ep, module, target); + int ret = OMPI_SUCCESS; + ret = check_sync_state(module, target, true); if (ret != OMPI_SUCCESS) { - /* fallback to using an atomic op to acquire a request handle */ - ret = opal_common_ucx_wpmem_fence(mem); - if (ret != OMPI_SUCCESS) { - OSC_UCX_VERBOSE(1, "opal_common_ucx_mem_fence failed: %d", ret); - OMPI_OSC_UCX_REQUEST_RETURN(ucx_req); - return OMPI_ERROR; - } - - ret = opal_common_ucx_wpmem_fetch_nb(mem, UCP_ATOMIC_FETCH_OP_FADD, - 0, target, &(module->req_result), - sizeof(uint64_t), remote_addr & (~0x7), - ompi_osc_ucx_req_completion, ucx_req, ep); - if (ret != OMPI_SUCCESS) { - OMPI_OSC_UCX_REQUEST_RETURN(ucx_req); - return ret; - } + return ret; } - *request = &ucx_req->super.super; + ret = ompi_osc_ucx_get(origin_addr, origin_count, origin_dt, target, target_disp, + target_count, target_dt, win); + if (ret != OMPI_SUCCESS) { + return ret; + } - return ret; + return osc_ucx_request_over_flush(module, win, target, ep, RGET_REQ, request); } int ompi_osc_ucx_raccumulate(const void *origin_addr, size_t origin_count,