Merge pull request pmodels#7355 from hzhou/2503_comm_vci

ch4: support comm_set_vcis in ucx and posix

Approved-by: Ken Reffenetti

Author: hzhou

src/mpid/ch4/netmod/ofi/ofi_init.c

@@ -549,6 +549,8 @@ categories :
 === END_MPI_T_CVAR_INFO_BLOCK ===
 */
 
+static bool ofi_initialized = false;
+
 static int update_global_limits(struct fi_info *prov);
 static void dump_global_settings(void);
 static int destroy_vci_context(int vci, int nic);
@@ -799,6 +801,8 @@ int MPIDI_OFI_init_world(void)
         MPIR_ERR_CHECK(mpi_errno);
     }
 
+    ofi_initialized = true;
+
   fn_exit:
     return mpi_errno;
   fn_fail:
@@ -879,10 +883,14 @@ static int flush_send_queue(void)
 {
     int mpi_errno = MPI_SUCCESS;
 
+    if (!ofi_initialized) {
+        goto fn_exit;
+    }
+
     MPIDI_OFI_dynamic_process_request_t *reqs;
     /* TODO - Iterate over each NIC in addition to each VNI when multi-NIC within the same
      * process is implemented. */
-    int num_vcis = (MPIDI_global.is_initialized ? MPIDI_OFI_global.num_vcis : 1);
+    int num_vcis = MPIDI_OFI_global.num_vcis;
     int num_reqs = num_vcis * 2;
     reqs = MPL_malloc(sizeof(MPIDI_OFI_dynamic_process_request_t) * num_reqs, MPL_MEM_OTHER);
 
@@ -937,12 +945,9 @@ int MPIDI_OFI_mpi_finalize_hook(void)
     MPIDI_OFI_mr_key_allocator_destroy();
 
     if (strcmp("sockets", MPIDI_OFI_global.prov_use[0]->fabric_attr->prov_name) == 0) {
-        /* sockets provider need flush any last lightweight send. Only do it if we initialized
-         * world. Sockets provider can't even send self messages otherwise. */
-        if (MPIDI_global.is_initialized) {
-            mpi_errno = flush_send_queue();
-            MPIR_ERR_CHECK(mpi_errno);
-        }
+        /* sockets provider need flush any last lightweight send. */
+        mpi_errno = flush_send_queue();
+        MPIR_ERR_CHECK(mpi_errno);
     } else if (MPIR_CVAR_NO_COLLECTIVE_FINALIZE) {
         /* skip collective work arounds */
     } else if (strcmp("verbs;ofi_rxm", MPIDI_OFI_global.prov_use[0]->fabric_attr->prov_name) == 0
@@ -980,12 +985,10 @@ int MPIDI_OFI_mpi_finalize_hook(void)
     /* Tearing down endpoints in reverse order they were created */
     for (int nic = MPIDI_OFI_global.num_nics - 1; nic >= 0; nic--) {
         for (int vci = MPIDI_OFI_global.num_vcis - 1; vci >= 0; vci--) {
-            if (MPIDI_global.is_initialized || (vci == 0 && nic == 0)) {
-                /* If the user has not freed all MPI objects, ofi might not shut down cleanly.
-                 * We intentionally ignore errors to avoid crashing in finalize. Debug builds
-                 * will warn about unfreed objects/memory. */
-                (void) destroy_vci_context(vci, nic);
-            }
+            /* If the user has not freed all MPI objects, ofi might not shut down cleanly.
+             * We intentionally ignore errors to avoid crashing in finalize. Debug builds
+             * will warn about unfreed objects/memory. */
+            (void) destroy_vci_context(vci, nic);
         }
     }
 

src/mpid/ch4/netmod/ucx/ucx_impl.h

@@ -129,6 +129,9 @@ MPL_STATIC_INLINE_PREFIX bool MPIDI_UCX_is_reachable_target(int rank, MPIR_Win *
 
 #define MPIDI_UCX_WIN_AV_TO_EP(av, vci, vci_target) MPIDI_UCX_AV((av)).dest[vci][vci_target]
 
+int MPIDI_UCX_init_world(void);
+int MPIDI_UCX_init_worker(int vci);
+
 /* am handler for message sent by ucp_am_send_nb */
 ucs_status_t MPIDI_UCX_am_handler(void *arg, void *data, size_t length, ucp_ep_h reply_ep,
                                   unsigned flags);

src/mpid/ch4/netmod/ucx/ucx_init.c

@@ -18,7 +18,7 @@ categories :
 === END_MPI_T_CVAR_INFO_BLOCK ===
 */
 
-static void request_init_callback(void *request);
+static bool ucx_initialized = false;
 
 static void request_init_callback(void *request)
 {
@@ -28,8 +28,6 @@ static void request_init_callback(void *request)
 
 }
 
-static void flush_all(void);
-
 int MPIDI_UCX_init_worker(int vci)
 {
     int mpi_errno = MPI_SUCCESS;
@@ -143,68 +141,6 @@ static int initial_address_exchange(void)
     goto fn_exit;
 }
 
-int MPIDI_UCX_all_vcis_address_exchange(void)
-{
-    int mpi_errno = MPI_SUCCESS;
-
-    int size = MPIR_Process.size;
-    int rank = MPIR_Process.rank;
-    int num_vcis = MPIDI_UCX_global.num_vcis;
-
-    /* ucx address lengths are non-uniform, use MPID_MAX_BC_SIZE */
-    size_t name_len = MPID_MAX_BC_SIZE;
-
-    int my_len = num_vcis * name_len;
-    char *all_names = MPL_malloc(size * my_len, MPL_MEM_ADDRESS);
-    MPIR_Assert(all_names);
-
-    char *my_names = all_names + rank * my_len;
-
-    /* put in my addrnames */
-    for (int i = 0; i < num_vcis; i++) {
-        char *vci_addrname = my_names + i * name_len;
-        memcpy(vci_addrname, MPIDI_UCX_global.ctx[i].if_address,
-               MPIDI_UCX_global.ctx[i].addrname_len);
-    }
-    /* Allgather */
-    MPIR_Comm *comm = MPIR_Process.comm_world;
-    mpi_errno = MPIR_Allgather_allcomm_auto(MPI_IN_PLACE, 0, MPIR_BYTE_INTERNAL,
-                                            all_names, my_len, MPIR_BYTE_INTERNAL, comm,
-                                            MPIR_ERR_NONE);
-    MPIR_ERR_CHECK(mpi_errno);
-
-    /* insert the addresses */
-    ucp_ep_params_t ep_params;
-    for (int vci_local = 0; vci_local < num_vcis; vci_local++) {
-        for (int r = 0; r < size; r++) {
-            MPIDI_UCX_addr_t *av = &MPIDI_UCX_AV(&MPIDIU_get_av(0, r));
-            for (int vci_remote = 0; vci_remote < num_vcis; vci_remote++) {
-                if (vci_local == 0 && vci_remote == 0) {
-                    /* don't overwrite existing addr, or bad things will happen */
-                    continue;
-                }
-                int idx = r * num_vcis + vci_remote;
-                ep_params.field_mask = UCP_EP_PARAM_FIELD_REMOTE_ADDRESS;
-                ep_params.address = (ucp_address_t *) (all_names + idx * name_len);
-
-                ucs_status_t ucx_status;
-                ucx_status = ucp_ep_create(MPIDI_UCX_global.ctx[vci_local].worker,
-                                           &ep_params, &av->dest[vci_local][vci_remote]);
-                MPIDI_UCX_CHK_STATUS(ucx_status);
-            }
-        }
-    }
-
-    /* Flush all pending wireup operations or it may interfere with RMA flush_ops count. */
-    flush_all();
-
-  fn_exit:
-    MPL_free(all_names);
-    return mpi_errno;
-  fn_fail:
-    goto fn_exit;
-}
-
 int MPIDI_UCX_init_local(int *tag_bits)
 {
     int mpi_errno = MPI_SUCCESS;
@@ -234,7 +170,7 @@ int MPIDI_UCX_init_local(int *tag_bits)
         UCP_PARAM_FIELD_REQUEST_SIZE |
         UCP_PARAM_FIELD_ESTIMATED_NUM_EPS | UCP_PARAM_FIELD_REQUEST_INIT;
 
-    if (MPIDI_UCX_global.num_vcis > 1) {
+    if (MPICH_IS_THREADED) {
         ucp_params.mt_workers_shared = 1;
         ucp_params.field_mask |= UCP_PARAM_FIELD_MT_WORKERS_SHARED;
     }
@@ -277,6 +213,8 @@ int MPIDI_UCX_init_world(void)
     mpi_errno = initial_address_exchange();
     MPIR_ERR_CHECK(mpi_errno);
 
+    ucx_initialized = true;
+
   fn_exit:
     return mpi_errno;
   fn_fail:
@@ -286,53 +224,23 @@ int MPIDI_UCX_init_world(void)
     goto fn_exit;
 }
 
-/* static functions for MPIDI_UCX_post_init */
-static void flush_cb(void *request, ucs_status_t status)
-{
-}
-
-static void flush_all(void)
-{
-    void *reqs[MPIDI_CH4_MAX_VCIS];
-    for (int vci = 0; vci < MPIDI_UCX_global.num_vcis; vci++) {
-        reqs[vci] = ucp_worker_flush_nb(MPIDI_UCX_global.ctx[vci].worker, 0, &flush_cb);
-    }
-    for (int vci = 0; vci < MPIDI_UCX_global.num_vcis; vci++) {
-        if (reqs[vci] == NULL) {
-            continue;
-        } else if (UCS_PTR_IS_ERR(reqs[vci])) {
-            continue;
-        } else {
-            ucs_status_t status;
-            do {
-                MPID_Progress_test(NULL);
-                status = ucp_request_check_status(reqs[vci]);
-            } while (status == UCS_INPROGRESS);
-            ucp_request_release(reqs[vci]);
-        }
-    }
-}
-
 int MPIDI_UCX_post_init(void)
 {
     int mpi_errno = MPI_SUCCESS;
 
-    MPIDI_global.is_initialized = 1;
-
     return mpi_errno;
 }
 
 int MPIDI_UCX_mpi_finalize_hook(void)
 {
     int mpi_errno = MPI_SUCCESS;
 
-    if (!MPIDI_global.is_initialized) {
-        /* Nothing to do */
-        return mpi_errno;
-    }
-
     ucs_status_ptr_t ucp_request;
-    ucs_status_ptr_t *pending;
+    ucs_status_ptr_t *pending = NULL;
+
+    if (!ucx_initialized) {
+        goto fn_exit;
+    }
 
     int n = MPIDI_UCX_global.num_vcis;
     pending = MPL_malloc(sizeof(ucs_status_ptr_t) * MPIR_Process.size * n * n, MPL_MEM_OTHER);

src/mpid/ch4/netmod/ucx/ucx_vci.c

@@ -5,36 +5,122 @@
 
 #include "mpidimpl.h"
 #include "ucx_impl.h"
+#include "mpidu_bc.h"
+
+static int all_vcis_address_exchange(MPIR_Comm * comm);
+static void flush_all(void);
 
 int MPIDI_UCX_comm_set_vcis(MPIR_Comm * comm, int num_vcis, int *all_num_vcis)
 {
     int mpi_errno = MPI_SUCCESS;
 
+    MPIR_Assert(MPIDI_UCX_global.num_vcis == 1);
     MPIDI_UCX_global.num_vcis = num_vcis;
 
-    /* set up local vcis */
+    mpi_errno = MPIR_Allgather_impl(&MPIDI_UCX_global.num_vcis, 1, MPIR_INT_INTERNAL,
+                                    all_num_vcis, 1, MPIR_INT_INTERNAL, comm, MPIR_ERR_NONE);
+    MPIR_ERR_CHECK(mpi_errno);
+
     for (int i = 1; i < MPIDI_UCX_global.num_vcis; i++) {
         mpi_errno = MPIDI_UCX_init_worker(i);
         MPIR_ERR_CHECK(mpi_errno);
     }
 
-    /* UCX netmod only support the same number of vcis on all procs */
-    for (int i = 0; i < comm->local_size; i++) {
-        all_num_vcis[i] = num_vcis;
-    }
-
-    /* address exchange */
-    if (num_vcis > 1) {
-        mpi_errno = MPIDI_UCX_all_vcis_address_exchange();
-        MPIR_ERR_CHECK(mpi_errno);
-    }
+    mpi_errno = all_vcis_address_exchange(comm);
+    MPIR_ERR_CHECK(mpi_errno);
 
     if (MPIR_CVAR_DEBUG_SUMMARY && comm->rank == 0) {
         printf("num_vcis: %d\n", MPIDI_UCX_global.num_vcis);
     }
+    /* Flush all pending wireup operations or it may interfere with RMA flush_ops count.
+     * Since this require progress in non-zero vcis, we need switch on is_initialized. */
+    flush_all();
 
   fn_exit:
     return mpi_errno;
   fn_fail:
     goto fn_exit;
 }
+
+static int all_vcis_address_exchange(MPIR_Comm * comm)
+{
+    int mpi_errno = MPI_SUCCESS;
+
+    int size = MPIR_Process.size;
+    int rank = MPIR_Process.rank;
+    int num_vcis = MPIDI_UCX_global.num_vcis;
+
+    /* ucx address lengths are non-uniform, use MPID_MAX_BC_SIZE */
+    size_t name_len = MPID_MAX_BC_SIZE;
+
+    int my_len = num_vcis * name_len;
+    char *all_names = MPL_malloc(size * my_len, MPL_MEM_ADDRESS);
+    MPIR_Assert(all_names);
+
+    char *my_names = all_names + rank * my_len;
+
+    /* put in my addrnames */
+    for (int i = 0; i < num_vcis; i++) {
+        char *vci_addrname = my_names + i * name_len;
+        memcpy(vci_addrname, MPIDI_UCX_global.ctx[i].if_address,
+               MPIDI_UCX_global.ctx[i].addrname_len);
+    }
+    /* Allgather */
+    mpi_errno = MPIR_Allgather_allcomm_auto(MPI_IN_PLACE, 0, MPIR_BYTE_INTERNAL,
+                                            all_names, my_len, MPIR_BYTE_INTERNAL,
+                                            comm, MPIR_ERR_NONE);
+    MPIR_ERR_CHECK(mpi_errno);
+
+    /* insert the addresses */
+    ucp_ep_params_t ep_params;
+    for (int vci_local = 0; vci_local < num_vcis; vci_local++) {
+        for (int r = 0; r < size; r++) {
+            MPIDI_UCX_addr_t *av = &MPIDI_UCX_AV(&MPIDIU_get_av(0, r));
+            for (int vci_remote = 0; vci_remote < num_vcis; vci_remote++) {
+                if (vci_local == 0 && vci_remote == 0) {
+                    /* don't overwrite existing addr, or bad things will happen */
+                    continue;
+                }
+                int idx = r * num_vcis + vci_remote;
+                ep_params.field_mask = UCP_EP_PARAM_FIELD_REMOTE_ADDRESS;
+                ep_params.address = (ucp_address_t *) (all_names + idx * name_len);
+
+                ucs_status_t ucx_status;
+                ucx_status = ucp_ep_create(MPIDI_UCX_global.ctx[vci_local].worker,
+                                           &ep_params, &av->dest[vci_local][vci_remote]);
+                MPIDI_UCX_CHK_STATUS(ucx_status);
+            }
+        }
+    }
+  fn_exit:
+    MPL_free(all_names);
+    return mpi_errno;
+  fn_fail:
+    goto fn_exit;
+}
+
+static void flush_cb(void *request, ucs_status_t status)
+{
+}
+
+static void flush_all(void)
+{
+    void *reqs[MPIDI_CH4_MAX_VCIS];
+    for (int vci = 0; vci < MPIDI_UCX_global.num_vcis; vci++) {
+        reqs[vci] = ucp_worker_flush_nb(MPIDI_UCX_global.ctx[vci].worker, 0, &flush_cb);
+    }
+    for (int vci = 0; vci < MPIDI_UCX_global.num_vcis; vci++) {
+        if (reqs[vci] == NULL) {
+            continue;
+        } else if (UCS_PTR_IS_ERR(reqs[vci])) {
+            continue;
+        } else {
+            ucs_status_t status;
+            do {
+                MPID_Progress_test(NULL);
+                status = ucp_request_check_status(reqs[vci]);
+            } while (status == UCS_INPROGRESS);
+            ucp_request_release(reqs[vci]);
+        }
+    }
+}

src/mpid/ch4/shm/posix/eager/include/posix_eager.h

@@ -12,6 +12,7 @@
 
 typedef int (*MPIDI_POSIX_eager_init_t) (int rank, int size);
 typedef int (*MPIDI_POSIX_eager_post_init_t) (void);
+typedef int (*MPIDI_POSIX_eager_set_vcis_t) (MPIR_Comm * comm, int num_vcis);
 typedef int (*MPIDI_POSIX_eager_finalize_t) (void);
 
 typedef int (*MPIDI_POSIX_eager_send_t) (int grank, MPIDI_POSIX_am_header_t * msg_hdr,
@@ -37,6 +38,7 @@ typedef size_t(*MPIDI_POSIX_eager_buf_limit_t) (void);
 typedef struct {
     MPIDI_POSIX_eager_init_t init;
     MPIDI_POSIX_eager_post_init_t post_init;
+    MPIDI_POSIX_eager_set_vcis_t set_vcis;
     MPIDI_POSIX_eager_finalize_t finalize;
 
     MPIDI_POSIX_eager_send_t send;
@@ -59,6 +61,7 @@ extern char MPIDI_POSIX_eager_strings[][MPIDI_MAX_POSIX_EAGER_STRING_LEN];
 
 int MPIDI_POSIX_eager_init(int rank, int size);
 int MPIDI_POSIX_eager_post_init(void);
+int MPIDI_POSIX_eager_set_vcis(MPIR_Comm * comm, int num_vcis);
 int MPIDI_POSIX_eager_finalize(void);
 
 MPL_STATIC_INLINE_PREFIX int MPIDI_POSIX_eager_send(int grank, MPIDI_POSIX_am_header_t * msg_hdr,