[flang][openacc] Avoid crash when collapse loop nest has extra directive

[clementval]

The compiler was crashing when the collapse loop nest could not be retrieved because of extra acc loop directive inside it.

[llvmbot]

@llvm/pr-subscribers-flang-openmp
@llvm/pr-subscribers-flang-semantics

@llvm/pr-subscribers-openacc

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

---

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

2 Files Affected:

• (modified) flang/lib/Semantics/resolve-directives.cpp (+14-3)
• (modified) flang/test/Semantics/OpenACC/acc-loop.f90 (+8)

diff --git a/flang/lib/Semantics/resolve-directives.cpp b/flang/lib/Semantics/resolve-directives.cpp
index bbb105e3516da18..45329b5cca9c03b 100644
--- a/flang/lib/Semantics/resolve-directives.cpp
+++ b/flang/lib/Semantics/resolve-directives.cpp
@@ -1183,13 +1183,23 @@ void AccAttributeVisitor::CheckAssociatedLoopIndex(
   }
 
   const auto getNextDoConstruct =
-      [this](const parser::Block &block) -> const parser::DoConstruct * {
+      [this](const parser::Block &block,
+          std::int64_t &level) -> const parser::DoConstruct * {
     for (const auto &entry : block) {
       if (const auto *doConstruct = GetDoConstructIf(entry)) {
         return doConstruct;
       } else if (parser::Unwrap<parser::CompilerDirective>(entry)) {
         // It is allowed to have a compiler directive associated with the loop.
         continue;
+      } else if (const auto &accLoop{
+                     parser::Unwrap<parser::OpenACCLoopConstruct>(entry)}) {
+        if (level == 0)
+          break;
+        const auto &beginDir{
+            std::get<parser::AccBeginLoopDirective>(accLoop->t)};
+        context_.Say(beginDir.source,
+            "LOOP directive not expected in COLLAPSE loop nest"_err_en_US);
+        level = 0;
       } else {
         break;
       }
@@ -1198,11 +1208,12 @@ void AccAttributeVisitor::CheckAssociatedLoopIndex(
   };
 
   const auto &outer{std::get<std::optional<parser::DoConstruct>>(x.t)};
-  for (const parser::DoConstruct *loop{&*outer}; loop && level > 0; --level) {
+  for (const parser::DoConstruct *loop{&*outer}; loop && level > 0;) {
     // Go through all nested loops to ensure index variable exists.
     GetLoopIndex(*loop);
     const auto &block{std::get<parser::Block>(loop->t)};
-    loop = getNextDoConstruct(block);
+    --level;
+    loop = getNextDoConstruct(block, level);
   }
   CHECK(level == 0);
 }
diff --git a/flang/test/Semantics/OpenACC/acc-loop.f90 b/flang/test/Semantics/OpenACC/acc-loop.f90
index 1c0a515441f432b..ee1e6e0cce09740 100644
--- a/flang/test/Semantics/OpenACC/acc-loop.f90
+++ b/flang/test/Semantics/OpenACC/acc-loop.f90
@@ -268,4 +268,12 @@ program openacc_loop_validity
   end do
   !$acc end loop
 
+  !$acc loop collapse(2)
+  do i = 1, 10
+    !ERROR: LOOP directive not expected in COLLAPSE loop nest
+    !$acc loop
+    do j = 1, 10
+    end do
+  end do
+
 end program openacc_loop_validity

[razvanlupusoru]

Generating an error for this case seems correct to me. This is because the clauses on acc loop are supposed to apply to all contained loops - and it would be ambiguous if multiple acc loop were contained.

Thanks Valentin!

[vzakhari]

LGTM. Thank you!