From 41f148e61d0b5a13f2406c4acf6e36ab24172a9d Mon Sep 17 00:00:00 2001
From: Alex Duran <alejandro.duran@intel.com>
Date: Thu, 9 Mar 2023 17:46:24 -0600
Subject: [PATCH] Fix an issue with th_task_state_memo_stack and proxy/helper
 tasks

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.

Differential Revision: https://reviews.llvm.org/D145736
---
 openmp/runtime/src/kmp_tasking.cpp | 46 ++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/openmp/runtime/src/kmp_tasking.cpp b/openmp/runtime/src/kmp_tasking.cpp
index e8c565c6ae918..71922ed708001 100644
--- a/openmp/runtime/src/kmp_tasking.cpp
+++ b/openmp/runtime/src/kmp_tasking.cpp
@@ -3934,6 +3934,44 @@ void __kmp_wait_to_unref_task_teams(void) {
   }
 }
 
+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) {
@@ -3953,6 +3991,14 @@ void __kmp_task_team_setup(kmp_info_t *this_thr, kmp_team_t *team, int always) {
                   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