[flang-commits] [flang] [flang][acc] Fix crash on collapse(force:N) with non-tightly nested loops (PR #191310)

[via flang-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

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

Author: khaki3

<details>
<summary>Changes</summary>

When collapse(force:N) is applied to non-tightly nested loops (loops
with code between them, CYCLE statements, etc.), the PFT evaluation
tree may not have N levels of nested DoConstruct evaluations.  Several
places in the OpenACC lowering called getNestedEvaluations() without
first checking hasNestedEvaluations(), causing assertion failures.

Add guards in:
- hasEarlyReturn(): check hasNestedEvaluations() before iterating
- createRegionOp(): guard the createEmptyRegionBlocks call
- Bridge.cpp collapse-force sinking: break instead of assert when
  inner DoConstruct is not found, guard getNestedEvaluations() calls

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


3 Files Affected:

- (modified) flang/lib/Lower/Bridge.cpp (+15-10) 
- (modified) flang/lib/Lower/OpenACC.cpp (+4-1) 
- (added) flang/test/Lower/OpenACC/acc-loop-collapse-force-non-tightly-nested.f90 (+39) 


``````````diff
diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index 63867fab57678..1e096fd5a1ead 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -3641,8 +3641,11 @@ class FirConverter : public Fortran::lower::AbstractConverter {
 
       if (curEval->lowerAsStructured()) {
         curEval = &curEval->getFirstNestedEvaluation();
-        for (uint64_t i = 1; i < loopCount; i++)
+        for (uint64_t i = 1; i < loopCount; i++) {
+          if (!curEval->hasNestedEvaluations())
+            break;
           curEval = &*std::next(curEval->getNestedEvaluations().begin());
+        }
       }
     }
 
@@ -3669,9 +3672,8 @@ class FirConverter : public Fortran::lower::AbstractConverter {
           }
           prologue.push_back(&*it);
         }
-        // Semantics guarantees collapseDepth does not exceed nest depth
-        // so childLoop must be found here.
-        assert(childLoop && "Expected inner DoConstruct for collapse");
+        if (!childLoop)
+          break;
         parent = childLoop;
         innermostLoopEval = childLoop;
       }
@@ -3692,16 +3694,19 @@ class FirConverter : public Fortran::lower::AbstractConverter {
       sink(prologue);
 
       // Lower innermost loop body, skipping sunk
-      for (Fortran::lower::pft::Evaluation &e :
-           innermostLoopEval->getNestedEvaluations())
-        if (!sunk.contains(&e))
-          genFIR(e);
+      if (innermostLoopEval && innermostLoopEval->hasNestedEvaluations())
+        for (Fortran::lower::pft::Evaluation &e :
+             innermostLoopEval->getNestedEvaluations())
+          if (!sunk.contains(&e))
+            genFIR(e);
 
       sink(epilogue);
     } else {
       // Normal lowering
-      for (Fortran::lower::pft::Evaluation &e : curEval->getNestedEvaluations())
-        genFIR(e);
+      if (curEval->hasNestedEvaluations())
+        for (Fortran::lower::pft::Evaluation &e :
+             curEval->getNestedEvaluations())
+          genFIR(e);
     }
     if (!isCacheConstruct)
       localSymbols.popScope();
diff --git a/flang/lib/Lower/OpenACC.cpp b/flang/lib/Lower/OpenACC.cpp
index 2154f38dca568..5148b085d4b92 100644
--- a/flang/lib/Lower/OpenACC.cpp
+++ b/flang/lib/Lower/OpenACC.cpp
@@ -1215,7 +1215,8 @@ createRegionOp(fir::FirOpBuilder &builder, mlir::Location loc,
 
   // If it is an unstructured region and is not the outer region of a combined
   // construct, create empty blocks for all evaluations.
-  if (eval.lowerAsUnstructured() && !outerCombined)
+  if (eval.lowerAsUnstructured() && !outerCombined &&
+      eval.hasNestedEvaluations())
     Fortran::lower::createEmptyRegionBlocks<mlir::acc::TerminatorOp,
                                             mlir::acc::YieldOp>(
         builder, eval.getNestedEvaluations());
@@ -1895,6 +1896,8 @@ buildACCLoopOp(Fortran::lower::AbstractConverter &converter,
 }
 
 static bool hasEarlyReturn(Fortran::lower::pft::Evaluation &eval) {
+  if (!eval.hasNestedEvaluations())
+    return false;
   bool hasReturnStmt = false;
   for (auto &e : eval.getNestedEvaluations()) {
     e.visit(Fortran::common::visitors{
diff --git a/flang/test/Lower/OpenACC/acc-loop-collapse-force-non-tightly-nested.f90 b/flang/test/Lower/OpenACC/acc-loop-collapse-force-non-tightly-nested.f90
new file mode 100644
index 0000000000000..9f21c1d61b91a
--- /dev/null
+++ b/flang/test/Lower/OpenACC/acc-loop-collapse-force-non-tightly-nested.f90
@@ -0,0 +1,39 @@
+! RUN: bbc -fopenacc -emit-hlfir %s -o - | FileCheck %s
+
+! Verify that collapse(force:4) on non-tightly nested loops (with code
+! between loops and CYCLE statements) does not crash the compiler.
+
+subroutine collapse_force_non_tight(a, sr, nq, np, desc)
+  implicit none
+  integer, intent(in) :: nq, np
+  integer, intent(in) :: desc(3)
+  real, intent(inout) :: a(np * nq)
+  real, intent(in) :: sr(max(np, nq))
+  integer :: j1, j2, jcol, j
+  integer :: i1, i2, irow, i
+
+!$acc parallel loop collapse(force:4) present(a, sr)
+  do j1 = 1, nq, desc(1)
+    do j2 = 1, desc(1)
+      jcol = j1 + j2 - 1
+      if (jcol > nq) cycle
+      j = j1 + j2 - 1
+
+      do i1 = 1, np, desc(2)
+        do i2 = 1, desc(2)
+          irow = i1 + i2 - 1
+          if (irow > np) cycle
+          i = i1 + i2 - 1
+
+          a(irow + (jcol - 1) * np) = a(irow + (jcol - 1) * np) * sr(i) * sr(j)
+        end do
+      end do
+    end do
+  end do
+!$acc end parallel loop
+end subroutine
+
+! CHECK: func.func @_QPcollapse_force_non_tight
+! CHECK: acc.parallel
+! CHECK: acc.loop
+! CHECK: collapse = [4]

``````````

</details>


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