[Flang][OpenMP] Fix semantics check for nested DISTRIBUTE

[skatrak]

Composite OpenMP constructs where DISTRIBUTE is the first leaf construct, as well as standalone DISTRIBUTE constructs, are allowed inside of TEAMS regions.

Before this patch, nesting a DISTRIBUTE construct inside of a combined TARGET TEAMS construct was disallowed, which it shouldn't be. Now both TEAMS and TARGET TEAMS constructs can be immediate parents of DISTRIBUTE constructs.

[llvmbot]

@llvm/pr-subscribers-flang-semantics

@llvm/pr-subscribers-flang-openmp

Author: Sergio Afonso (skatrak)

Changes

Composite OpenMP constructs where DISTRIBUTE is the first leaf construct, as well as standalone DISTRIBUTE constructs, are allowed inside of TEAMS regions.

Before this patch, nesting a DISTRIBUTE construct inside of a combined TARGET TEAMS construct was disallowed, which it shouldn't be. Now both TEAMS and TARGET TEAMS constructs can be immediate parents of DISTRIBUTE constructs.

---

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

2 Files Affected:

• (modified) flang/lib/Semantics/check-omp-structure.cpp (+2-1)
• (modified) flang/test/Semantics/OpenMP/nested-distribute.f90 (+7)

diff --git a/flang/lib/Semantics/check-omp-structure.cpp b/flang/lib/Semantics/check-omp-structure.cpp
index 2493eb3ed3676..e9637b7bb591f 100644
--- a/flang/lib/Semantics/check-omp-structure.cpp
+++ b/flang/lib/Semantics/check-omp-structure.cpp
@@ -280,7 +280,8 @@ void OmpStructureChecker::HasInvalidDistributeNesting(
       violation = true;
     } else {
       // `distribute` region has to be strictly nested inside `teams`
-      if (!llvm::omp::topTeamsSet.test(GetContextParent().directive)) {
+      if (!OmpDirectiveSet{llvm::omp::OMPD_teams, llvm::omp::OMPD_target_teams}
+               .test(GetContextParent().directive)) {
         violation = true;
       }
     }
diff --git a/flang/test/Semantics/OpenMP/nested-distribute.f90 b/flang/test/Semantics/OpenMP/nested-distribute.f90
index 5103790392897..ba8c3bf04b337 100644
--- a/flang/test/Semantics/OpenMP/nested-distribute.f90
+++ b/flang/test/Semantics/OpenMP/nested-distribute.f90
@@ -74,6 +74,13 @@ program main
    !$omp end distribute
   !$omp end teams
 
+  !$omp target teams
+    !$omp distribute
+    do i = 1, 10
+    end do
+    !$omp end distribute
+  !$omp end target teams
+
   !$omp teams 
       !ERROR: Only `DISTRIBUTE` or `PARALLEL` regions are allowed to be strictly nested inside `TEAMS` region.
       !$omp task

[kiranchandramohan]

LGTM.