[flang] Allow compiler directives for collapsed loops in OpenACC

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.

Reviewed By: clementval

Differential Revision: https://reviews.llvm.org/D132846

flang/lib/Semantics/resolve-directives.cpp

@@ -868,6 +868,21 @@ void AccAttributeVisitor::PrivatizeAssociatedLoopIndex(
   }
   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
@@ -879,8 +894,7 @@ void AccAttributeVisitor::PrivatizeAssociatedLoopIndex(
     }
 
     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);
 }

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

@@ -0,0 +1,21 @@
+! 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