Merge pull request #12577 from tvegas1/ucx_thread_mode

pml/ucx: Propagate MPI serialized thread mode

Author: bosilca

ompi/mca/pml/ucx/pml_ucx.c

@@ -283,6 +283,20 @@ int mca_pml_ucx_close(void)
     return OMPI_SUCCESS;
 }
 
+static ucs_thread_mode_t mca_pml_ucx_thread_mode(int ompi_mode)
+{
+    switch (ompi_mode) {
+    case MPI_THREAD_MULTIPLE:
+        return UCS_THREAD_MODE_MULTI;
+    case MPI_THREAD_SERIALIZED:
+        return UCS_THREAD_MODE_SERIALIZED;
+    case MPI_THREAD_FUNNELED:
+    case MPI_THREAD_SINGLE:
+    default:
+        return UCS_THREAD_MODE_SINGLE;
+    }
+}
+
 int mca_pml_ucx_init(int enable_mpi_threads)
 {
     ucp_worker_params_t params;
@@ -292,12 +306,11 @@ int mca_pml_ucx_init(int enable_mpi_threads)
 
     PML_UCX_VERBOSE(1, "mca_pml_ucx_init");
 
-    /* TODO check MPI thread mode */
     params.field_mask  = UCP_WORKER_PARAM_FIELD_THREAD_MODE;
     if (enable_mpi_threads) {
         params.thread_mode = UCS_THREAD_MODE_MULTI;
     } else {
-        params.thread_mode = UCS_THREAD_MODE_SINGLE;
+        params.thread_mode = mca_pml_ucx_thread_mode(ompi_mpi_thread_provided);
     }
 
 #if HAVE_DECL_UCP_WORKER_FLAG_IGNORE_REQUEST_LEAK

oshmem/mca/spml/ucx/spml_ucx_component.c

@@ -322,6 +322,8 @@ static int spml_ucx_init(void)
     wkr_params.field_mask  = UCP_WORKER_PARAM_FIELD_THREAD_MODE;
     if (oshmem_mpi_thread_requested == SHMEM_THREAD_MULTIPLE) {
         wkr_params.thread_mode = UCS_THREAD_MODE_MULTI;
+    } else if (oshmem_mpi_thread_requested == SHMEM_THREAD_SERIALIZED) {
+        wkr_params.thread_mode = UCS_THREAD_MODE_SERIALIZED;
     } else {
         wkr_params.thread_mode = UCS_THREAD_MODE_SINGLE;
     }