[Openmp-commits] [openmp] [OpenMP] Fix node destruction race in __kmpc_omp_taskwait_deps_51 (PR #86130)

[Ulrich Weigand via Openmp-commits]

https://github.com/uweigand created https://github.com/llvm/llvm-project/pull/86130

The __kmpc_omp_taskwait_deps_51 allocates a kmp_depnode_t node on its stack, and there is currently a race condition where another thread might still be accessing that node after the function has returned and its stack frame was released.

While the function does wait until the node's npredecessors count has reached zero before exiting, there is still a window where the function that last decremented the npredecessors count assumes the node is still accessible.

For heap-allocated kmp_depnode_t nodes, this normally works via a separate ndeps count that only reaches zero at the point where no accesses to the node are expected at all; in fact, at this point the heap allocation will be freed.

For this case of a stack-allocated kmp_depnode_t node, it therefore makes sense to similarly respect the ndeps count; we need to wait until this reaches 1 (not 0, because it is not heap-allocated so there's always one extra count to prevent it from being freed), before we can safely deallocate our stack frame.

As this is expected to be a short race window of only a few instructions, it should be fine to just use a busy wait loop checking the ndeps count.

Fixes: https://github.com/llvm/llvm-project/issues/85963

>From c8954a27f68fef79a5a4cc3e149598fa5db7d546 Mon Sep 17 00:00:00 2001
From: Ulrich Weigand <ulrich.weigand at de.ibm.com>
Date: Thu, 21 Mar 2024 15:56:24 +0100
Subject: [PATCH] [OpenMP] Fix node destruction race in
 __kmpc_omp_taskwait_deps_51

The __kmpc_omp_taskwait_deps_51 allocates a kmp_depnode_t node on
its stack, and there is currently a race condition where another
thread might still be accessing that node after the function has
returned and its stack frame was released.

While the function does wait until the node's npredecessors count
has reached zero before exiting, there is still a window where the
function that last decremented the npredecessors count assumes the
node is still accessible.

For heap-allocated kmp_depnode_t nodes, this normally works via a
separate ndeps count that only reaches zero at the point where no
accesses to the node are expected at all; in fact, at this point
the heap allocation will be freed.

For this case of a stack-allocated kmp_depnode_t node, it therefore
makes sense to similarly respect the ndeps count; we need to wait
until this reaches 1 (not 0, because it is not heap-allocated so
there's always one extra count to prevent it from being freed),
before we can safely deallocate our stack frame.

As this is expected to be a short race window of only a few
instructions, it should be fine to just use a busy wait loop
checking the ndeps count.

Fixes: https://github.com/llvm/llvm-project/issues/85963
---
 openmp/runtime/src/kmp_taskdeps.cpp | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/openmp/runtime/src/kmp_taskdeps.cpp b/openmp/runtime/src/kmp_taskdeps.cpp
index f7529481393f97..e575ad8b08a55f 100644
--- a/openmp/runtime/src/kmp_taskdeps.cpp
+++ b/openmp/runtime/src/kmp_taskdeps.cpp
@@ -1030,6 +1030,12 @@ void __kmpc_omp_taskwait_deps_51(ident_t *loc_ref, kmp_int32 gtid,
                        __kmp_task_stealing_constraint);
   }
 
+  // Wait until the last __kmp_release_deps is finished before we free the
+  // current stack frame holding the "node" variable; once its nrefs count
+  // reaches 1, we're sure nobody else can try to reference it again.
+  while (node.dn.nrefs > 1)
+    KMP_YIELD(TRUE);
+
 #if OMPT_SUPPORT
   __ompt_taskwait_dep_finish(current_task, taskwait_task_data);
 #endif /* OMPT_SUPPORT */