Fix initialization of th_task_state on each thread on expanding hot teams.

[tlwilmar]

Summary

The th_task_state was initialized from the master thread's value, or from its memo stack, but this causes problems because neither of those may have the right value at the right time. However, other threads in the team are guaranteed to have the right values, so we change the initialize the new threads' th_task_state from the th_task_state of the last of the older threads in the hot team.

Fix #56307.

[jlpeyton]

LGTM

[rogfer01]

Hi @tlwilmar could you elaborate a bit on this part as it is very unobvious at first

However, other threads in the team are guaranteed to have the right values, so we change the initialize the new threads' th_task_state from the th_task_state of the last of the older threads in the hot team.

Thanks!

[tlwilmar]

@rogfer01 th_task_state is used to refer to one of two task teams that are reused for each hot team. These two task teams exist because of a need to start using a new one while there may still be activity on the old task team, so we flip between the two. Each thread has a th_task_state so it knows which task team is in use for a particular parallel region. When a hot team is expanded in size for a larger parallel region, new threads added to the team need to have their th_task_state initialized to exactly what the other threads in the team are using, so that they are all in sync. So instead, we just initialize from a thread that was in the team before the resize.

[rogfer01]

In D142247#4115990, @tlwilmar wrote:

@rogfer01 th_task_state is used to refer to one of two task teams that are reused for each hot team. These two task teams exist because of a need to start using a new one while there may still be activity on the old task team, so we flip between the two. Each thread has a th_task_state so it knows which task team is in use for a particular parallel region. When a hot team is expanded in size for a larger parallel region, new threads added to the team need to have their th_task_state initialized to exactly what the other threads in the team are using, so that they are all in sync. So instead, we just initialize from a thread that was in the team before the resize.

Thanks for the explanation @tlwilmar!

openmp/runtime/src/kmp_runtime.cpp

#endif // KMP_NESTED_HOT_TEAMS
                             root->r.r_uber_thread->th.th_ident);
 
       /* reinitialize the threads */
       KMP_DEBUG_ASSERT(team->t.t_nproc == new_nproc);
       for (f = 0; f < team->t.t_nproc; ++f)
         __kmp_initialize_info(team->t.t_threads[f], team, f,
                               __kmp_gtid_from_tid(f, team));
 
-      if (level) { // set th_task_state for new threads in nested hot team
-        // __kmp_initialize_info() no longer zeroes th_task_state, so we should
-        // only need to set the th_task_state for the new threads. th_task_state
-        // for primary thread will not be accurate until after this in
-        // __kmp_fork_call(), so we look to the primary thread's memo_stack to
-        // get the correct value.
-        for (f = old_nproc; f < team->t.t_nproc; ++f)
-          team->t.t_threads[f]->th.th_task_state =
-              team->t.t_threads[0]->th.th_task_state_memo_stack[level];
-      } else { // set th_task_state for new threads in non-nested hot team
-        // copy primary thread's state
-        kmp_uint8 old_state = team->t.t_threads[0]->th.th_task_state;
-        for (f = old_nproc; f < team->t.t_nproc; ++f)
-          team->t.t_threads[f]->th.th_task_state = old_state;
-      }
+      // set th_task_state for new threads in hot team with older thread's state
+      kmp_uint8 old_state = team->t.t_threads[old_nproc - 1]->th.th_task_state;
 
 #ifdef KMP_DEBUG
       for (f = 0; f < team->t.t_nproc; ++f) {
         KMP_DEBUG_ASSERT(team->t.t_threads[f] &&
                          team->t.t_threads[f]->th.th_team_nproc ==
                              team->t.t_nproc);
       }
 #endif