[Openmp-commits] [llvm] [openmp] [OpenMP] Publish target-init team state with acq_rel barrier (PR #208701)

[via Openmp-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-offload

Author: Abid Qadeer (abidh)

<details>
<summary>Changes</summary>

In __kmpc_target_init, thread 0 initializes the team state in LDS (HasThreadState = 0, ThreadStates = nullptr) and then releases the other threads with synchronize::threadsAligned(atomic::relaxed). Because the relaxed ordering emits no memory fence, the initializing LDS stores are not guaranteed to be visible to the worker threads before they proceed. A worker thread can then read the uninitialized team state and produce wrong results or crash. The fix is to Use atomic::acq_rel so that the memory fences are emitted.

Reproducing the problem requires an `if` clause with a non-constant condition on the target region. The worker's read of HasThreadState originates from the serialized-region check in __kmpc_parallel_60. When the `if` condition is constant (or absent), the optimizer can prove that branch is dead and eliminates the HasThreadState load entirely, so the race has nothing to observe. A non-constant `if` keeps if_expr, and hence the LDS read alive, exposing the problem.

I have added a Fortran offload regression test that launches a target region many times to make the timing-dependent race likely to surface.

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

Assisted-by: Cursor

---
Full diff: https://github.com/llvm/llvm-project/pull/208701.diff


2 Files Affected:

- (added) offload/test/offloading/fortran/target-teams-if-runtime-init.f90 (+61) 
- (modified) openmp/device/src/Kernel.cpp (+1-1) 


``````````diff
diff --git a/offload/test/offloading/fortran/target-teams-if-runtime-init.f90 b/offload/test/offloading/fortran/target-teams-if-runtime-init.f90
new file mode 100644
index 0000000000000..edbd5fdea1f8e
--- /dev/null
+++ b/offload/test/offloading/fortran/target-teams-if-runtime-init.f90
@@ -0,0 +1,61 @@
+! Regression test for a race at kernel start-up.
+
+! On entry to a target region, thread 0 initializes the team state in LDS
+! (HasThreadState = 0, ThreadStates = nullptr) inside __kmpc_target_init and
+! then releases the other threads with a barrier. That barrier was emitted with
+! `atomic::relaxed`, so the initializing LDS stores are not guaranteed to be
+! visible to the worker threads before they proceed. A worker could then read
+! uninitialized team state, which showed up as wrong results (or a crash).
+!
+! The race is timing dependent, so the proven reproducer pattern is repeated
+! over a large iteration space to make a bad ordering likely to surface. A
+! correct runtime passes every launch deterministically.
+
+! REQUIRES: flang, amdgpu
+
+! RUN: %libomptarget-compile-fortran-generic
+! RUN: %libomptarget-run-generic 2>&1 | %fcheck-generic
+
+module control_mod
+  implicit none
+  ! Not a parameter on purpose: a runtime `if` value is required to triggered
+  ! the failure.
+  logical :: should_offload = .true.
+end module
+
+program target_teams_if_runtime_init
+  use control_mod
+  implicit none
+  integer, parameter :: n = 30000
+  integer, parameter :: nlaunch = 50
+  integer :: data(n)
+  integer :: i, launch, errors
+
+  errors = 0
+
+  do launch = 1, nlaunch
+    data = 0
+
+    !$omp target enter data map(alloc: data)
+
+    !$omp target teams distribute parallel do map(present, tofrom: data) &
+    !$omp&   if(should_offload)
+    do i = 1, n
+      data(i) = i
+    end do
+    !$omp end target teams distribute parallel do
+
+    !$omp target exit data map(from: data)
+
+    do i = 1, n
+      if (data(i) /= i) then
+        errors = errors + 1
+      end if
+    end do
+  end do
+
+  print *, "number of errors: ", errors
+
+end program target_teams_if_runtime_init
+
+! CHECK: number of errors: 0
diff --git a/openmp/device/src/Kernel.cpp b/openmp/device/src/Kernel.cpp
index d6b8659436156..600aa0c8528b7 100644
--- a/openmp/device/src/Kernel.cpp
+++ b/openmp/device/src/Kernel.cpp
@@ -112,7 +112,7 @@ int32_t __kmpc_target_init(KernelEnvironmentTy &KernelEnvironment,
   if (IsSPMD) {
     initializeRuntime(/*IsSPMD=*/true, KernelEnvironment,
                       KernelLaunchEnvironment);
-    synchronize::threadsAligned(atomic::relaxed);
+    synchronize::threadsAligned(atomic::acq_rel);
   } else {
     initializeRuntime(/*IsSPMD=*/false, KernelEnvironment,
                       KernelLaunchEnvironment);

``````````

</details>


https://github.com/llvm/llvm-project/pull/208701