Skip to content

Commit eee891f

Browse files
devrealdevreal
committed
osc/ucx: implement rput and rget using ucp_worker_flush_nb
Fallback to the old method of acquiring a request from an atomic operation is preserved. Some networks might provide better performance if the request-based operations do not rely on atomic operations. Some minor fixes to opal_common_ucx_winfo_flush included in this commit. Signed-off-by: Joseph Schuchart <[email protected]>
1 parent 59928bb commit eee891f

File tree

3 files changed

+85
-23
lines changed

3 files changed

+85
-23
lines changed

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

Lines changed: 36 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1115,21 +1115,28 @@ int ompi_osc_ucx_rput(const void *origin_addr, int origin_count,
11151115
return ret;
11161116
}
11171117

1118-
ret = opal_common_ucx_wpmem_fence(mem);
1119-
if (ret != OMPI_SUCCESS) {
1120-
OSC_UCX_VERBOSE(1, "opal_common_ucx_mem_fence failed: %d", ret);
1121-
return OMPI_ERROR;
1122-
}
1123-
11241118
mca_osc_ucx_component.num_incomplete_req_ops++;
1125-
/* TODO: investigate whether ucp_worker_flush_nb is a better choice here */
1126-
ret = opal_common_ucx_wpmem_fetch_nb(module->state_mem, UCP_ATOMIC_FETCH_OP_FADD,
1127-
0, target, &(module->req_result),
1128-
sizeof(uint64_t), remote_addr & (~0x7),
1129-
req_completion, ucx_req);
1119+
ret = opal_common_ucx_wpmem_flush_ep_nb(mem, target, req_completion, ucx_req);
1120+
11301121
if (ret != OMPI_SUCCESS) {
11311122
OMPI_OSC_UCX_REQUEST_RETURN(ucx_req);
11321123
return ret;
1124+
1125+
/* fallback to using an atomic op to acquire a request handle */
1126+
ret = opal_common_ucx_wpmem_fence(mem);
1127+
if (ret != OMPI_SUCCESS) {
1128+
OSC_UCX_VERBOSE(1, "opal_common_ucx_mem_fence failed: %d", ret);
1129+
return OMPI_ERROR;
1130+
}
1131+
1132+
ret = opal_common_ucx_wpmem_fetch_nb(mem, UCP_ATOMIC_FETCH_OP_FADD,
1133+
0, target, &(module->req_result),
1134+
sizeof(uint64_t), remote_addr & (~0x7),
1135+
req_completion, ucx_req);
1136+
if (ret != OMPI_SUCCESS) {
1137+
OMPI_OSC_UCX_REQUEST_RETURN(ucx_req);
1138+
return ret;
1139+
}
11331140
}
11341141

11351142
*request = &ucx_req->super;
@@ -1170,21 +1177,28 @@ int ompi_osc_ucx_rget(void *origin_addr, int origin_count,
11701177
return ret;
11711178
}
11721179

1173-
ret = opal_common_ucx_wpmem_fence(mem);
1174-
if (ret != OMPI_SUCCESS) {
1175-
OSC_UCX_VERBOSE(1, "opal_common_ucx_mem_fence failed: %d", ret);
1176-
return OMPI_ERROR;
1177-
}
1178-
11791180
mca_osc_ucx_component.num_incomplete_req_ops++;
1180-
/* TODO: investigate whether ucp_worker_flush_nb is a better choice here */
1181-
ret = opal_common_ucx_wpmem_fetch_nb(module->state_mem, UCP_ATOMIC_FETCH_OP_FADD,
1182-
0, target, &(module->req_result),
1183-
sizeof(uint64_t), remote_addr & (~0x7),
1184-
req_completion, ucx_req);
1181+
ret = opal_common_ucx_wpmem_flush_ep_nb(mem, target, req_completion, ucx_req);
1182+
11851183
if (ret != OMPI_SUCCESS) {
11861184
OMPI_OSC_UCX_REQUEST_RETURN(ucx_req);
11871185
return ret;
1186+
1187+
/* fallback to using an atomic op to acquire a request handle */
1188+
ret = opal_common_ucx_wpmem_fence(mem);
1189+
if (ret != OMPI_SUCCESS) {
1190+
OSC_UCX_VERBOSE(1, "opal_common_ucx_mem_fence failed: %d", ret);
1191+
return OMPI_ERROR;
1192+
}
1193+
1194+
ret = opal_common_ucx_wpmem_fetch_nb(mem, UCP_ATOMIC_FETCH_OP_FADD,
1195+
0, target, &(module->req_result),
1196+
sizeof(uint64_t), remote_addr & (~0x7),
1197+
req_completion, ucx_req);
1198+
if (ret != OMPI_SUCCESS) {
1199+
OMPI_OSC_UCX_REQUEST_RETURN(ucx_req);
1200+
return ret;
1201+
}
11881202
}
11891203

