diff --git a/config/ompi_check_ucx.m4 b/config/ompi_check_ucx.m4 index a8a3f597cc..4e3dfcad90 100644 --- a/config/ompi_check_ucx.m4 +++ b/config/ompi_check_ucx.m4 @@ -45,7 +45,7 @@ AC_DEFUN([OMPI_CHECK_UCX],[ [OPAL_CHECK_PACKAGE([ompi_check_ucx], [ucp/api/ucp.h], [ucp], - [ucp_ep_flush_nb], + [ucp_cleanup], [], [], [], @@ -77,7 +77,7 @@ AC_DEFUN([OMPI_CHECK_UCX],[ OPAL_CHECK_PACKAGE([ompi_check_ucx], [ucp/api/ucp.h], [ucp], - [ucp_ep_flush_nb], + [ucp_cleanup], [], [$ompi_check_ucx_dir], [$ompi_check_ucx_libdir], @@ -108,6 +108,18 @@ AC_DEFUN([OMPI_CHECK_UCX],[ [AC_DEFINE([HAVE_UCP_TAG_SEND_NBR],[1], [have ucp_tag_send_nbr()])], [], [#include ]) + AC_CHECK_DECLS([ucp_ep_flush_nb], + [AC_DEFINE([HAVE_UCP_EP_FLUSH_NB],[1], + [have ucp_ep_flush_nb()])], [], + [#include ]) + AC_CHECK_DECLS([ucp_worker_flush_nb], + [AC_DEFINE([HAVE_UCP_WORKER_FLUSH_NB],[1], + [have ucp_worker_flush_nb()])], [], + [#include ]) + AC_CHECK_DECLS([ucp_request_check_status], + [AC_DEFINE([HAVE_UCP_REQUEST_CHECK_STATUS],[1], + [have ucp_request_check_status()])], [], + [#include ]) CPPFLAGS=$old_CPPFLAGS OPAL_SUMMARY_ADD([[Transports]],[[Open UCX]],[$1],[$ompi_check_ucx_happy])])]) diff --git a/opal/mca/common/ucx/common_ucx.h b/opal/mca/common/ucx/common_ucx.h index 3a2d208dc8..51f42a006c 100644 --- a/opal/mca/common/ucx/common_ucx.h +++ b/opal/mca/common/ucx/common_ucx.h @@ -33,6 +33,9 @@ ucs_status_t opal_common_ucx_wait_request(ucs_status_ptr_t request, ucp_worker_h { ucs_status_t status; int i; +#ifndef HAVE_UCP_REQUEST_CHECK_STATUS + ucp_tag_recv_info_t info; +#endif /* check for request completed or failed */ if (OPAL_LIKELY(UCS_OK == request)) { @@ -44,7 +47,13 @@ ucs_status_t opal_common_ucx_wait_request(ucs_status_ptr_t request, ucp_worker_h while (1) { /* call UCX progress */ for (i = 0; i < opal_common_ucx_progress_iterations; i++) { - if (UCS_INPROGRESS != (status = ucp_request_check_status(request))) { + if (UCS_INPROGRESS != (status = +#if HAVE_UCP_REQUEST_CHECK_STATUS + ucp_request_check_status(request) +#else + ucp_request_test(request, &info) +#endif + )) { ucp_request_free(request); return status; } @@ -59,19 +68,27 @@ ucs_status_t opal_common_ucx_wait_request(ucs_status_ptr_t request, ucp_worker_h static inline ucs_status_t opal_common_ucx_ep_flush(ucp_ep_h ep, ucp_worker_h worker) { +#if HAVE_UCP_EP_FLUSH_NB ucs_status_ptr_t status; status = ucp_ep_flush_nb(ep, 0, opal_common_ucx_empty_complete_cb); return opal_common_ucx_wait_request(status, worker); +#else + return ucp_ep_flush(ep); +#endif } static inline ucs_status_t opal_common_ucx_worker_flush(ucp_worker_h worker) { +#if HAVE_UCP_WORKER_FLUSH_NB ucs_status_ptr_t status; status = ucp_worker_flush_nb(worker, 0, opal_common_ucx_empty_complete_cb); return opal_common_ucx_wait_request(status, worker); +#else + return ucp_worker_flush(worker); +#endif } static inline