[Openmp-commits] [openmp] Draft: [OpenMP] Fix td_tdg_task_id underflow with taskloop and taskgraph (PR #150877)

[Josep Pinot via Openmp-commits]

https://github.com/jpinot updated https://github.com/llvm/llvm-project/pull/150877

>From c53fcbe5527d1aee3aca96eaf823d9092f8669c3 Mon Sep 17 00:00:00 2001
From: jpinot <josep.pinot at bsc.es>
Date: Thu, 17 Jul 2025 12:22:01 +0200
Subject: [PATCH 1/2] [OpenMP] Fix td_tdg_task_id underflow with taskloop and
 taskgraph

This patch addresses an issue where the td_tdg_task_id could underflow,
leading to a negative task ID, when a taskloop region was encountered
before a taskgraph clause.
---
 openmp/runtime/src/kmp_tasking.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/openmp/runtime/src/kmp_tasking.cpp b/openmp/runtime/src/kmp_tasking.cpp
index e4d92a78fd6b9..f4e5107627b73 100644
--- a/openmp/runtime/src/kmp_tasking.cpp
+++ b/openmp/runtime/src/kmp_tasking.cpp
@@ -4952,7 +4952,8 @@ static void __kmp_taskloop(ident_t *loc, int gtid, kmp_task_t *task, int if_val,
   }
 
 #if OMPX_TASKGRAPH
-  KMP_ATOMIC_DEC(&__kmp_tdg_task_id);
+  if (taskdata->is_taskgraph)
+    KMP_ATOMIC_DEC(&__kmp_tdg_task_id);
 #endif
   // =========================================================================
   // calculate loop parameters

>From 81ef84f6f7b64391883048cf826c2d60173e3bd0 Mon Sep 17 00:00:00 2001
From: jpinot <josep.pinot at bsc.es>
Date: Wed, 30 Jul 2025 09:57:56 +0200
Subject: [PATCH 2/2] [wip][openmp] Move _kmp_tdg_task_id inside kmp_tdg_info

---
 openmp/runtime/src/kmp.h           |  1 +
 openmp/runtime/src/kmp_tasking.cpp | 14 +++++++-------
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/openmp/runtime/src/kmp.h b/openmp/runtime/src/kmp.h
index 818edf9060ad1..96c8544c2a994 100644
--- a/openmp/runtime/src/kmp.h
+++ b/openmp/runtime/src/kmp.h
@@ -2667,6 +2667,7 @@ typedef struct kmp_tdg_info {
   kmp_tdg_status_t tdg_status =
       KMP_TDG_NONE; // Status of the TDG (recording, ready...)
   std::atomic<kmp_int32> num_tasks; // Number of TDG nodes
+  std::atomic<kmp_int32> __kmp_tdg_task_id_next; // Task id of next node
   kmp_bootstrap_lock_t
       graph_lock; // Protect graph attributes when updated via taskloop_recur
   // Taskloop reduction related
diff --git a/openmp/runtime/src/kmp_tasking.cpp b/openmp/runtime/src/kmp_tasking.cpp
index f4e5107627b73..b1b79d635fb91 100644
--- a/openmp/runtime/src/kmp_tasking.cpp
+++ b/openmp/runtime/src/kmp_tasking.cpp
@@ -1430,7 +1430,7 @@ kmp_task_t *__kmp_task_alloc(ident_t *loc_ref, kmp_int32 gtid,
     taskdata->is_taskgraph = 1;
     taskdata->tdg = __kmp_global_tdgs[__kmp_curr_tdg_idx];
     taskdata->td_task_id = KMP_GEN_TASK_ID();
-    taskdata->td_tdg_task_id = KMP_ATOMIC_INC(&__kmp_tdg_task_id);
+    taskdata->td_tdg_task_id = KMP_ATOMIC_INC(&tdg->__kmp_tdg_task_id_next);
   }
 #endif
   KA_TRACE(20, ("__kmp_task_alloc(exit): T#%d created task %p parent=%p\n",
@@ -4436,7 +4436,7 @@ kmp_task_t *__kmp_task_dup_alloc(kmp_info_t *thread, kmp_task_t *task_src
 #if OMPX_TASKGRAPH
   if (taskdata->is_taskgraph && !taskloop_recur &&
       __kmp_tdg_is_recording(taskdata_src->tdg->tdg_status))
-    taskdata->td_tdg_task_id = KMP_ATOMIC_INC(&__kmp_tdg_task_id);
+    taskdata->td_tdg_task_id = KMP_ATOMIC_INC(&taskdata_src->tdg->__kmp_tdg_task_id_next);
 #endif
   taskdata->td_task_id = KMP_GEN_TASK_ID();
   if (task->shareds != NULL) { // need setup shareds pointer
@@ -4951,10 +4951,10 @@ static void __kmp_taskloop(ident_t *loc, int gtid, kmp_task_t *task, int if_val,
     __kmpc_taskgroup(loc, gtid);
   }
 
-#if OMPX_TASKGRAPH
-  if (taskdata->is_taskgraph)
-    KMP_ATOMIC_DEC(&__kmp_tdg_task_id);
-#endif
+/* #if OMPX_TASKGRAPH */
+/*   if (taskdata->is_taskgraph) */
+/*     KMP_ATOMIC_DEC(&__kmp_tdg_task_id); */
+/* #endif */
   // =========================================================================
   // calculate loop parameters
   kmp_taskloop_bounds_t task_bounds(task, lb, ub);
@@ -5425,7 +5425,7 @@ void __kmp_end_record(kmp_int32 gtid, kmp_tdg_info_t *tdg) {
     KMP_ATOMIC_ST_RLX(&this_record_map[i].npredecessors_counter,
                       this_record_map[i].npredecessors);
   }
-  KMP_ATOMIC_ST_RLX(&__kmp_tdg_task_id, 0);
+  /* KMP_ATOMIC_ST_RLX(&__kmp_tdg_task_id, 0); */
 
   if (__kmp_tdg_dot)
     __kmp_print_tdg_dot(tdg, gtid);