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

[via Openmp-commits]

Author: Ulrich Weigand
Date: 2024-03-28T12:15:39+01:00
New Revision: b999e631c03640f7d1d93b5319da496ed4e0df55

URL: https://github.com/llvm/llvm-project/commit/b999e631c03640f7d1d93b5319da496ed4e0df55
DIFF: https://github.com/llvm/llvm-project/commit/b999e631c03640f7d1d93b5319da496ed4e0df55.diff

LOG: [OpenMP] Fix node destruction race in __kmpc_omp_taskwait_deps_51 (#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

Added: 
    

Modified: 
    openmp/runtime/src/kmp_taskdeps.cpp

Removed: 
    


################################################################################
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 */