[flang] Allow compiler directives for collapsed loops in OpenACC

[unterumarmung]

Summary

If one tries to compile the attached test case code with flan,
the one will get an internal compiler error on CHECK(level == 0)
at the end of PrivatizeAssociatedLoopIndex function.
Other compilers (gfortran and nvfortran) build this code just fine.
This change fixes the ICE.

[clementval]

LGTM.

flang/lib/Semantics/resolve-directives.cpp

877

[unterumarmung]

Thank you for the quick review!

flang/lib/Semantics/resolve-directives.cpp

 void AccAttributeVisitor::PrivatizeAssociatedLoopIndex(
     const parser::OpenACCLoopConstruct &x) {
   std::int64_t level{GetContext().associatedLoopLevel};
   if (level <= 0) { // collpase value was negative or 0
     return;
   }
   Symbol::Flag ivDSA{Symbol::Flag::AccPrivate};
 
+  const auto getNextDoConstruct =
+      [this](const parser::Block &block) -> 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 {
+        break;
+      }
+    }
+    return nullptr;
+  };
+
   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 (auto *symbol{ResolveAcc(iv, ivDSA, currScope())}) {
       symbol->set(Symbol::Flag::AccPreDetermined);
       iv.symbol = symbol; // adjust the symbol within region
       AddToContextObjectWithDSA(*symbol, ivDSA);
     }
 
     const auto &block{std::get<parser::Block>(loop->t)};
-    const auto it{block.begin()};
-    loop = it != block.end() ? GetDoConstructIf(*it) : nullptr;
+    loop = getNextDoConstruct(block);
   }
   CHECK(level == 0);
 }
 
 void AccAttributeVisitor::EnsureAllocatableOrPointer(
     const llvm::acc::Clause clause, const parser::AccObjectList &objectList) {
   for (const auto &accObject : objectList.v) {
     common::visit(

flang/test/Semantics/OpenACC/acc-resolve03.f90

+! RUN: %flang_fc1 -fopenacc %s
+! A regression test to check that
+! arbitrary compiler directives do not generate errors
+! inside OpenACC collapsed loops
+subroutine foo
+  integer, parameter :: loop_bound = 42
+  integer :: a
+  integer :: b
+  integer :: c
+
+  !$acc parallel
+  do a = 0, loop_bound
+    !$acc loop collapse(2)
+    do b = 0, loop_bound
+      !dir$ ivdep
+      do c = 0, loop_bound
+      enddo
+    enddo
+  enddo
+  !$acc end parallel
+end subroutine foo