Skip to content

UCX Notified Communication #20

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
joe-explr wants to merge 10 commits into
base: notified-rma
Choose a base branch
from
Open

UCX Notified Communication #20

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
2c880b8
UCX Notified Communication
joe-explr Mar 31, 2026
37ecf0a
UCX Notified Communication changes
joe-explr Jul 21, 2026
21395bc
Get and set Notify changes
joe-explr Jul 24, 2026
d944d8a
osc/ucx: use MPI_ERR_RMA_NOTIFICATION for invalid notification index
joe-explr Aug 12, 2026
26bcd0d
osc/ucx: complete notified communication except thresholds
joe-explr Aug 12, 2026
caaa0bc
osc/ucx: give notification counters their own registration
joe-explr Aug 12, 2026
74e978a
osc/ucx: progress the worker in MPI_WIN_GET_NOTIFY_VALUE
joe-explr Aug 12, 2026
fe59a33
Fix MPI_WIN_SET_NUM_NOTIFY hang on an invalid argument
joe-explr Aug 13, 2026
6fa28ad
osc/ucx: grow notification counters on demand
joe-explr Aug 13, 2026
b140aa3
osc/ucx: make notified request completion cover the notification
joe-explr Aug 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 24 additions & 2 deletions ompi/mca/osc/sm/osc_sm_comm.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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;
}

Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down
100 changes: 100 additions & 0 deletions ompi/mca/osc/ucx/osc_ucx.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +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
/* 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 {
Expand All @@ -43,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;
Expand Down Expand Up @@ -122,6 +130,24 @@ 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), 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 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 */
size_t *sizes; /* used if not every process has the same size */
uint64_t *addrs;
uint64_t *state_addrs;
Expand Down Expand Up @@ -149,6 +175,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;
Expand Down Expand Up @@ -277,6 +308,75 @@ 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_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,
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);

/* returns the size at the peer */
static inline size_t ompi_osc_ucx_get_size(ompi_osc_ucx_module_t *module, int rank)
{
Expand Down
Loading
Loading