UCX: Propagate MPI serialized for all worker creations

ompi/mca/pml/ucx/pml_ucx.c

@@ -283,20 +283,6 @@ 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;
@@ -310,7 +296,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 = mca_pml_ucx_thread_mode(ompi_mpi_thread_provided);
+        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

@@ -27,6 +27,8 @@
 #include "opal/util/argv.h"
 #include "opal/util/printf.h"
 
+#include "mpi.h"
+
 #include <fnmatch.h>
 #include <stdio.h>
 #include <ucm/api/ucm.h>
@@ -50,6 +52,23 @@ static void opal_common_ucx_mem_release_cb(void *buf, size_t length, void *cbdat
     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

@@ -128,6 +128,7 @@ OPAL_DECLSPEC int opal_common_ucx_del_procs_nofence(opal_common_ucx_del_proc_t *
                                                     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);
 
 /**
  * Load an integer value of \c size bytes from \c ptr and cast it to uint64_t.

oshmem/mca/spml/ucx/spml_ucx.c

@@ -1047,8 +1047,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;
     }