11901204
*request = &ucx_req->super;

‎opal/mca/common/ucx/common_ucx_wpool.c

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -761,7 +761,7 @@ OPAL_DECLSPEC int opal_common_ucx_winfo_flush(opal_common_ucx_winfo_t *winfo, in
761761
((opal_common_ucx_request_t *) req)->winfo = winfo;
762762
}
763763

764-
if (OPAL_COMMON_UCX_FLUSH_B) {
764+
if (OPAL_COMMON_UCX_FLUSH_B == type) {
765765
rc = opal_common_ucx_wait_request_mt(req, "ucp_ep_flush_nb");
766766
} else {
767767
*req_ptr = req;
@@ -820,13 +820,57 @@ OPAL_DECLSPEC int opal_common_ucx_wpmem_flush(opal_common_ucx_wpmem_t *mem,
820820
if (rc != OPAL_SUCCESS) {
821821
MCA_COMMON_UCX_ERROR("opal_common_ucx_flush failed: %d", rc);
822822
rc = OPAL_ERROR;
823+
break;
823824
}
824825
}
825826
opal_mutex_unlock(&ctx->mutex);
826827

827828
return rc;
828829
}
829830

831+
832+
OPAL_DECLSPEC int opal_common_ucx_wpmem_flush_ep_nb(opal_common_ucx_wpmem_t *mem,
833+
int target,
834+
opal_common_ucx_user_req_handler_t user_req_cb,
835+
void *user_req_ptr)
836+
{
837+
#if HAVE_DECL_UCP_EP_FLUSH_NB
838+
int rc = OPAL_SUCCESS;
839+
ucp_ep_h ep = NULL;
840+
ucp_rkey_h rkey = NULL;
841+
opal_common_ucx_winfo_t *winfo = NULL;
842+
843+
if (NULL == mem) {
844+
return OPAL_SUCCESS;
845+
}
846+
847+
rc = opal_common_ucx_tlocal_fetch(mem, target, &ep, &rkey, &winfo);
848+
if (OPAL_UNLIKELY(OPAL_SUCCESS != rc)) {
849+
MCA_COMMON_UCX_ERROR("tlocal_fetch failed: %d", rc);
850+
return rc;
851+
}
852+
853+
opal_mutex_lock(&winfo->mutex);
854+
opal_common_ucx_request_t *req;
855+
req = ucp_worker_flush_nb(winfo->worker, 0, opal_common_ucx_req_completion);
856+
if (UCS_PTR_IS_PTR(req)) {
857+
req->ext_req = user_req_ptr;
858+
req->ext_cb = user_req_cb;
859+
req->winfo = winfo;
860+
} else {
861+
if (user_req_cb != NULL) {
862+
(*user_req_cb)(user_req_ptr);
863+
}
864+
}
865+
opal_mutex_unlock(&winfo->mutex);
866+
return rc;
867+
#else
868+
return OPAL_ERR_NOT_SUPPORTED;
869+
#endif // HAVE_DECL_UCP_EP_FLUSH_NB
870+
871+
}
872+
873+
830874
OPAL_DECLSPEC int opal_common_ucx_wpmem_fence(opal_common_ucx_wpmem_t *mem)
831875
{
832876
ucs_status_t status = UCS_OK;

‎opal/mca/common/ucx/common_ucx_wpool.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,10 @@ OPAL_DECLSPEC void opal_common_ucx_wpmem_free(opal_common_ucx_wpmem_t *mem);
247247

248248
OPAL_DECLSPEC int opal_common_ucx_wpmem_flush(opal_common_ucx_wpmem_t *mem,
249249
opal_common_ucx_flush_scope_t scope, int target);
250+
OPAL_DECLSPEC int opal_common_ucx_wpmem_flush_ep_nb(opal_common_ucx_wpmem_t *mem,
251+
int target,
252+
opal_common_ucx_user_req_handler_t user_req_cb,
253+
void *user_req_ptr);
250254
OPAL_DECLSPEC int opal_common_ucx_wpmem_fence(opal_common_ucx_wpmem_t *mem);
251255

252256
OPAL_DECLSPEC int opal_common_ucx_winfo_flush(opal_common_ucx_winfo_t *winfo, int target,

0 commit comments

Comments
 (0)