[Openmp-commits] [openmp] [OpenMP] Have hidden helper team allocate new OS threads only (PR #87119)

[Jonathan Peyton via Openmp-commits]

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

The hidden helper team pre-allocates the gtid space [1, num_hidden_helpers] (inclusive). If regular host threads are allocated, then put back in the thread pool, then the hidden helper team is initialized, the hidden helper team tries to allocate the threads from the thread pool with gtids higher than [1, num_hidden_helpers]. Instead, have the hidden helper team fork OS threads so the correct gtid range used for hidden helper threads.

Fixes: #87117

>From e32f5c18316959cac4035e2987af39d1233cc5e4 Mon Sep 17 00:00:00 2001
From: Jonathan Peyton <jonathan.l.peyton at intel.com>
Date: Fri, 29 Mar 2024 17:14:10 -0500
Subject: [PATCH] [OpenMP] Have hidden helper team allocate new OS threads only

The hidden helper team pre-allocates the gtid space [1,
num_hidden_helpers] (inclusive). If regular host threads are allocated,
then put back in the thread pool, then the hidden helper team is
initialized, the hidden helper team tries to allocate the threads from
the thread pool with gtids higher than [1, num_hidden_helpers]. Instead,
have the hidden helper team fork OS threads so the correct gtid range
used for hidden helper threads.

Fixes: #87117
---
 openmp/runtime/src/kmp_runtime.cpp            |  8 +++--
 .../tasking/hidden_helper_task/issue-87117.c  | 36 +++++++++++++++++++
 2 files changed, 41 insertions(+), 3 deletions(-)
 create mode 100644 openmp/runtime/test/tasking/hidden_helper_task/issue-87117.c

diff --git a/openmp/runtime/src/kmp_runtime.cpp b/openmp/runtime/src/kmp_runtime.cpp
index a60bdb968371e0..9219d17d511a5c 100644
--- a/openmp/runtime/src/kmp_runtime.cpp
+++ b/openmp/runtime/src/kmp_runtime.cpp
@@ -4431,8 +4431,10 @@ kmp_info_t *__kmp_allocate_thread(kmp_root_t *root, kmp_team_t *team,
 #endif
   KMP_MB();
 
-  /* first, try to get one from the thread pool */
-  if (__kmp_thread_pool) {
+  /* first, try to get one from the thread pool unless allocating thread is
+   * the main hidden helper thread. The hidden helper team should always
+   * allocate new OS threads. */
+  if (__kmp_thread_pool && !KMP_HIDDEN_HELPER_TEAM(team)) {
     new_thr = CCAST(kmp_info_t *, __kmp_thread_pool);
     __kmp_thread_pool = (volatile kmp_info_t *)new_thr->th.th_next_pool;
     if (new_thr == __kmp_thread_pool_insert_pt) {
@@ -4497,7 +4499,7 @@ kmp_info_t *__kmp_allocate_thread(kmp_root_t *root, kmp_team_t *team,
   }
 
   /* no, well fork a new one */
-  KMP_ASSERT(__kmp_nth == __kmp_all_nth);
+  KMP_ASSERT(KMP_HIDDEN_HELPER_TEAM(team) || __kmp_nth == __kmp_all_nth);
   KMP_ASSERT(__kmp_all_nth < __kmp_threads_capacity);
 
 #if KMP_USE_MONITOR
diff --git a/openmp/runtime/test/tasking/hidden_helper_task/issue-87117.c b/openmp/runtime/test/tasking/hidden_helper_task/issue-87117.c
new file mode 100644
index 00000000000000..23080982f49e19
--- /dev/null
+++ b/openmp/runtime/test/tasking/hidden_helper_task/issue-87117.c
@@ -0,0 +1,36 @@
+// RUN: %libomp-compile
+// RUN: env KMP_HOT_TEAMS_MODE=0 KMP_HOT_TEAMS_MAX_LEVEL=1 %libomp-run
+//
+// Force the defaults of:
+// KMP_HOT_TEAMS_MODE=0 means free extra threads after parallel
+//   involving non-hot team
+// KMP_HOT_TEAMS_MAX_LEVEL=1 means only the initial outer team
+//   is a hot team.
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <omp.h>
+
+int main() {
+  int a;
+  omp_set_max_active_levels(2);
+// This nested parallel creates extra threads on the thread pool
+#pragma omp parallel num_threads(2)
+  {
+#pragma omp parallel num_threads(2)
+    {
+#pragma omp atomic
+      a++;
+    }
+  }
+
+// Causes assert if hidden helper thread tries to allocate from thread pool
+// instead of creating new OS threads
+#pragma omp parallel num_threads(1)
+  {
+#pragma omp target nowait
+    { a++; }
+  }
+
+  return EXIT_SUCCESS;
+}