[Openmp-commits] [openmp] [OpenMP] Fix nested parallel with tasking (PR #87309)

[Jonathan Peyton via Openmp-commits]

https://github.com/jpeyton52 created https://github.com/llvm/llvm-project/pull/87309

When a nested parallel region ends, the runtime calls __kmp_join_call(). During this call, the primary thread of the nested parallel region will reset its tid (retval of omp_get_thread_num()) to what it was in the outer parallel region. A data race occurs with the current code when another worker thread from the nested inner parallel region tries to steal tasks from the primary thread's task deque. The worker thread reads the tid value directly from the primary thread's data structure and may read the wrong value.

This change just uses the calculated victim_tid from execute_tasks() directly in the steal_task() routine rather than reading tid from the data structure.

Fixes: #87307

>From dd581d02e977bdd0cc0f4c0a7dcaf71601d32d93 Mon Sep 17 00:00:00 2001
From: Jonathan Peyton <jonathan.l.peyton at intel.com>
Date: Mon, 1 Apr 2024 22:14:49 -0500
Subject: [PATCH] [OpenMP] Fix nested parallel with tasking

When a nested parallel region ends it calls __kmp_join_call(). During
this call, the primary thread of the nested parallel region will reset
its tid (retval of omp_get_thread_num()) to what it was in the outer
parallel region. A data race occurs with the current code when another
worker thread from the nested parallel region tries to steal tasks from
the primary thread's task deque. The worker thread reads the tid value
directly from the primary thread's data structure and may read the wrong
value.

Fixes: #87307
---
 openmp/runtime/src/kmp_tasking.cpp        | 15 ++++----
 openmp/runtime/test/tasking/issue-87307.c | 43 +++++++++++++++++++++++
 2 files changed, 52 insertions(+), 6 deletions(-)
 create mode 100644 openmp/runtime/test/tasking/issue-87307.c

diff --git a/openmp/runtime/src/kmp_tasking.cpp b/openmp/runtime/src/kmp_tasking.cpp
index 155e17ba7ec874..6303bb0d63f0f4 100644
--- a/openmp/runtime/src/kmp_tasking.cpp
+++ b/openmp/runtime/src/kmp_tasking.cpp
@@ -3219,7 +3219,7 @@ static kmp_task_t *__kmp_remove_my_task(kmp_info_t *thread, kmp_int32 gtid,
 // __kmp_steal_task: remove a task from another thread's deque
 // Assume that calling thread has already checked existence of
 // task_team thread_data before calling this routine.
-static kmp_task_t *__kmp_steal_task(kmp_info_t *victim_thr, kmp_int32 gtid,
+static kmp_task_t *__kmp_steal_task(kmp_int32 victim_tid, kmp_int32 gtid,
                                     kmp_task_team_t *task_team,
                                     std::atomic<kmp_int32> *unfinished_threads,
                                     int *thread_finished,
@@ -3229,15 +3229,18 @@ static kmp_task_t *__kmp_steal_task(kmp_info_t *victim_thr, kmp_int32 gtid,
   kmp_taskdata_t *current;
   kmp_thread_data_t *victim_td, *threads_data;
   kmp_int32 target;
-  kmp_int32 victim_tid;
+  kmp_info_t *victim_thr;
 
   KMP_DEBUG_ASSERT(__kmp_tasking_mode != tskm_immediate_exec);
 
   threads_data = task_team->tt.tt_threads_data;
   KMP_DEBUG_ASSERT(threads_data != NULL); // Caller should check this condition
+  KMP_DEBUG_ASSERT(victim_tid >= 0);
+  KMP_DEBUG_ASSERT(victim_tid < task_team->tt.tt_nproc);
 
-  victim_tid = victim_thr->th.th_info.ds.ds_tid;
   victim_td = &threads_data[victim_tid];
+  victim_thr = victim_td->td.td_thr;
+  (void)victim_thr; // Use in TRACE messages which aren't always enabled.
 
   KA_TRACE(10, ("__kmp_steal_task(enter): T#%d try to steal from T#%d: "
                 "task_team=%p ntasks=%d head=%u tail=%u\n",
@@ -3452,9 +3455,9 @@ static inline int __kmp_execute_tasks_template(
 
         if (!asleep) {
           // We have a victim to try to steal from
-          task = __kmp_steal_task(other_thread, gtid, task_team,
-                                  unfinished_threads, thread_finished,
-                                  is_constrained);
+          task =
+              __kmp_steal_task(victim_tid, gtid, task_team, unfinished_threads,
+                               thread_finished, is_constrained);
         }
         if (task != NULL) { // set last stolen to victim
           if (threads_data[tid].td.td_deque_last_stolen != victim_tid) {
diff --git a/openmp/runtime/test/tasking/issue-87307.c b/openmp/runtime/test/tasking/issue-87307.c
new file mode 100644
index 00000000000000..f889ae20329eb3
--- /dev/null
+++ b/openmp/runtime/test/tasking/issue-87307.c
@@ -0,0 +1,43 @@
+// RUN: %libomp-compile-and-run
+#include <stdio.h>
+#include <stdlib.h>
+#include <omp.h>
+
+int a;
+
+void inc_a() {
+#pragma omp task
+  {
+#pragma omp atomic
+    a++;
+  }
+}
+
+int main() {
+  int n;
+  int nth_outer;
+  omp_set_max_active_levels(2);
+  omp_set_dynamic(0);
+
+  for (n = 0; n < 200; ++n) {
+    a = 0;
+#pragma omp parallel num_threads(8)
+    {
+      if (omp_get_thread_num() == 0)
+        nth_outer = omp_get_num_threads();
+#pragma omp parallel num_threads(2)
+      {
+        int i;
+#pragma omp master
+        for (i = 0; i < 50; ++i)
+          inc_a();
+      }
+    }
+    if (a != nth_outer * 50) {
+      fprintf(stderr, "error: a (%d) != %d\n", a, nth_outer * 50);
+      return EXIT_FAILURE;
+    }
+  }
+
+  return EXIT_SUCCESS;
+}