Fix an issue with th_task_state_memo_stack and proxy/helper tasks

[tlwilmar]

Summary

When proxy or helper tasks were used in inactive parallel regions, no memo of the th_task_state was stored in the stack, so th_task_state became invalid. This change inserts an item in the memo stack to track these th_task_states.

Patch by Alex Duran.

[tlwilmar]

Ping.

[tianshilei1992]

LG

openmp/runtime/src/kmp_tasking.cpp

if (done) {
       break;
     }
 
     // If oversubscribed or have waited a bit, yield.
     KMP_YIELD_OVERSUB_ELSE_SPIN(spins, time);
   }
 }
 
+void __kmp_shift_task_state_stack(kmp_info_t *this_thr, kmp_uint8 value) {
+  // Shift values from th_task_state_top+1 to task_state_stack_sz
+  if (this_thr->th.th_task_state_top + 1 >=
+      this_thr->th.th_task_state_stack_sz) { // increase size
+    kmp_uint32 new_size = 2 * this_thr->th.th_task_state_stack_sz;
+    kmp_uint8 *old_stack, *new_stack;
+    kmp_uint32 i;
+    new_stack = (kmp_uint8 *)__kmp_allocate(new_size);
+    for (i = 0; i <= this_thr->th.th_task_state_top; ++i) {
+      new_stack[i] = this_thr->th.th_task_state_memo_stack[i];
+    }
+    // If we need to reallocate do the shift at the same time.
+    for (; i < this_thr->th.th_task_state_stack_sz; ++i) {
+      new_stack[i + 1] = this_thr->th.th_task_state_memo_stack[i];
+    }
+    for (i = this_thr->th.th_task_state_stack_sz; i < new_size;
+         ++i) { // zero-init rest of stack
+      new_stack[i] = 0;
+    }
+    old_stack = this_thr->th.th_task_state_memo_stack;
+    this_thr->th.th_task_state_memo_stack = new_stack;
+    this_thr->th.th_task_state_stack_sz = new_size;
+    __kmp_free(old_stack);
+  } else {
+    kmp_uint8 *end;
+    kmp_uint32 i;
+
+    end = &this_thr->th
+               .th_task_state_memo_stack[this_thr->th.th_task_state_stack_sz];
+
+    for (i = this_thr->th.th_task_state_stack_sz - 1;
+         i > this_thr->th.th_task_state_top; i--, end--)
+      end[0] = end[-1];
+  }
+  this_thr->th.th_task_state_memo_stack[this_thr->th.th_task_state_top + 1] =
+      value;
+}
+
 // __kmp_task_team_setup:  Create a task_team for the current team, but use
 // an already created, unused one if it already exists.
 void __kmp_task_team_setup(kmp_info_t *this_thr, kmp_team_t *team, int always) {
   KMP_DEBUG_ASSERT(__kmp_tasking_mode != tskm_immediate_exec);
 
   // If this task_team hasn't been created yet, allocate it. It will be used in
   // the region after the next.
   // If it exists, it is the current task team and shouldn't be touched yet as
   // it may still be in use.
   if (team->t.t_task_team[this_thr->th.th_task_state] == NULL &&
       (always || team->t.t_nproc > 1)) {
     team->t.t_task_team[this_thr->th.th_task_state] =
         __kmp_allocate_task_team(this_thr, team);
     KA_TRACE(20, ("__kmp_task_team_setup: Primary T#%d created new task_team %p"
                   " for team %d at parity=%d\n",
                   __kmp_gtid_from_thread(this_thr),
                   team->t.t_task_team[this_thr->th.th_task_state], team->t.t_id,
                   this_thr->th.th_task_state));
   }
+  if (this_thr->th.th_task_state == 1 && always && team->t.t_nproc == 1) {
+    // fix task state stack to adjust for proxy and helper tasks
+    KA_TRACE(20, ("__kmp_task_team_setup: Primary T#%d needs to shift stack"
+                  " for team %d at parity=%d\n",
+                  __kmp_gtid_from_thread(this_thr), team->t.t_id,
+                  this_thr->th.th_task_state));
+    __kmp_shift_task_state_stack(this_thr, this_thr->th.th_task_state);
+  }
 
   // After threads exit the release, they will call sync, and then point to this
   // other task_team; make sure it is allocated and properly initialized. As
   // threads spin in the barrier release phase, they will continue to use the
   // previous task_team struct(above), until they receive the signal to stop
   // checking for tasks (they can't safely reference the kmp_team_t struct,
   // which could be reallocated by the primary thread). No task teams are formed
   // for serialized teams.