[Openmp-commits] [openmp] 6085f3f - [OpenMP] Address __kmp_dist_for_static_init issue (#129902)

[via Openmp-commits]

Author: Sergio Afonso
Date: 2025-03-17T11:44:29Z
New Revision: 6085f3f6a80dc3be97d31a4c31a22fd77dbb6e83

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

LOG: [OpenMP] Address __kmp_dist_for_static_init issue (#129902)

This patch attempts to provide a fix for an issue that appears when the
`__kmp_dist_for_static_init` function is called from a serialized team.

This is triggered by code generated by flang for `distribute parallel
do` constructs whenever an `if` clause for the `parallel` leaf construct
is present. This results in the introduction of a call to
`__kmpc_fork_call_if` in place of `__kmpc_fork_call`. When it evaluates
to `false`, it defers execution to `__kmp_serialized_parallel`, which
creates a new serial team that is picked up by
`__kmp_dist_for_static_init`, resulting in an incorrect `team` pointer
that causes the `nteams == (kmp_uint32)team->t.t_parent->t.t_nproc`
assertion to fail.

The sequence of calls replicating this issue can be summarized as:
  - `__kmpc_fork_teams`
  - `__kmpc_fork_call_if`
  - `__kmpc_dist_for_static_init_*`

Since I am not familiar with the implementation of the OpenMP runtime,
it is possible that the above sequence of calls is incorrect, or that
the bug can be better fixed in another way, so I am open to discussing
this.

The following Fortran program can be compiled with flang to show the
issue:

```f90
! Compile and run: flang -fopenmp test.f90 -o test && ./test
! Check LLVM IR: flang -fc1 -emit-llvm -fopenmp test.f90 -o -

program main
  implicit none
  integer, parameter :: n = 10
  integer :: i, idx(n)

  !$omp teams
  !$omp distribute parallel do if(.false.)
  do i=1,n
    idx(i) = i
  end do
  !$omp end teams

  print *, idx
end program
```

Added: 
    

Modified: 
    openmp/runtime/src/kmp_sched.cpp

Removed: 
    


################################################################################
diff  --git a/openmp/runtime/src/kmp_sched.cpp b/openmp/runtime/src/kmp_sched.cpp
index 2e0dfac6eeb3b..2b1bb6f595f9a 100644
--- a/openmp/runtime/src/kmp_sched.cpp
+++ b/openmp/runtime/src/kmp_sched.cpp
@@ -542,6 +542,12 @@ static void __kmp_dist_for_static_init(ident_t *loc, kmp_int32 gtid,
   nth = th->th.th_team_nproc;
   team = th->th.th_team;
   KMP_DEBUG_ASSERT(th->th.th_teams_microtask); // we are in the teams construct
+  // skip optional serialized teams to prevent this from using the wrong teams
+  // information when called after __kmp_serialized_parallel
+  // TODO: make __kmp_serialized_parallel eventually call __kmp_fork_in_teams
+  // to address this edge case
+  while (team->t.t_parent && team->t.t_serialized)
+    team = team->t.t_parent;
   nteams = th->th.th_teams_size.nteams;
   team_id = team->t.t_master_tid;
   KMP_DEBUG_ASSERT(nteams == (kmp_uint32)team->t.t_parent->t.t_nproc);