pml/ucx: Propagate MPI serialized thread mode

The UCX PML needs to differentiate between mpi serialized and single thread
mode according to MPI standard.

Using UCX with multiple threads, even when providing mutual exclusion, is not
completely equivalent to only using UCX with one same thread.

Signed-off-by: Thomas Vegas <tvegas@nvidia.com>

Author: tvegas1

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;
     }