[flang][openacc] Support labeled DO loop after acc loop directive

[clementval]

Make the DoConstruct in OpenACCLoopConstruct optional and move the labeled do construct in in the canonicalization step.

[llvmbot]

@llvm/pr-subscribers-flang-parser

Changes

Make the `DoConstruct` in `OpenACCLoopConstruct` optional and move the labeled do construct in in the canonicalization step. -- Full diff: https://github.com//pull/66294.diff

7 Files Affected:

• (modified) flang/include/flang/Parser/parse-tree.h (+5-1)
• (modified) flang/lib/Parser/openacc-parsers.cpp (+1-3)
• (modified) flang/lib/Parser/unparse.cpp (+1-1)
• (modified) flang/lib/Semantics/canonicalize-acc.cpp (+26-10)
• (modified) flang/lib/Semantics/resolve-directives.cpp (+2-2)
• (modified) flang/test/Lower/OpenACC/acc-loop.f90 (+6)
• (modified) flang/test/Semantics/OpenACC/acc-loop-validity.f90 (+6-2)

diff --git a/flang/include/flang/Parser/parse-tree.h b/flang/include/flang/Parser/parse-tree.h
index 5d92ecb05841767..cb4bb59bf312c7c 100644
--- a/flang/include/flang/Parser/parse-tree.h
+++ b/flang/include/flang/Parser/parse-tree.h
@@ -4247,7 +4247,11 @@ struct OpenACCDeclarativeConstruct {
 EMPTY_CLASS(AccEndLoop);
 struct OpenACCLoopConstruct {
   TUPLE_CLASS_BOILERPLATE(OpenACCLoopConstruct);
-  std::tuple<AccBeginLoopDirective, DoConstruct, std::optional<AccEndLoop>> t;
+  OpenACCLoopConstruct(AccBeginLoopDirective &&a)
+      : t({std::move(a), std::nullopt, std::nullopt}) {}
+  std::tuple<AccBeginLoopDirective, std::optional<DoConstruct>,
+      std::optional<AccEndLoop>>
+      t;
 };
 
 struct OpenACCStandaloneConstruct {
diff --git a/flang/lib/Parser/openacc-parsers.cpp b/flang/lib/Parser/openacc-parsers.cpp
index afa12b88019bd9a..e1cbe6b86df91e9 100644
--- a/flang/lib/Parser/openacc-parsers.cpp
+++ b/flang/lib/Parser/openacc-parsers.cpp
@@ -154,9 +154,7 @@ TYPE_PARSER(construct<AccEndLoop>(startAccLine >> "END LOOP"_tok))
 
 TYPE_PARSER(construct<OpenACCLoopConstruct>(
     sourced(Parser<AccBeginLoopDirective>{} / endAccLine),
-    withMessage("A DO loop must follow the loop construct"_err_en_US,
-        Parser<DoConstruct>{}),
-    maybe(Parser<AccEndLoop>{} / endAccLine)))
+    maybe(Parser<DoConstruct>{}), maybe(Parser<AccEndLoop>{} / endAccLine)))
 
 // 2.15.1 Routine directive
 TYPE_PARSER(sourced(construct<OpenACCRoutineConstruct>(verbatim("ROUTINE"_tok),
diff --git a/flang/lib/Parser/unparse.cpp b/flang/lib/Parser/unparse.cpp
index 398545d315e5cd5..6cac3fb5859f87b 100644
--- a/flang/lib/Parser/unparse.cpp
+++ b/flang/lib/Parser/unparse.cpp
@@ -1947,7 +1947,7 @@ class UnparseVisitor {
     Walk(std::get<AccBeginLoopDirective>(x.t));
     Put("\n");
     EndOpenACC();
-    Walk(std::get<DoConstruct>(x.t));
+    Walk(std::get<std::optional<DoConstruct>>(x.t));
   }
   void Unparse(const AccBeginLoopDirective &x) {
     Walk(std::get<AccLoopDirective>(x.t));
diff --git a/flang/lib/Semantics/canonicalize-acc.cpp b/flang/lib/Semantics/canonicalize-acc.cpp
index e79ab997637b083..3887e479cbce5d5 100644
--- a/flang/lib/Semantics/canonicalize-acc.cpp
+++ b/flang/lib/Semantics/canonicalize-acc.cpp
@@ -109,21 +109,37 @@ class CanonicalizationOfAcc {
 
   void RewriteOpenACCLoopConstruct(parser::OpenACCLoopConstruct &x,
       parser::Block &block, parser::Block::iterator it) {
+    parser::Block::iterator nextIt;
     auto &beginDir{std::get<parser::AccBeginLoopDirective>(x.t)};
     auto &dir{std::get<parser::AccLoopDirective>(beginDir.t)};
-    const auto &doCons{std::get<parser::DoConstruct>(x.t)};
+    auto &nestedDo{std::get<std::optional<parser::DoConstruct>>(x.t)};
+
+    if (!nestedDo) {
+      nextIt = it;
+      if (++nextIt != block.end()) {
+        if (auto *doCons{parser::Unwrap<parser::DoConstruct>(*nextIt)}) {
+          nestedDo = std::move(*doCons);
+          nextIt = block.erase(nextIt);
+        }
+      }
+    }
 
-    if (!doCons.GetLoopControl()) {
-      messages_.Say(dir.source,
-          "DO loop after the %s directive must have loop control"_err_en_US,
-          parser::ToUpperCaseLetters(dir.source.ToString()));
+    if (nestedDo) {
+      if (!nestedDo->GetLoopControl()) {
+        messages_.Say(dir.source,
+            "DO loop after the %s directive must have loop control"_err_en_US,
+            parser::ToUpperCaseLetters(dir.source.ToString()));
+        return;
+      }
+      CheckDoConcurrentClauseRestriction<parser::OpenACCLoopConstruct,
+          parser::AccBeginLoopDirective>(x, *nestedDo);
+      CheckTileClauseRestriction<parser::OpenACCLoopConstruct,
+          parser::AccBeginLoopDirective>(x, *nestedDo);
       return;
     }
-
-    CheckDoConcurrentClauseRestriction<parser::OpenACCLoopConstruct,
-        parser::AccBeginLoopDirective>(x, doCons);
-    CheckTileClauseRestriction<parser::OpenACCLoopConstruct,
-        parser::AccBeginLoopDirective>(x, doCons);
+    messages_.Say(dir.source,
+        "A DO loop must follow the %s directive"_err_en_US,
+        parser::ToUpperCaseLetters(dir.source.ToString()));
   }
 
   void RewriteOpenACCCombinedConstruct(parser::OpenACCCombinedConstruct &x,
diff --git a/flang/lib/Semantics/resolve-directives.cpp b/flang/lib/Semantics/resolve-directives.cpp
index 9a2689c2e6c81d5..63a3e1cd31549c0 100644
--- a/flang/lib/Semantics/resolve-directives.cpp
+++ b/flang/lib/Semantics/resolve-directives.cpp
@@ -1153,8 +1153,8 @@ void AccAttributeVisitor::PrivatizeAssociatedLoopIndex(
     return nullptr;
   };
 
-  const auto &outer{std::get<parser::DoConstruct>(x.t)};
-  for (const parser::DoConstruct *loop{&outer}; loop && level > 0; --level) {
+  const auto &outer{std::get<std::optional<parser::DoConstruct>>(x.t)};
+  for (const parser::DoConstruct *loop{&*outer}; loop && level > 0; --level) {
     // go through all the nested do-loops and resolve index variables
     const parser::Name *iv{GetLoopIndex(*loop)};
     if (iv) {
diff --git a/flang/test/Lower/OpenACC/acc-loop.f90 b/flang/test/Lower/OpenACC/acc-loop.f90
index f902728b67100f3..ec5eff5da78e7ed 100644
--- a/flang/test/Lower/OpenACC/acc-loop.f90
+++ b/flang/test/Lower/OpenACC/acc-loop.f90
@@ -315,4 +315,10 @@ program acc_loop
 ! CHECK: %[[CACHE:.*]] = acc.cache varPtr(%{{.*}} : !fir.ref<!fir.array<10xf32>>) bounds(%{{.*}}) -> !fir.ref<!fir.array<10xf32>> {name = "b"}
 ! CHECK: acc.loop cache(%[[CACHE]] : !fir.ref<!fir.array<10xf32>>)
 
+  !$acc loop
+  do 100 i=0, n
+  100 continue
+! CHECK: acc.loop
+! CHECK: fir.do_loop
+
 end program
diff --git a/flang/test/Semantics/OpenACC/acc-loop-validity.f90 b/flang/test/Semantics/OpenACC/acc-loop-validity.f90
index 906602640efacc7..205e67a4728fcbe 100644
--- a/flang/test/Semantics/OpenACC/acc-loop-validity.f90
+++ b/flang/test/Semantics/OpenACC/acc-loop-validity.f90
@@ -4,12 +4,16 @@ program openacc_clause_validity
 
   implicit none
 
-  integer :: i
+  integer :: i, n
 
   i = 0
 
+  !ERROR: A DO loop must follow the LOOP directive
   !$acc loop
-  !ERROR: A DO loop must follow the loop construct
   i = 1
 
+  !$acc loop
+  do 100 i=0, n
+  100 continue
+
 end

[vzakhari]

LGTM. Thank you!