Passing estimated_num_procs to UCX init in PML and SPML.

Signed-off-by: Xin Zhao <[email protected]>

Author: xinzhao3

ompi/mca/pml/ucx/pml_ucx.c

@@ -136,15 +136,16 @@ int mca_pml_ucx_open(void)
                              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_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();
 
     status = ucp_init(&params, config, &ompi_pml_ucx.ucp_context);
     ucp_config_release(config);

oshmem/mca/spml/ucx/spml_ucx_component.c

@@ -126,8 +126,9 @@ static int mca_spml_ucx_component_open(void)
     }
 
     memset(&params, 0, sizeof(params));
-    params.field_mask = UCP_PARAM_FIELD_FEATURES;
+    params.field_mask = UCP_PARAM_FIELD_FEATURES|UCP_PARAM_FIELD_ESTIMATED_NUM_EPS;
     params.features   = UCP_FEATURE_RMA|UCP_FEATURE_AMO32|UCP_FEATURE_AMO64;
+    params.estimated_num_eps = oshmem_num_procs();
 
     err = ucp_init(&params, ucp_config, &mca_spml_ucx.ucp_context);
     ucp_config_release(ucp_config);