OSHMEM/MCA/SPML/UCX: WIP

Signed-off-by: Roie Danino <[email protected]>

Author: roiedanino

oshmem/mca/spml/ucx/spml_ucx.c

@@ -292,6 +292,17 @@ static void mca_spml_ucx_rkey_store_put(mca_spml_ucx_rkey_store_t *store,
     ucp_rkey_destroy(rkey);
 }
 
+static void mca_spml_ucx_team_world_init()
+{
+    int rc = mca_spml_ucx_team_split_strided(NULL, 0, 1, oshmem_num_procs(), NULL, 0,
+                                             &SHMEM_TEAM_WORLD);
+    
+    if (rc != OSHMEM_SUCCESS) {
+        SPML_UCX_ERROR("mca_spml_ucx_team_split_strided failed (SHMEM_TEAM_WORLD creation)");
+        oshmem_shmem_abort(-1);
+    }
+}
+
 int mca_spml_ucx_enable(bool enable)
 {
     SPML_UCX_VERBOSE(50, "*** ucx ENABLED ****");
@@ -315,8 +326,7 @@ void mca_spml_ucx_peer_mkey_cache_init(mca_spml_ucx_ctx_t *ucx_ctx, int pe)
 int mca_spml_ucx_peer_mkey_cache_add(ucp_peer_t *ucp_peer, int index)
 {
     /* Allocate an array to hold the pointers to the ucx_cached_mkey */
-    if (index >= (int)ucp_peer->mkeys_cnt){
-        int old_size = ucp_peer->mkeys_cnt;
+    if (index >= (int)ucp_peer->mkeys_cnt){        int old_size = ucp_peer->mkeys_cnt;
         ucp_peer->mkeys_cnt = index + 1;
         ucp_peer->mkeys = realloc(ucp_peer->mkeys, sizeof(ucp_peer->mkeys[0]) * ucp_peer->mkeys_cnt);
         if (NULL == ucp_peer->mkeys) {
@@ -451,6 +461,14 @@ int mca_spml_ucx_ctx_mkey_del(mca_spml_ucx_ctx_t *ucx_ctx, int pe, uint32_t segn
     return OSHMEM_SUCCESS;
 }
 
+static void mca_spml_ucx_team_world_destroy()
+{
+    if (SHMEM_TEAM_WORLD != NULL) {
+        mca_spml_ucx_team_destroy(SHMEM_TEAM_WORLD);
+        SHMEM_TEAM_WORLD = NULL;
+    }
+}
+
 int mca_spml_ucx_del_procs(oshmem_group_t* group, size_t nprocs)
 {
     size_t ucp_workers = mca_spml_ucx.ucp_workers;
@@ -460,6 +478,8 @@ int mca_spml_ucx_del_procs(oshmem_group_t* group, size_t nprocs)
 
     oshmem_shmem_barrier();
 
+    mca_spml_ucx_team_world_destroy();
+
     if (!mca_spml_ucx_ctx_default.ucp_peers) {
         return OSHMEM_SUCCESS;
     }
@@ -1163,8 +1183,6 @@ int mca_spml_ucx_ctx_create(long options, shmem_ctx_t *ctx)
         }
         SHMEM_MUTEX_UNLOCK(mca_spml_ucx.internal_mutex);
     }
- 
-    mca_spml_ucx_team_world_init();
 
     (*ctx) = (shmem_ctx_t)ucx_ctx;
     return OSHMEM_SUCCESS;
@@ -1183,8 +1201,6 @@ void mca_spml_ucx_ctx_destroy(shmem_ctx_t ctx)
         opal_progress_unregister(spml_ucx_ctx_progress);
     }
 
-    mca_spml_ucx_team_world_destroy();
-
     SHMEM_MUTEX_UNLOCK(mca_spml_ucx.internal_mutex);
 }
 
@@ -1754,25 +1770,6 @@ size_t mca_spml_ucx_test_some_vector(void *ivars, int cmp,
     RUNTIME_SHMEM_NOT_IMPLEMENTED_API_ABORT_RET_SIZE_T();
 }
 
-void mca_spml_ucx_team_world_init()
-{
-    int rc = mca_spml_ucx_team_split_strided(NULL, 0, 1, oshmem_num_procs(), NULL, 0,
-                                             &SHMEM_TEAM_WORLD);
-    
-    if (rc != OSHMEM_SUCCESS) {
-        SPML_UCX_ERROR("mca_spml_ucx_team_split_strided failed (SHMEM_TEAM_WORLD creation)");
-        oshmem_shmem_abort(-1);
-    }
-}
-
-void mca_spml_ucx_team_world_destroy()
-{
-    if (SHMEM_TEAM_WORLD != NULL) {
-        mca_spml_ucx_team_destroy(SHMEM_TEAM_WORLD);
-        SHMEM_TEAM_WORLD = NULL;
-    }
-}
-
 /* This routine is not implemented */
 int mca_spml_ucx_team_sync(shmem_team_t team)
 {
@@ -1820,16 +1817,16 @@ int mca_spml_ucx_team_translate_pe(shmem_team_t src_team, int src_pe,
     mca_spml_ucx_team_t *ucx_dest_team = (mca_spml_ucx_team_t*) dest_team;
     int global_pe;
 
-    if (src_pe == SPML_UCX_PE_NOT_IN_TEAM || (src_team == dest_team)) {
+    if ((src_pe == SPML_UCX_PE_NOT_IN_TEAM) || (src_team == dest_team)) {
         return src_pe;
     }
 
-    if (src_team == dest_team) {
-        return src_pe;
-    }
-    
     global_pe = ucx_src_team->start + src_pe * ucx_src_team->stride;
 
+    SPML_UCX_WARN("team_translate_pe(src_team=%p, src_pe=%d, dest_team=%p), global pe: %d, "
+                  "src_team->start: %d, src pe: %d, src_team->stride: %d",
+                  src_team, src_pe, dest_team, global_pe, ucx_src_team->start, src_pe, ucx_src_team->stride);
+
     if (dest_team == SHMEM_TEAM_WORLD) {
         return global_pe;
     }
@@ -1849,32 +1846,44 @@ int mca_spml_ucx_team_split_strided(shmem_team_t parent_team, int start, int
     mca_spml_ucx_team_t *ucx_parent_team;
     mca_spml_ucx_team_t *ucx_new_team;
     int my_pe;
-    int n_pes;
 
     SPML_UCX_ASSERT(((start + size * stride) <= oshmem_num_procs()) && (start < size) && (stride > 0) && (size > 0));
 
+    SPML_UCX_WARN("team_split_strided(parent_team=%p, start=%d, stride=%d, size=%d, config=%p, "
+                  "config_mask=%ld, new_team=%p)",
+                  parent_team, start, stride, size, config, config_mask, new_team);
+
+    ucx_new_team = (mca_spml_ucx_team_t *)malloc(sizeof(mca_spml_ucx_team_t));
+    ucx_new_team->start = start;
+    ucx_new_team->stride = stride;
+
     if (parent_team == NULL) {
         my_pe = shmem_my_pe();
     } else {
         ucx_parent_team = (mca_spml_ucx_team_t*) parent_team;
+
         SPML_UCX_VALIDATE_TEAM(parent_team);
         if (mca_spml_ucx_is_pe_in_strided_team(ucx_parent_team->my_pe, start, stride, size)) {
             my_pe = (ucx_parent_team->my_pe - start) / stride;
+            SPML_UCX_WARN("split: my_pe at parent team: %d, start: %d, stride: %d, size: %d, "
+                          "my_pe at new team: %d", ucx_parent_team->my_pe, start, stride, size, my_pe);
         } else {
             /* not in team, according to spec it should be SHMEM_TEAM_INVALID but its value is NULL which
             can be also interpreted as 0 (first pe), therefore -1 is used */
+
+            SPML_UCX_WARN("pe #%d is not part of the new team", ucx_parent_team->my_pe);
             my_pe = SPML_UCX_PE_NOT_IN_TEAM;
         }
+
+        /* In order to simplify pe translations start and stride are calculated with respect to
+         * world_team */
+        ucx_new_team->start += ucx_parent_team->start;
+        ucx_new_team->stride *= ucx_parent_team->stride;
     }
 
-    ucx_new_team = (mca_spml_ucx_team_t *)malloc(sizeof(mca_spml_ucx_team_t));
     ucx_new_team->n_pes = size;
     ucx_new_team->my_pe = my_pe;
 
-    /* In order to simplify pe translations start and stride are calculated with respect to
-     * world_team */
-    ucx_new_team->start = ucx_parent_team->start + start;
-    ucx_new_team->stride = ucx_parent_team->stride * stride;
     ucx_new_team->config = calloc(1, sizeof(mca_spml_ucx_team_config_t));
 
     if (config != NULL) {

oshmem/mca/spml/ucx/spml_ucx.h

@@ -130,7 +130,6 @@ typedef struct mca_spml_ucx_team_config {
 } mca_spml_ucx_team_config_t;
 
 typedef struct mca_spml_ucx_team {
-    shmem_team_t                super;
     int                         n_pes;
     int                         my_pe;
     int                         stride;
@@ -324,9 +323,6 @@ mca_spml_ucx_mem_map_flags_symmetric_rkey(struct mca_spml_ucx *spml_ucx);
 extern void mca_spml_ucx_rkey_store_init(mca_spml_ucx_rkey_store_t *store);
 extern void mca_spml_ucx_rkey_store_cleanup(mca_spml_ucx_rkey_store_t *store);
 
-void mca_spml_ucx_team_world_init();
-void mca_spml_ucx_team_world_destroy();
-
 static inline int
 mca_spml_ucx_peer_mkey_get(ucp_peer_t *ucp_peer, int index, spml_ucx_cached_mkey_t **out_rmkey)
 {