[flang][OpenMP] Skip compiler directives in getCollapsedLoopEval

flang/lib/Lower/OpenMP/OpenMP.cpp

@@ -25,6 +25,7 @@
 #include "flang/Lower/ConvertVariable.h"
 #include "flang/Lower/DirectivesCommon.h"
 #include "flang/Lower/OpenMP/Clauses.h"
+#include "flang/Lower/PFTBuilder.h"
 #include "flang/Lower/StatementContext.h"
 #include "flang/Lower/Support/ReductionProcessor.h"
 #include "flang/Lower/SymbolMap.h"
@@ -568,14 +569,9 @@ getCollapsedLoopEval(lower::pft::Evaluation &eval, int collapseValue) {
   if (collapseValue == 0)
     return &eval;
 
-  lower::pft::Evaluation *curEval = &eval.getFirstNestedEvaluation();
-  for (int i = 1; i < collapseValue; i++) {
-    // The nested evaluations should be DoConstructs (i.e. they should form
-    // a loop nest). Each DoConstruct is a tuple <NonLabelDoStmt, Block,
-    // EndDoStmt>.
-    assert(curEval->isA<parser::DoConstruct>());
-    curEval = &*std::next(curEval->getNestedEvaluations().begin());
-  }
+  lower::pft::Evaluation *curEval = &eval;
+  for (int i = 0; i < collapseValue; i++)
+    curEval = getNestedDoConstruct(*curEval);
   return curEval;
 }
 

flang/lib/Lower/OpenMP/Utils.cpp

@@ -796,7 +796,7 @@ static void processTileSizesFromOpenMPConstruct(
   }
 }
 
-static pft::Evaluation *getNestedDoConstruct(pft::Evaluation &eval) {
+pft::Evaluation *getNestedDoConstruct(pft::Evaluation &eval) {
   for (pft::Evaluation &nested : eval.getNestedEvaluations()) {
     // In an OpenMPConstruct there can be compiler directives:
     // 1 <<OpenMPConstruct>>

flang/lib/Lower/OpenMP/Utils.h

@@ -167,6 +167,8 @@ void genObjectList(const ObjectList &objects,
 void lastprivateModifierNotSupported(const omp::clause::Lastprivate &lastp,
                                      mlir::Location loc);
 
+pft::Evaluation *getNestedDoConstruct(pft::Evaluation &eval);
+
 int64_t collectLoopRelatedInfo(
     lower::AbstractConverter &converter, mlir::Location currentLocation,
     lower::pft::Evaluation &eval, const omp::List<omp::Clause> &clauses,

flang/test/Lower/OpenMP/compiler-directives-loop.f90

@@ -0,0 +1,31 @@
+!RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=60 %s -o - | FileCheck %s
+
+! Check that we generate proper body of the do-construct.
+
+!CHECK: omp.loop_nest (%[[ARG1:arg[0-9]+]]) : i32 = (%c1_i32) to (%c10_i32) inclusive step (%c1_i32_1) {
+!CHECK:   %[[V0:[0-9]+]]:2 = hlfir.declare %arg0 {uniq_name = "_QFEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+!CHECK:   hlfir.assign %[[ARG1]] to %[[V0]]#0 : i32, !fir.ref<i32>
+!CHECK:   %[[V1:[0-9]+]] = fir.load %[[V0]]#0 : !fir.ref<i32>
+!CHECK:   %[[V2:[0-9]+]] = fir.convert %[[V1]] : (i32) -> f32
+!CHECK:   %[[V3:[0-9]+]] = fir.load %[[V0]]#0 : !fir.ref<i32>
+!CHECK:   %[[V4:[0-9]+]] = fir.convert %[[V3]] : (i32) -> i64
+!CHECK:   %[[V5:[0-9]+]] = hlfir.designate %3#0 (%[[V4]])  : (!fir.ref<!fir.array<10xf32>>, i64) -> !fir.ref<f32>
+!CHECK:   hlfir.assign %[[V2]] to %[[V5]] : f32, !fir.ref<f32>
+!CHECK:   omp.yield
+!CHECK: }
+
+program omp_cdir_codegen
+  implicit none
+  integer, parameter :: n = 10
+  real :: a(n)
+  integer :: i
+
+!$omp parallel do
+!dir$ unroll
+  do i = 1, n
+    a(i) = real(i)
+  end do
+!$omp end parallel do
+
+  print *, 'a(1)=', a(1), ' a(n)=', a(n)
+end program omp_cdir_codegen