diff --git a/config/ompi_check_ucx.m4 b/config/ompi_check_ucx.m4 index 42e53f9ce8..7f04ba3a52 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 ffb7d61834..fb7b7f8461 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/oshmem/mca/spml/ucx/spml_ucx_component.c b/oshmem/mca/spml/ucx/spml_ucx_component.c index 0f0ce2a15a..3d29bd4e5d 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) {