UCX: Propagate MPI serialized for all worker creations

Move MPI to UCX thread mode function to common source. Also use
serialized mode for all oshmem initializations.

Signed-off-by: Thomas Vegas <[email protected]>
(cherry picked from commit 44cc663)

Author: tvegas1

ompi/mca/pml/ucx/pml_ucx.c

@@ -293,7 +293,8 @@ int mca_pml_ucx_init(int enable_mpi_threads)
     if (enable_mpi_threads) {
         params.thread_mode = UCS_THREAD_MODE_MULTI;
     } else {
-        params.thread_mode = UCS_THREAD_MODE_SINGLE;
+        params.thread_mode =
+            opal_common_ucx_thread_mode(ompi_mpi_thread_provided);
     }
 
 #if HAVE_DECL_UCP_WORKER_FLAG_IGNORE_REQUEST_LEAK

opal/mca/common/ucx/common_ucx.c

@@ -25,7 +25,10 @@
 #include "opal/util/argv.h"
 #include "opal/util/printf.h"
 
+#include "mpi.h"
+
 #include <ucm/api/ucm.h>
+
 #include <fnmatch.h>
 #include <stdio.h>
 
@@ -49,6 +52,23 @@ static void opal_common_ucx_mem_release_cb(void *buf, size_t length,
     ucm_vm_munmap(buf, length);
 }
 
+ucs_thread_mode_t opal_common_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:
+        return UCS_THREAD_MODE_SINGLE;
+    default:
+        MCA_COMMON_UCX_WARN("Unknown MPI thread mode %d, using multithread",
+                            ompi_mode);
+        return UCS_THREAD_MODE_MULTI;
+    }
+}
+
 OPAL_DECLSPEC void opal_common_ucx_mca_var_register(const mca_base_component_t *component)
 {
     char *default_tls = "rc_verbs,ud_verbs,rc_mlx5,dc_mlx5,ud_mlx5,cuda_ipc,rocm_ipc";

opal/mca/common/ucx/common_ucx.h

@@ -124,6 +124,7 @@ OPAL_DECLSPEC int opal_common_ucx_del_procs(opal_common_ucx_del_proc_t *procs, s
 OPAL_DECLSPEC int opal_common_ucx_del_procs_nofence(opal_common_ucx_del_proc_t *procs, size_t count,
                                                size_t my_rank, size_t max_disconnect, ucp_worker_h worker);
 OPAL_DECLSPEC void opal_common_ucx_mca_var_register(const mca_base_component_t *component);
+OPAL_DECLSPEC ucs_thread_mode_t opal_common_ucx_thread_mode(int ompi_mode);
 
 static inline
 ucs_status_t opal_common_ucx_request_status(ucs_status_ptr_t request)

oshmem/mca/spml/ucx/spml_ucx.c

@@ -1016,8 +1016,11 @@ static int mca_spml_ucx_ctx_create_common(long options, mca_spml_ucx_ctx_t **ucx
     ucx_ctx->strong_sync = mca_spml_ucx_ctx_default.strong_sync;      
 
     params.field_mask  = UCP_WORKER_PARAM_FIELD_THREAD_MODE;
-    if (oshmem_mpi_thread_provided == SHMEM_THREAD_SINGLE || options & SHMEM_CTX_PRIVATE || options & SHMEM_CTX_SERIALIZED) {
+    if (oshmem_mpi_thread_provided == SHMEM_THREAD_SINGLE ||
+        oshmem_mpi_thread_provided == SHMEM_THREAD_FUNNELED || options & SHMEM_CTX_PRIVATE) {
         params.thread_mode = UCS_THREAD_MODE_SINGLE;
+    } else if (oshmem_mpi_thread_provided == SHMEM_THREAD_SERIALIZED || options & SHMEM_CTX_SERIALIZED) {
+        params.thread_mode = UCS_THREAD_MODE_SERIALIZED;
     } else {
         params.thread_mode = UCS_THREAD_MODE_MULTI;
     }