Merge pull request #6864 from hoopoepg/topic/ucx-ppn-hint

UCX: added PPN hint for UCX context

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 <ucp/api/ucp.h>])
                   AC_CHECK_DECLS([UCP_WORKER_ATTR_FIELD_ADDRESS_FLAGS],

ompi/mca/pml/ucx/pml_ucx.c

@@ -213,6 +213,11 @@ int mca_pml_ucx_open(void)
                                      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(&params, config, &ompi_pml_ucx.ucp_context);
     ucp_config_release(config);
 

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) {

oshmem/mca/spml/ucx/spml_ucx_component.c

@@ -212,8 +212,12 @@ static int spml_ucx_init(void)
     opal_common_ucx_mca_register();
 
     memset(&params, 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(&params, ucp_config, &mca_spml_ucx.ucp_context);
     ucp_config_release(ucp_config);
     if (UCS_OK != err) {