teams: preallocate num_contexts if config_mask set

Sandia-OpenSHMEM · Apr 30, 2024 · dad6a11 · dad6a11
1 parent aa1e722
commit dad6a11
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/shmem_team.c b/src/shmem_team.c
@@ -338,19 +338,21 @@ int shmem_internal_team_split_strided(shmem_internal_team_t *parent_team, int PE
         if (config_mask == 0) {
             shmem_team_config_t default_config;
             default_config.num_contexts = 0;
+            myteam->contexts_len = 0;
             myteam->config_mask = 0;
             memcpy(&myteam->config, &default_config, sizeof(shmem_team_config_t));
         } else {
             if (config_mask != SHMEM_TEAM_NUM_CONTEXTS) {
                 RAISE_WARN_MSG("Invalid team_split_strided config_mask (%ld)\n", config_mask);
                 return -1;
             } else {
-                myteam->config      = *config;
-                myteam->config_mask = config_mask;
+                myteam->config       = *config;
+                myteam->config_mask  = config_mask;
+                myteam->contexts_len = config->num_contexts;
+                myteam->contexts     = malloc(config->num_contexts * sizeof(shmem_transport_ctx_t*));
             }
         }
 
-        myteam->contexts_len = 0;
         myteam->psync_idx = -1;
 
         shmem_internal_op_to_all(psync_pool_avail_reduced,