[flang-commits] [flang] [Flang][OpenMP] Fix semantics check for nested DISTRIBUTE (PR #91592)

[Sergio Afonso via flang-commits]

https://github.com/skatrak created https://github.com/llvm/llvm-project/pull/91592

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.

>From 7fed43dd11b4cfaa98812213242ec6eca3c39d03 Mon Sep 17 00:00:00 2001
From: Sergio Afonso <safonsof at amd.com>
Date: Thu, 9 May 2024 13:47:33 +0100
Subject: [PATCH] [Flang][OpenMP] Fix semantics check for nested DISTRIBUTE

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.
---
 flang/lib/Semantics/check-omp-structure.cpp       | 3 ++-
 flang/test/Semantics/OpenMP/nested-distribute.f90 | 7 +++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

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