ch4/ofi: relax av insertion order of non-root vcis

We used restrict all remote vcis to be inserted on each local vci in the
exact same order. This forces us to initialize all vcis at init time.

Relax the restriction by allowing local non-root vci to have a subset of
remote vcis to be inserted than on the local root vci. This allows
enabling multiple vcis on a sub communicator rather than always
comm_world.

Because we insert all remote endpoints to all local non-root-endpoints at the
same time, thus follow the exact same insertion order, they will share
the same av table index except for the local root endpoint because it
has inserted other remote root endpoints at init time. The local root to
remote non-root endpoints will have a fixed offset from that of local
non-root.

Author: hzhou

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

@@ -7,6 +7,24 @@
 #include "ofi_impl.h"
 #include "ofi_init.h"
 
+/* NOTE on av insertion order:
+ *
+ * Each nic-vci is an endpoint with a unique address, and inside libfabric maintains
+ * one av table. Thus to fully store the address mapping, we'll need a multi-dim table as
+ *     av_table[src_vci][src_nic][dest_rank][dest_vci][dest_nic]
+ * Note, this table is for illustration, and different from MPIDI_OFI_addr_t.
+ *
+ * However, if we insert the addresses in the same order between local endpoints, then the
+ * av table indexes will be *identical*. Then, we can omit the dimension of [src_vci][src_nic].
+ * I.e. we only need av_table[rank][vci][nic], saving two dimensions of local vcis and local nics.
+ *
+ * To achieve that, we need always insert each remote address on *all* local endpoints together.
+ * Because we separate root addr (av_table[0][0][rank][0][0]) separately, we allow the root
+ * address to be inserted separately from the rest. The rest of the addresses are only
+ * needed when multiple vcis/nics are enabled. But we require for each remote rank, all remote
+ * endpoints to be inserted all at once.
+ */
+
 int MPIDI_OFI_vci_init(void)
 {
     MPIDI_OFI_global.num_nics = 1;
@@ -173,7 +191,7 @@ static int setup_additional_vcis(void)
 #endif
 
 #define GET_AV_AND_ADDRNAMES(rank) \
-    MPIDI_OFI_addr_t *av ATTRIBUTE((unused)) = &MPIDI_OFI_AV(&MPIDIU_get_av(0, rank)); \
+    MPIDI_av_entry_t *av ATTRIBUTE((unused)) = &MPIDIU_get_av(0, rank); \
     char *r_names = all_names + rank * max_vcis * num_nics * name_len;
 
 #define DO_AV_INSERT(ctx_idx, nic, vci) \
@@ -237,9 +255,9 @@ static int addr_exchange_all_ctx(MPIR_Comm * comm)
     int mpi_errno = MPI_SUCCESS;
     MPIR_CHKLMEM_DECL();
 
-    MPIR_Comm *comm = MPIR_Process.comm_world;
-    int size = MPIR_Process.size;
-    int rank = MPIR_Process.rank;
+    MPIR_Assert(comm == MPIR_Process.comm_world);
+    int size = comm->local_size;
+    int rank = comm->rank;
 
     int max_vcis;
     int *all_num_vcis;
@@ -270,9 +288,10 @@ static int addr_exchange_all_ctx(MPIR_Comm * comm)
         goto fn_exit;
     }
 
