OSHMEM/MCA/SPML/UCX: removed RC checks where not needed, fixed bugs

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

Author: roiedanino

oshmem/mca/spml/ucx/spml_ucx.c

@@ -292,17 +292,6 @@ 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 ****");
@@ -461,14 +450,6 @@ 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;
@@ -478,8 +459,6 @@ 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;
     }
@@ -1779,15 +1758,21 @@ int mca_spml_ucx_team_sync(shmem_team_t team)
 int mca_spml_ucx_team_my_pe(shmem_team_t team)
 {
     mca_spml_ucx_team_t *ucx_team = (mca_spml_ucx_team_t *)team;
-    SPML_UCX_VALIDATE_TEAM(team);
+
+    if (team == SHMEM_TEAM_WORLD) {
+        return shmem_my_pe();
+    }
 
     return ucx_team->my_pe;
 }
 
 int mca_spml_ucx_team_n_pes(shmem_team_t team)
 {
     mca_spml_ucx_team_t *ucx_team = (mca_spml_ucx_team_t *)team;
-    SPML_UCX_VALIDATE_TEAM(team);
+
+    if (team == SHMEM_TEAM_WORLD) {
+        return shmem_n_pes();
+    }
 
     return ucx_team->n_pes;
 }
@@ -1823,10 +1808,6 @@ int mca_spml_ucx_team_translate_pe(shmem_team_t src_team, int 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;
     }
@@ -1845,42 +1826,38 @@ 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 base_pe;
+    int base_start;
+    int base_stride;
     int my_pe;
 
     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();
+    if (parent_team == SHMEM_TEAM_WORLD) {
+        base_pe = shmem_my_pe();
+        base_start = 0;
+        base_stride = 1;
     } else {
         ucx_parent_team = (mca_spml_ucx_team_t*) parent_team;
+        base_pe = ucx_parent_team->my_pe;
+        base_start = ucx_parent_team->start;
+        base_stride = ucx_parent_team->stride;
+    }
 
-        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;
+    if (mca_spml_ucx_is_pe_in_strided_team(base_pe, start, stride, size)) {
+        my_pe = (base_pe - start) / stride;
+    } 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 */
+        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 = (mca_spml_ucx_team_t *)malloc(sizeof(mca_spml_ucx_team_t));
+    ucx_new_team->start = base_start + start;
+    ucx_new_team->stride = base_stride * stride;
+
     ucx_new_team->n_pes = size;
     ucx_new_team->my_pe = my_pe;
 
@@ -1890,10 +1867,10 @@ int mca_spml_ucx_team_split_strided(shmem_team_t parent_team, int start, int
         memcpy(&ucx_new_team->config->super, config, sizeof(shmem_team_config_t));
     }
 
-    ucx_new_team->config = config;
-    ucx_new_team->parent_team = parent_team;
+    ucx_new_team->config = (mca_spml_ucx_team_config_t*)config;
+    ucx_new_team->parent_team = (mca_spml_ucx_team_t*)parent_team;
 
-    *new_team = ucx_new_team;
+    *new_team = (shmem_team_t)ucx_new_team;
 
     return OSHMEM_SUCCESS;
 }
@@ -1904,21 +1881,29 @@ int mca_spml_ucx_team_split_2d(shmem_team_t parent_team, int xrange, const
         shmem_team_t *yaxis_team)
 {
     mca_spml_ucx_team_t *ucx_parent_team = (mca_spml_ucx_team_t*) parent_team;
-    int yrange = ucx_parent_team->n_pes / xrange;
-    int pe_x = ucx_parent_team->my_pe % xrange;
-    int pe_y = ucx_parent_team->my_pe / xrange;
+    int parent_n_pes = (parent_team == SHMEM_TEAM_WORLD) ?
+                                        oshmem_num_procs() :
+                                        ucx_parent_team->n_pes;
+    int parent_my_pe = (parent_team == SHMEM_TEAM_WORLD) ?
+                                        shmem_my_pe() :
+                                        ucx_parent_team->my_pe;
+    int yrange = parent_n_pes / xrange;
+    int pe_x = parent_my_pe % xrange;
+    int pe_y = parent_my_pe / xrange;
     int rc;
 
     /* Create x-team of my_pe */
-    rc = mca_spml_ucx_team_split_strided(parent_team, pe_y * xrange, 1, xrange, xaxis_config, xaxis_mask, xaxis_team);
+    rc = mca_spml_ucx_team_split_strided(parent_team, pe_y * xrange, 1, xrange,
+                                         xaxis_config, xaxis_mask, xaxis_team);
 
     if (rc != OSHMEM_SUCCESS) {
         SPML_UCX_ERROR("mca_spml_ucx_team_split_strided failed (x-axis team creation)");
         return rc;
     }
 
     /* Create y-team of my_pe */
-    rc = mca_spml_ucx_team_split_strided(parent_team, pe_x, xrange, yrange, yaxis_config, yaxis_mask, yaxis_team);
+    rc = mca_spml_ucx_team_split_strided(parent_team, pe_x, xrange, yrange,
+                                         yaxis_config, yaxis_mask, yaxis_team);
     if (rc != OSHMEM_SUCCESS) {
         SPML_UCX_ERROR("mca_spml_ucx_team_split_strided failed (y-axis team creation)");
         goto out_free_xaxis;

oshmem/shmem/c/shmem_team.c

@@ -51,14 +51,9 @@ void shmem_team_sync(shmem_team_t team)
 
 int shmem_team_my_pe(shmem_team_t team)
 {
-    int rc = 0;
-
     RUNTIME_CHECK_INIT();
 
-    rc = MCA_SPML_CALL(team_my_pe(team));
-    RUNTIME_CHECK_IMPL_RC(rc);
-
-    return rc;
+    return MCA_SPML_CALL(team_my_pe(team));
 }
 
 int shmem_team_n_pes(shmem_team_t team)
@@ -85,15 +80,11 @@ int shmem_team_get_config(shmem_team_t team, long config_mask, shmem_team_config
 }
 int shmem_team_translate_pe(shmem_team_t src_team, int src_pe, shmem_team_t dest_team)
 {
-    int rc = 0;
-
     RUNTIME_CHECK_INIT();
 
-    rc = MCA_SPML_CALL(team_translate_pe(src_team, src_pe, dest_team));
-    RUNTIME_CHECK_IMPL_RC(rc);
-
-    return rc;
+    return MCA_SPML_CALL(team_translate_pe(src_team, src_pe, dest_team));
 }
+
 int shmem_team_split_strided (shmem_team_t parent_team, int start, int stride,
         int size, const shmem_team_config_t *config, long config_mask,
         shmem_team_t *new_team)