Skip to content

Commit 4538fe9

Browse files
ATOMIC/UCX: make atomic post blocking
According to API definition ucp_atomic_op_nbx() may return request handle. In this case progress till request moves to OK state and release it. Signed-off-by: Artemy Kovalyov <artemyko@mellanox.com>
1 parent 8e51a72 commit 4538fe9

File tree

1 file changed

+9
-7
lines changed

1 file changed

+9
-7
lines changed

oshmem/mca/atomic/ucx/atomic_ucx_module.c

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -55,13 +55,15 @@ int mca_atomic_ucx_op(shmem_ctx_t ctx,
5555
ucp_atomic_post_op_t op)
5656
#endif
5757
{
58-
ucs_status_t status;
5958
spml_ucx_mkey_t *ucx_mkey;
6059
uint64_t rva;
6160
mca_spml_ucx_ctx_t *ucx_ctx = (mca_spml_ucx_ctx_t *)ctx;
6261
#if HAVE_DECL_UCP_ATOMIC_OP_NBX
6362
ucs_status_ptr_t status_ptr;
63+
#else
64+
ucs_status_t status;
6465
#endif
66+
int res;
6567

6668
assert((8 == size) || (4 == size));
6769

@@ -71,20 +73,20 @@ int mca_atomic_ucx_op(shmem_ctx_t ctx,
7173
status_ptr = ucp_atomic_op_nbx(ucx_ctx->ucp_peers[pe].ucp_conn,
7274
op, &value, 1, rva, ucx_mkey->rkey,
7375
&mca_spml_ucp_request_params[size >> 3]);
74-
if (OPAL_LIKELY(!UCS_PTR_IS_ERR(status_ptr))) {
75-
mca_spml_ucx_remote_op_posted(ucx_ctx, pe);
76-
}
77-
status = UCS_PTR_STATUS(status_ptr);
76+
res = opal_common_ucx_wait_request(status_ptr, ucx_ctx->ucp_worker[0],
77+
"ucp_atomic_op_nbx post");
7878
#else
7979
status = ucp_atomic_post(ucx_ctx->ucp_peers[pe].ucp_conn,
8080
op, value, size, rva,
8181
ucx_mkey->rkey);
82+
res = ucx_status_to_oshmem(status);
8283
#endif
83-
if (OPAL_LIKELY(UCS_OK == status)) {
84+
85+
if (OPAL_LIKELY(OSHMEM_SUCCESS == res)) {
8486
mca_spml_ucx_remote_op_posted(ucx_ctx, pe);
8587
}
8688

87-
return ucx_status_to_oshmem(status);
89+
return res;
8890
}
8991

9092
static inline

0 commit comments

Comments
 (0)