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

[Sergio Afonso via Openmp-commits]

https://github.com/skatrak created https://github.com/llvm/llvm-project/pull/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
```

>From 4aeb31f4ee305a56ae6eb2fb71c0c49fb7cfac36 Mon Sep 17 00:00:00 2001
From: Sergio Afonso <safonsof at amd.com>
Date: Wed, 5 Mar 2025 16:56:22 +0000
Subject: [PATCH] [OpenMP] Address __kmp_dist_for_static_init issue

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 previous 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
```
---
 openmp/runtime/src/kmp_sched.cpp | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/openmp/runtime/src/kmp_sched.cpp b/openmp/runtime/src/kmp_sched.cpp
index 2e0dfac6eeb3b..c93a0eab25c53 100644
--- a/openmp/runtime/src/kmp_sched.cpp
+++ b/openmp/runtime/src/kmp_sched.cpp
@@ -542,6 +542,10 @@ 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 being called after
+  // __kmp_serialized_parallel from using the wrong teams information
+  while (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);