[Openmp-commits] [openmp] [OpenMP][libomp] Fix tasking debug assert (PR #95823)

[Jonathan Peyton via Openmp-commits]

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

The debug assert is meant to check that the index is a valid which means the runtime needs to check against the size of the array instead of the number of threads. A free()-ed thread put back in the thread pool may index into anywhere inside the task team's available array from 0 to tt_max_threads potentially.

Fixes: #94260

>From 59a9c298e9e3919a2b804ac1a231945e76d8bb70 Mon Sep 17 00:00:00 2001
From: Jonathan Peyton <jonathan.l.peyton at intel.com>
Date: Mon, 17 Jun 2024 11:01:29 -0500
Subject: [PATCH] [OpenMP][libomp] Fix tasking debug assert

The debug assert is meant to check that the index is a valid
which means the runtime needs to check against the size of the array
instead of the number of threads. A free()-ed thread put back in the
thread pool may index into anywhere inside the task team's available
array from 0 to tt_max_threads potentially.

Fixes: #94260
---
 openmp/runtime/src/kmp_tasking.cpp            |  2 +-
 openmp/runtime/test/tasking/issue-94260-1.cpp | 66 +++++++++++++++++++
 openmp/runtime/test/tasking/issue-94260-2.c   | 51 ++++++++++++++
 3 files changed, 118 insertions(+), 1 deletion(-)
 create mode 100644 openmp/runtime/test/tasking/issue-94260-1.cpp
 create mode 100644 openmp/runtime/test/tasking/issue-94260-2.c

diff --git a/openmp/runtime/src/kmp_tasking.cpp b/openmp/runtime/src/kmp_tasking.cpp
index 936afb91ac2f8..03ce0dd752af1 100644
--- a/openmp/runtime/src/kmp_tasking.cpp
+++ b/openmp/runtime/src/kmp_tasking.cpp
@@ -3245,7 +3245,7 @@ static kmp_task_t *__kmp_steal_task(kmp_int32 victim_tid, kmp_int32 gtid,
   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);
+  KMP_DEBUG_ASSERT(victim_tid < task_team->tt.tt_max_threads);
 
   victim_td = &threads_data[victim_tid];
   victim_thr = victim_td->td.td_thr;
diff --git a/openmp/runtime/test/tasking/issue-94260-1.cpp b/openmp/runtime/test/tasking/issue-94260-1.cpp
new file mode 100644
index 0000000000000..6f2bf2ae0a859
--- /dev/null
+++ b/openmp/runtime/test/tasking/issue-94260-1.cpp
@@ -0,0 +1,66 @@
+// RUN: %libomp-cxx-compile-and-run
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <omp.h>
+
+// The number of times to run each test
+#define NTIMES 2
+
+// Every thread creates a single "increment" task
+void test_tasks() {
+  for (int i = 0; i < 100; ++i)
+#pragma omp task
+  {
+    int tid = omp_get_thread_num();
+  }
+}
+
+// Testing single level of parallelism with increment tasks
+void test_base(int nthreads) {
+#ifdef VERBOSE
+#pragma omp master
+  printf("    test_base(%d)\n", nthreads);
+#endif
+#pragma omp parallel num_threads(nthreads)
+  { test_tasks(); }
+}
+
+// Testing nested parallel with increment tasks
+// first = nthreads of outer parallel
+// second = nthreads of nested parallel
+void test_nest(int first, int second) {
+#ifdef VERBOSE
+#pragma omp master
+  printf("   test_nest(%d, %d)\n", first, second);
+#endif
+#pragma omp parallel num_threads(first)
+  {
+    for (int i = 0; i < 100; ++i)
+#pragma omp task
+    {
+      int tid = omp_get_thread_num();
+    }
+    test_base(second);
+  }
+}
+
+template <typename... Args>
+void run_ntimes(int n, void (*func)(Args...), Args... args) {
+  for (int i = 0; i < n; ++i) {
+    func(args...);
+  }
+}
+
+int main() {
+  omp_set_max_active_levels(5);
+
+  for (int i = 0; i < 100; ++i) {
+    run_ntimes(NTIMES, test_nest, 4, 3);
+    run_ntimes(NTIMES, test_nest, 2, 1);
+  }
+
+  printf("PASS\n");
+  return EXIT_SUCCESS;
+}
diff --git a/openmp/runtime/test/tasking/issue-94260-2.c b/openmp/runtime/test/tasking/issue-94260-2.c
new file mode 100644
index 0000000000000..c2f43f91f70f9
--- /dev/null
+++ b/openmp/runtime/test/tasking/issue-94260-2.c
@@ -0,0 +1,51 @@
+// RUN: %libomp-compile-and-run
+
+#include <stdio.h>
+#include <omp.h>
+
+int test_omp_parallel_num_threads() {
+  int num_failed;
+  int threads;
+  int nthreads;
+  int max_threads = 0;
+
+  num_failed = 0;
+#pragma omp task
+  {}
+
+/* first we check how many threads are available */
+#pragma omp parallel
+  {
+#pragma omp task
+    {}
+#pragma omp master
+    max_threads = omp_get_num_threads();
+  }
+
+  /* we increase the number of threads from one to maximum:*/
+  for (threads = 1; threads <= max_threads; threads++) {
+    nthreads = 0;
+#pragma omp parallel reduction(+ : num_failed) num_threads(threads)
+    {
+#pragma omp task
+      {}
+      num_failed = num_failed + !(threads == omp_get_num_threads());
+#pragma omp atomic
+      nthreads += 1;
+    }
+    num_failed = num_failed + !(nthreads == threads);
+  }
+  return (!num_failed);
+}
+
+int main() {
+  int i;
+  int num_failed = 0;
+
+  for (i = 0; i < 100; i++) {
+    if (!test_omp_parallel_num_threads()) {
+      num_failed++;
+    }
+  }
+  return num_failed;
+}