[Openmp-commits] [openmp] [OpenMP] Introduce KMP_TASK_STEALING to allow disabling task stealing. (PR #76891)

[Ye Luo via Openmp-commits]

https://github.com/ye-luo created https://github.com/llvm/llvm-project/pull/76891

None

>From b443dbe2da8c3fd55b1f01ccd81e384ca1dad8e8 Mon Sep 17 00:00:00 2001
From: Ye Luo <yeluo at anl.gov>
Date: Wed, 3 Jan 2024 22:05:04 -0600
Subject: [PATCH] [OpenMP] Introduce KMP_TASK_STEALING to allow disabling task
 stealing.

---
 openmp/runtime/src/kmp.h            |  1 +
 openmp/runtime/src/kmp_global.cpp   |  1 +
 openmp/runtime/src/kmp_settings.cpp | 25 ++++++++++++++++++++-----
 openmp/runtime/src/kmp_tasking.cpp  |  3 ++-
 4 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/openmp/runtime/src/kmp.h b/openmp/runtime/src/kmp.h
index 3dbf8c71c48d38..39e6f1fc6f92f3 100644
--- a/openmp/runtime/src/kmp.h
+++ b/openmp/runtime/src/kmp.h
@@ -2410,6 +2410,7 @@ typedef enum kmp_tasking_mode {
 
 extern kmp_tasking_mode_t
     __kmp_tasking_mode; /* determines how/when to execute tasks */
+extern int __kmp_task_stealing;
 extern int __kmp_task_stealing_constraint;
 extern int __kmp_enable_task_throttling;
 extern kmp_int32 __kmp_default_device; // Set via OMP_DEFAULT_DEVICE if
diff --git a/openmp/runtime/src/kmp_global.cpp b/openmp/runtime/src/kmp_global.cpp
index b132f38fd3b084..1df9ff370194a5 100644
--- a/openmp/runtime/src/kmp_global.cpp
+++ b/openmp/runtime/src/kmp_global.cpp
@@ -352,6 +352,7 @@ omp_memspace_handle_t const llvm_omp_target_device_mem_space =
    be redefined to have exactly 32 bits. */
 KMP_BUILD_ASSERT(sizeof(kmp_tasking_flags_t) == 4);
 
+int __kmp_task_stealing = 1; /* Allow task stealing by default */
 int __kmp_task_stealing_constraint = 1; /* Constrain task stealing by default */
 int __kmp_enable_task_throttling = 1;
 
diff --git a/openmp/runtime/src/kmp_settings.cpp b/openmp/runtime/src/kmp_settings.cpp
index e731bf45e8eee1..f276397caf087a 100644
--- a/openmp/runtime/src/kmp_settings.cpp
+++ b/openmp/runtime/src/kmp_settings.cpp
@@ -1358,15 +1358,27 @@ static void __kmp_stg_print_tasking(kmp_str_buf_t *buffer, char const *name,
 
 static void __kmp_stg_parse_task_stealing(char const *name, char const *value,
                                           void *data) {
-  __kmp_stg_parse_int(name, value, 0, 1,
-                      (int *)&__kmp_task_stealing_constraint);
+  __kmp_stg_parse_int(name, value, 0, 1, (int *)&__kmp_task_stealing);
 } // __kmp_stg_parse_task_stealing
 
 static void __kmp_stg_print_task_stealing(kmp_str_buf_t *buffer,
                                           char const *name, void *data) {
-  __kmp_stg_print_int(buffer, name, __kmp_task_stealing_constraint);
+  __kmp_stg_print_int(buffer, name, __kmp_task_stealing);
 } // __kmp_stg_print_task_stealing
 
+static void __kmp_stg_parse_task_stealing_constraint(char const *name,
+                                                     char const *value,
+                                                     void *data) {
+  __kmp_stg_parse_int(name, value, 0, 1,
+                      (int *)&__kmp_task_stealing_constraint);
+} // __kmp_stg_parse_task_stealing_constraint
+
+static void __kmp_stg_print_task_stealing_constraint(kmp_str_buf_t *buffer,
+                                                     char const *name,
+                                                     void *data) {
+  __kmp_stg_print_int(buffer, name, __kmp_task_stealing_constraint);
+} // __kmp_stg_print_task_stealing_constraint
+
 static void __kmp_stg_parse_max_active_levels(char const *name,
                                               char const *value, void *data) {
   kmp_uint64 tmp_dflt = 0;
@@ -5529,8 +5541,10 @@ static kmp_setting_t __kmp_stg_table[] = {
 
     {"KMP_TASKING", __kmp_stg_parse_tasking, __kmp_stg_print_tasking, NULL, 0,
      0},
-    {"KMP_TASK_STEALING_CONSTRAINT", __kmp_stg_parse_task_stealing,
+    {"KMP_TASK_STEALING", __kmp_stg_parse_task_stealing,
      __kmp_stg_print_task_stealing, NULL, 0, 0},
+    {"KMP_TASK_STEALING_CONSTRAINT", __kmp_stg_parse_task_stealing_constraint,
+     __kmp_stg_print_task_stealing_constraint, NULL, 0, 0},
     {"OMP_MAX_ACTIVE_LEVELS", __kmp_stg_parse_max_active_levels,
      __kmp_stg_print_max_active_levels, NULL, 0, 0},
     {"OMP_DEFAULT_DEVICE", __kmp_stg_parse_default_device,
@@ -5742,7 +5756,8 @@ static kmp_setting_t __kmp_stg_table[] = {
 #if OMPX_TASKGRAPH
     {"KMP_MAX_TDGS", __kmp_stg_parse_max_tdgs, __kmp_std_print_max_tdgs, NULL,
      0, 0},
-    {"KMP_TDG_DOT", __kmp_stg_parse_tdg_dot, __kmp_stg_print_tdg_dot, NULL, 0, 0},
+    {"KMP_TDG_DOT", __kmp_stg_parse_tdg_dot, __kmp_stg_print_tdg_dot, NULL, 0,
+     0},
 #endif
 
 #if OMPT_SUPPORT
diff --git a/openmp/runtime/src/kmp_tasking.cpp b/openmp/runtime/src/kmp_tasking.cpp
index 6e8b948efa064f..d0174b706d9ff5 100644
--- a/openmp/runtime/src/kmp_tasking.cpp
+++ b/openmp/runtime/src/kmp_tasking.cpp
@@ -3401,7 +3401,8 @@ static inline int __kmp_execute_tasks_template(
       if (task == NULL && use_own_tasks) { // check own queue next
         task = __kmp_remove_my_task(thread, gtid, task_team, is_constrained);
       }
-      if ((task == NULL) && (nthreads > 1)) { // Steal a task finally
+      if (__kmp_task_stealing && (task == NULL) &&
+          (nthreads > 1)) { // Steal a task finally
         int asleep = 1;
         use_own_tasks = 0;
         // Try to steal from the last place I stole from successfully.