[flang-commits] [flang] [flang][OpenACC] Skip non-loop evaluations when descending a collapsed/tiled DO nest (PR #223579)

[via flang-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: Ron Green [NVIDIA] (ronGreenNV)

<details>
<summary>Changes</summary>

## Summary
- The collapse/tile loop-nest descent in `Bridge.cpp` (used for OpenACC `collapse`/`tile` and CUDA Fortran `!$cuf kernel do`) assumed each loop level's nested evaluations have the fixed shape `{NonLabelDoStmt, DoConstruct, EndDoStmt}` and unconditionally advanced to the second nested evaluation to find the next inner `DO`.
- A compiler directive (e.g. `!DIR$ IVDEP`) placed between loop levels inserts an extra sibling evaluation there, so the descent landed on the directive instead of the inner `DoConstruct` and lowered the directive's (empty) nested evaluations as the loop body — silently dropping the real body with no diagnostic.
- Adds a helper that searches a level's nested evaluations for the actual `DoConstruct`, matching the approach already used by `visitLoopControl` in `OpenACC.cpp`, and uses it at all three affected descent sites (OpenACC collapse/tile, and the two CUDA Fortran kernel-do sites).

## Test plan
- New lit test `flang/test/Lower/OpenACC/acc-loop-collapse-directive-between-loops.f90` covering `collapse(2)` with `!DIR$ IVDEP` and `collapse(3)` with `!DIR$ NOVECTOR` between loop levels.
- Verified the new test fails without the fix (directive absorbed as loop body, no `hlfir.assign` for the real body) and passes with it.
- Existing `acc-loop-collapse-force-lowering.f90`, `acc-loop-collapse-force-non-tightly-nested.f90`, `do-concurrent-collapse.f90`, `cuda-kernel-loop-directive.cuf`, `cuda-kernel-do-reduction.cuf` all still pass.

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


2 Files Affected:

- (modified) flang/lib/Lower/Bridge.cpp (+31-6) 
- (added) flang/test/Lower/OpenACC/acc-loop-collapse-directive-between-loops.f90 (+55) 


``````````diff
diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index 4b70fbe9ea106..2b5bdeab6d118 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -502,6 +502,20 @@ class TypeInfoConverter {
 using IncrementLoopNestInfo = llvm::SmallVector<IncrementLoopInfo, 8>;
 } // namespace
 
+/// Find the first nested DoConstruct evaluation directly under \p eval,
+/// skipping over any other sibling evaluations (e.g. a CompilerDirective
+/// such as !DIR$ IVDEP) that may appear between loop levels of a collapsed
+/// or tiled loop nest. Returns nullptr if none is found.
+static Fortran::lower::pft::Evaluation *
+findNestedDoConstructEvaluation(Fortran::lower::pft::Evaluation &eval) {
+  if (!eval.hasNestedEvaluations())
+    return nullptr;
+  for (Fortran::lower::pft::Evaluation &child : eval.getNestedEvaluations())
+    if (child.getIf<Fortran::parser::DoConstruct>())
+      return &child;
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // FirConverter
 //===----------------------------------------------------------------------===//
@@ -3638,9 +3652,11 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         const auto *outerDo = curEval->getIf<Fortran::parser::DoConstruct>();
         if (!(outerDo && outerDo->IsDoConcurrent()))
           for (uint64_t i = 1; i < loopCount; i++) {
-            if (!curEval->hasNestedEvaluations())
+            Fortran::lower::pft::Evaluation *nextDo =
+                findNestedDoConstructEvaluation(*curEval);
+            if (!nextDo)
               break;
-            curEval = &*std::next(curEval->getNestedEvaluations().begin());
+            curEval = nextDo;
           }
       }
     }
@@ -3983,8 +3999,12 @@ class FirConverter : public Fortran::lower::AbstractConverter {
 
         ivTypes.push_back(idxTy);
         ivLocs.push_back(crtLoc);
-        if (i < nestedLoops - 1)
-          loopEval = &*std::next(loopEval->getNestedEvaluations().begin());
+        if (i < nestedLoops - 1) {
+          Fortran::lower::pft::Evaluation *nextDo =
+              findNestedDoConstructEvaluation(*loopEval);
+          assert(nextDo && "expected a nested DO CONSTRUCT");
+          loopEval = nextDo;
+        }
       }
     }
 
@@ -4012,8 +4032,13 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     if (crtEval->lowerAsStructured()) {
       crtEval = &crtEval->getFirstNestedEvaluation();
       if (!outerDoConstruct->IsDoConcurrent())
-        for (int64_t i = 1; i < nestedLoops; i++)
-          crtEval = &*std::next(crtEval->getNestedEvaluations().begin());
+        for (int64_t i = 1; i < nestedLoops; i++) {
+          Fortran::lower::pft::Evaluation *nextDo =
+              findNestedDoConstructEvaluation(*crtEval);
+          if (!nextDo)
+            break;
+          crtEval = nextDo;
+        }
     }
 
     // Generate loop body
diff --git a/flang/test/Lower/OpenACC/acc-loop-collapse-directive-between-loops.f90 b/flang/test/Lower/OpenACC/acc-loop-collapse-directive-between-loops.f90
new file mode 100644
index 0000000000000..89f551ae2492f
--- /dev/null
+++ b/flang/test/Lower/OpenACC/acc-loop-collapse-directive-between-loops.f90
@@ -0,0 +1,55 @@
+! RUN: bbc -fopenacc -emit-hlfir %s -o - | FileCheck %s
+
+! Verify that a compiler directive (e.g. !DIR$ IVDEP) appearing between the
+! levels of a collapsed loop nest does not get mistaken for the next nested
+! DO CONSTRUCT. The directive is an extra sibling evaluation between the
+! NonLabelDoStmt and the inner DoConstruct; the collapse descent must skip
+! over it rather than absorbing it as the loop body.
+
+subroutine collapse2_directive_between_loops(n, a)
+  integer, intent(in) :: n
+  integer :: a(n,n)
+  integer :: i, j
+
+  !$acc parallel loop collapse(2) copy(a)
+  do i = 1, n
+!DIR$ IVDEP
+    do j = 1, n
+      a(j,i) = 1
+    end do
+  end do
+  !$acc end parallel loop
+end subroutine
+
+! CHECK-LABEL: func.func @_QPcollapse2_directive_between_loops(
+! CHECK: acc.parallel
+! CHECK: acc.loop combined(parallel)
+! CHECK: hlfir.designate
+! CHECK: hlfir.assign
+! CHECK: acc.yield
+! CHECK: collapse([2])
+
+subroutine collapse3_directive_between_loops(n, a)
+  integer, intent(in) :: n
+  integer :: a(n,n,n)
+  integer :: i, j, k
+
+  !$acc parallel loop collapse(3) copy(a)
+  do i = 1, n
+!DIR$ NOVECTOR
+    do j = 1, n
+      do k = 1, n
+        a(k,j,i) = 1
+      end do
+    end do
+  end do
+  !$acc end parallel loop
+end subroutine
+
+! CHECK-LABEL: func.func @_QPcollapse3_directive_between_loops(
+! CHECK: acc.parallel
+! CHECK: acc.loop combined(parallel)
+! CHECK: hlfir.designate
+! CHECK: hlfir.assign
+! CHECK: acc.yield
+! CHECK: collapse([3])

``````````

</details>


https://github.com/llvm/llvm-project/pull/223579