-    /* allocate additional av addrs */
+    /* allocate all_dest[] in av entry */
     for (int i = 0; i < size; i++) {
-        MPIDI_av_entry_t *av = &MPIDIU_get_av(0, i);
+        MPIDI_av_entry_t *av = MPIDIU_comm_rank_to_av(comm, i);
+        MPIR_Assert(MPIDI_OFI_AV(av).all_dest == NULL);
         MPIDI_OFI_AV(av).all_dest = MPL_malloc(max_vcis * num_nics * sizeof(fi_addr_t),
                                                MPL_MEM_ADDRESS);
         MPIR_ERR_CHKANDJUMP(!MPIDI_OFI_AV(av).all_dest, mpi_errno, MPI_ERR_OTHER, "**nomem");
@@ -300,94 +319,49 @@ static int addr_exchange_all_ctx(MPIR_Comm * comm)
     mpi_errno = MPIR_Allgather_fallback(MPI_IN_PLACE, 0, MPIR_BYTE_INTERNAL,
                                         all_names, my_len, MPIR_BYTE_INTERNAL, comm, MPIR_ERR_NONE);
 
-    /* Step 2: insert and store non-root nic/vci on the root context */
-    int root_ctx_idx = MPIDI_OFI_get_ctx_index(0, 0);
+    /* insert and store non-root nic/vci on the root context */
     for (int r = 0; r < size; r++) {
+        fi_addr_t expect_addr = FI_ADDR_NOTAVAIL;
+        fi_addr_t root_offset = 0;
         GET_AV_AND_ADDRNAMES(r);
         for (int nic = 0; nic < num_nics; nic++) {
             for (int vci = 0; vci < NUM_VCIS_FOR_RANK(r); vci++) {
-                SKIP_ROOT(nic, vci);
-                DO_AV_INSERT(root_ctx_idx, nic, vci);
-                MPIDI_OFI_AV_ADDR(av, 0, 0, vci, nic) = addr;
-            }
-        }
-    }
-
-    /* Step 3: insert all nic/vci on non-root context, following exact order as step 1 and 2 */
-
-    int *is_node_roots = NULL;
-    if (MPIR_CVAR_CH4_ROOTS_ONLY_PMI) {
-        MPIR_CHKLMEM_MALLOC(is_node_roots, size * sizeof(int));
-        for (int r = 0; r < size; r++) {
-            is_node_roots[r] = 0;
-        }
-        for (int i = 0; i < MPIR_Process.num_nodes; i++) {
-            is_node_roots[MPIR_Process.node_root_map[i]] = 1;
-        }
-    }
-
-    for (int nic_local = 0; nic_local < num_nics; nic_local++) {
-        for (int vci_local = 0; vci_local < num_vcis; vci_local++) {
-            SKIP_ROOT(nic_local, vci_local);
-            int ctx_idx = MPIDI_OFI_get_ctx_index(vci_local, nic_local);
-
-            /* -- same order as step 1 -- */
-            if (MPIR_CVAR_CH4_ROOTS_ONLY_PMI) {
-                /* node roots */
-                for (int r = 0; r < size; r++) {
-                    if (is_node_roots[r]) {
-                        GET_AV_AND_ADDRNAMES(r);
-                        DO_AV_INSERT(ctx_idx, 0, 0);
-                        MPIR_Assert(MPIDI_OFI_AV_ROOT_ADDR(av) == addr);
-                    }
-                }
-                /* non-node-root */
-                for (int r = 0; r < size; r++) {
-                    if (!is_node_roots[r]) {
-                        GET_AV_AND_ADDRNAMES(r);
-                        DO_AV_INSERT(ctx_idx, 0, 0);
-                        MPIR_Assert(MPIDI_OFI_AV_ROOT_ADDR(av) == addr);
-                    }
-                }
-            } else {
-                /* !MPIR_CVAR_CH4_ROOTS_ONLY_PMI */
-                for (int r = 0; r < size; r++) {
-                    GET_AV_AND_ADDRNAMES(r);
-                    DO_AV_INSERT(ctx_idx, 0, 0);
-                    MPIR_Assert(MPIDI_OFI_AV_ROOT_ADDR(av) == addr);
-                }
-            }
-
-            /* -- same order as step 2 -- */
-            for (int r = 0; r < size; r++) {
-                GET_AV_AND_ADDRNAMES(r);
-                for (int nic = 0; nic < num_nics; nic++) {
-                    for (int vci = 0; vci < NUM_VCIS_FOR_RANK(r); vci++) {
-                        SKIP_ROOT(nic, vci);
+                /* for each local endpoints */
+                for (int nic_local = 0; nic_local < num_nics; nic_local++) {
+                    for (int vci_local = 0; vci_local < num_vcis; vci_local++) {
+                        /* skip root */
+                        if (nic == 0 && vci == 0 && nic_local == 0 && vci_local == 0) {
+                            continue;
+                        }
+                        int ctx_idx = MPIDI_OFI_get_ctx_index(vci_local, nic_local);
                         DO_AV_INSERT(ctx_idx, nic, vci);
-                        MPIR_Assert(MPIDI_OFI_AV_ADDR(av, 0, 0, vci, nic) == addr);
+                        /* we expect all resulting addr to be the same except for local root endpoint, which
+                         * will have an offset */
+                        if (expect_addr == FI_ADDR_NOTAVAIL) {
+                            expect_addr = addr;
+                        } else if (nic_local == 0 && vci_local == 0) {
+                            if (root_offset == 0) {
+                                root_offset = addr - expect_addr;
+                            } else {
+                                MPIR_Assert(addr == expect_addr + root_offset);
+                            }
+                        } else {
+                            MPIR_Assert(addr == expect_addr);
+                        }
                     }
                 }
+                MPIR_Assert(expect_addr != FI_ADDR_NOTAVAIL);
+                MPIDI_OFI_AV_ADDR_NO_OFFSET(av, vci, nic) = expect_addr;
+                /* next */
+                expect_addr++;
             }
         }
+        MPIDI_OFI_AV(av).root_offset = root_offset;
     }
+
     mpi_errno = MPIR_Barrier_fallback(comm, MPIR_ERR_NONE);
     MPIR_ERR_CHECK(mpi_errno);
 
-    /* check */
-#if MPIDI_CH4_MAX_VCIS > 1
-    if (MPIDI_OFI_ENABLE_AV_TABLE) {
-        for (int r = 0; r < size; r++) {
-            MPIDI_OFI_addr_t *av ATTRIBUTE((unused)) = &MPIDI_OFI_AV(&MPIDIU_get_av(0, r));
-            for (int nic = 0; nic < num_nics; nic++) {
-                for (int vci = 0; vci < NUM_VCIS_FOR_RANK(r); vci++) {
-                    MPIR_Assert(MPIDI_OFI_AV_ADDR(av, 0, 0, vci, nic) ==
-                                get_av_table_index(r, nic, vci, all_num_vcis));
-                }
-            }
-        }
-    }
-#endif
   fn_exit:
     MPIR_CHKLMEM_FREEALL();
     return mpi_errno;