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

@@ -197,22 +197,27 @@ int mca_pml_ucx_open(void)
     }
 
     /* Initialize UCX context */
-    params.field_mask      = UCP_PARAM_FIELD_FEATURES |
+    params.field_mask        = UCP_PARAM_FIELD_FEATURES |
-                             UCP_PARAM_FIELD_REQUEST_SIZE |
+                               UCP_PARAM_FIELD_REQUEST_SIZE |
-                             UCP_PARAM_FIELD_REQUEST_INIT |
+                               UCP_PARAM_FIELD_REQUEST_INIT |
-                             UCP_PARAM_FIELD_REQUEST_CLEANUP |
+                               UCP_PARAM_FIELD_REQUEST_CLEANUP |
-                             UCP_PARAM_FIELD_TAG_SENDER_MASK |
+                               UCP_PARAM_FIELD_TAG_SENDER_MASK |
-                             UCP_PARAM_FIELD_MT_WORKERS_SHARED |
+                               UCP_PARAM_FIELD_MT_WORKERS_SHARED |
-                             UCP_PARAM_FIELD_ESTIMATED_NUM_EPS;
+                               UCP_PARAM_FIELD_ESTIMATED_NUM_EPS;
-    params.features        = UCP_FEATURE_TAG;
+    params.features          = UCP_FEATURE_TAG;
-    params.request_size    = sizeof(ompi_request_t);
+    params.request_size      = sizeof(ompi_request_t);
-    params.request_init    = mca_pml_ucx_request_init;
+    params.request_init      = mca_pml_ucx_request_init;
-    params.request_cleanup = mca_pml_ucx_request_cleanup;
+    params.request_cleanup   = mca_pml_ucx_request_cleanup;
-    params.tag_sender_mask = PML_UCX_SPECIFIC_SOURCE_MASK;
+    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(&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.field_mask        = UCP_PARAM_FIELD_FEATURES          |
-    params.features   = UCP_FEATURE_RMA|UCP_FEATURE_AMO32|UCP_FEATURE_AMO64;
+                               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) {