diff --git a/config/ompi_check_ucx.m4 b/config/ompi_check_ucx.m4 index 044b599dc3..b783fd1662 100644 --- a/config/ompi_check_ucx.m4 +++ b/config/ompi_check_ucx.m4 @@ -120,7 +120,8 @@ AC_DEFUN([OMPI_CHECK_UCX],[ UCP_ATOMIC_POST_OP_XOR, UCP_ATOMIC_FETCH_OP_FAND, UCP_ATOMIC_FETCH_OP_FOR, - UCP_ATOMIC_FETCH_OP_FXOR], + UCP_ATOMIC_FETCH_OP_FXOR, + UCP_PARAM_FIELD_ESTIMATED_NUM_PPN], [], [], [#include ]) AC_CHECK_DECLS([UCP_WORKER_ATTR_FIELD_ADDRESS_FLAGS], diff --git a/ompi/mca/pml/ucx/pml_ucx.c b/ompi/mca/pml/ucx/pml_ucx.c index db6a9fa56d..cd5aa32810 100644 --- a/ompi/mca/pml/ucx/pml_ucx.c +++ b/ompi/mca/pml/ucx/pml_ucx.c @@ -197,22 +197,27 @@ int mca_pml_ucx_open(void) } /* Initialize UCX context */ - params.field_mask = UCP_PARAM_FIELD_FEATURES | - UCP_PARAM_FIELD_REQUEST_SIZE | - UCP_PARAM_FIELD_REQUEST_INIT | - UCP_PARAM_FIELD_REQUEST_CLEANUP | - UCP_PARAM_FIELD_TAG_SENDER_MASK | - UCP_PARAM_FIELD_MT_WORKERS_SHARED | - UCP_PARAM_FIELD_ESTIMATED_NUM_EPS; - params.features = UCP_FEATURE_TAG; - params.request_size = sizeof(ompi_request_t); - params.request_init = mca_pml_ucx_request_init; - params.request_cleanup = mca_pml_ucx_request_cleanup; - params.tag_sender_mask = PML_UCX_SPECIFIC_SOURCE_MASK; + params.field_mask = UCP_PARAM_FIELD_FEATURES | + UCP_PARAM_FIELD_REQUEST_SIZE | + UCP_PARAM_FIELD_REQUEST_INIT | + UCP_PARAM_FIELD_REQUEST_CLEANUP | + UCP_PARAM_FIELD_TAG_SENDER_MASK | + UCP_PARAM_FIELD_MT_WORKERS_SHARED | + UCP_PARAM_FIELD_ESTIMATED_NUM_EPS; + params.features = UCP_FEATURE_TAG; + params.request_size = sizeof(ompi_request_t); + params.request_init = mca_pml_ucx_request_init; + params.request_cleanup = mca_pml_ucx_request_cleanup; + params.tag_sender_mask = PML_UCX_SPECIFIC_SOURCE_MASK; params.mt_workers_shared = 0; /* we do not need mt support for context since it will be protected by worker */ params.estimated_num_eps = ompi_proc_world_size(); +#if HAVE_DECL_UCP_PARAM_FIELD_ESTIMATED_NUM_PPN + params.estimated_num_ppn = opal_process_info.num_local_peers + 1; + params.field_mask |= UCP_PARAM_FIELD_ESTIMATED_NUM_PPN; +#endif + status = ucp_init(¶ms, config, &ompi_pml_ucx.ucp_context); ucp_config_release(config); diff --git a/opal/mca/common/ucx/common_ucx_wpool.c b/opal/mca/common/ucx/common_ucx_wpool.c index 5f85fd5691..0ef56c44a2 100644 --- a/opal/mca/common/ucx/common_ucx_wpool.c +++ b/opal/mca/common/ucx/common_ucx_wpool.c @@ -168,6 +168,11 @@ opal_common_ucx_wpool_init(opal_common_ucx_wpool_t *wpool, context_params.request_init = opal_common_ucx_req_init; context_params.request_size = sizeof(opal_common_ucx_request_t); +#if HAVE_DECL_UCP_PARAM_FIELD_ESTIMATED_NUM_PPN + context_params.estimated_num_ppn = opal_process_info.num_local_peers + 1; + context_params.field_mask |= UCP_PARAM_FIELD_ESTIMATED_NUM_PPN; +#endif + status = ucp_init(&context_params, config, &wpool->ucp_ctx); ucp_config_release(config); if (UCS_OK != status) { diff --git a/oshmem/mca/spml/ucx/spml_ucx_component.c b/oshmem/mca/spml/ucx/spml_ucx_component.c index 04cc3da87c..0bfdc1d61e 100644 --- a/oshmem/mca/spml/ucx/spml_ucx_component.c +++ b/oshmem/mca/spml/ucx/spml_ucx_component.c @@ -212,8 +212,12 @@ static int spml_ucx_init(void) opal_common_ucx_mca_register(); memset(¶ms, 0, sizeof(params)); - params.field_mask = UCP_PARAM_FIELD_FEATURES|UCP_PARAM_FIELD_ESTIMATED_NUM_EPS|UCP_PARAM_FIELD_MT_WORKERS_SHARED; - params.features = UCP_FEATURE_RMA|UCP_FEATURE_AMO32|UCP_FEATURE_AMO64; + params.field_mask = UCP_PARAM_FIELD_FEATURES | + UCP_PARAM_FIELD_ESTIMATED_NUM_EPS | + UCP_PARAM_FIELD_MT_WORKERS_SHARED; + params.features = UCP_FEATURE_RMA | + UCP_FEATURE_AMO32 | + UCP_FEATURE_AMO64; params.estimated_num_eps = ompi_proc_world_size(); if (oshmem_mpi_thread_requested == SHMEM_THREAD_MULTIPLE) { params.mt_workers_shared = 1; @@ -221,6 +225,11 @@ static int spml_ucx_init(void) params.mt_workers_shared = 0; } +#if HAVE_DECL_UCP_PARAM_FIELD_ESTIMATED_NUM_PPN + params.estimated_num_ppn = opal_process_info.num_local_peers + 1; + params.field_mask |= UCP_PARAM_FIELD_ESTIMATED_NUM_PPN; +#endif + err = ucp_init(¶ms, ucp_config, &mca_spml_ucx.ucp_context); ucp_config_release(ucp_config); if (UCS_OK != err) {