[flang][OpenMP] Handle unstructured CF in compound loop constructs

[ergawy]

Fixes a bug in handling unstructured control-flow in compound loop constructs. The fix makes sure that unstructured CF does not get lowered until we reach the last item of the compound construct. This way, we avoid moving block of unstructured loops in-between the middle items of the construct and messing (i.e. adding operations) to these block while doing so.

[llvmbot]

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

@llvm/pr-subscribers-flang-openmp

Author: Kareem Ergawy (ergawy)

Changes

Fixes a bug in handling unstructured control-flow in compound loop constructs. The fix makes sure that unstructured CF does not get lowered until we reach the last item of the compound construct. This way, we avoid moving block of unstructured loops in-between the middle items of the construct and messing (i.e. adding operations) to these block while doing so.

---

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

2 Files Affected:

• (modified) flang/lib/Lower/OpenMP/OpenMP.cpp (+5-2)
• (modified) flang/test/Lower/OpenMP/loop-compound.f90 (+35-1)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 60c83586e468b6..8195f4a897a90b 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -590,9 +590,11 @@ static void createBodyOfOp(mlir::Operation &op, const OpWithBodyGenInfo &info,
   mlir::Operation *marker = insertMarker(firOpBuilder);
 
   // If it is an unstructured region, create empty blocks for all evaluations.
-  if (info.eval.lowerAsUnstructured())
+  if (lower::omp::isLastItemInQueue(item, queue) &&
+      info.eval.lowerAsUnstructured()) {
     lower::createEmptyRegionBlocks<mlir::omp::TerminatorOp, mlir::omp::YieldOp>(
         firOpBuilder, info.eval.getNestedEvaluations());
+  }
 
   // Start with privatization, so that the lowering of the nested
   // code will use the right symbols.
@@ -966,7 +968,8 @@ static void genBodyOfTargetOp(
 
   // Create blocks for unstructured regions. This has to be done since
   // blocks are initially allocated with the function as the parent region.
-  if (eval.lowerAsUnstructured()) {
+  if (lower::omp::isLastItemInQueue(item, queue) &&
+      eval.lowerAsUnstructured()) {
     lower::createEmptyRegionBlocks<mlir::omp::TerminatorOp, mlir::omp::YieldOp>(
         firOpBuilder, eval.getNestedEvaluations());
   }
diff --git a/flang/test/Lower/OpenMP/loop-compound.f90 b/flang/test/Lower/OpenMP/loop-compound.f90
index e76edfe052f745..8c025ab237920e 100644
--- a/flang/test/Lower/OpenMP/loop-compound.f90
+++ b/flang/test/Lower/OpenMP/loop-compound.f90
@@ -5,7 +5,7 @@
 ! RUN: %flang_fc1 -fopenmp -emit-hlfir %s -o - | FileCheck %s
 
 program main
-  integer :: i
+  integer :: i,j
 
   ! TODO When composite constructs are supported add:
   ! - TASKLOOP SIMD
@@ -231,4 +231,38 @@ program main
   do i = 1, 10
   end do
   !$omp end teams distribute simd
+
+  ! ----------------------------------------------------------------------------
+  ! Unstructured control-flow in loop
+  ! ----------------------------------------------------------------------------
+  ! CHECK:      omp.target
+  ! CHECK:      omp.teams
+  ! CHECK:      omp.parallel
+  ! CHECK:      omp.distribute
+  ! CHECK-NEXT: omp.wsloop
+  ! CHECK-NEXT: omp.loop_nest
+  !
+  ! Verify the conrol-flow of the unstructured inner loop.
+  ! CHECK:        cf.br ^[[BB1:.*]]
+  ! CHECK:      ^[[BB1]]:
+  ! CHECK:        cf.br ^[[BB2:.*]]
+  ! CHECK:      ^[[BB2]]:
+  ! CHECK:        cf.cond_br %{{.*}}, ^[[BB3:.*]], ^[[BB6:.*]]
+  ! CHECK:      ^[[BB3]]:
+  ! CHECK:        cf.cond_br %{{.*}}, ^[[BB4:.*]], ^[[BB5:.*]]
+  ! CHECK:      ^[[BB4]]:
+  ! CHECK:        cf.br ^[[BB6]]
+  ! CHECK:      ^[[BB5]]:
+  ! CHECK:        cf.br ^[[BB2]]
+  ! CHECK:      ^[[BB6]]:
+  ! CHECK-NEXT:   omp.yield
+  !$omp target teams distribute parallel do
+  do i = 1, 10
+    outerloop: do j = i-1, i+i
+      if (j == i) then
+        exit outerloop
+      end if
+    end do outerloop
+  end do
+  !$omp end target teams distribute parallel do
 end program main

[skatrak]

Thank you Kareem, this LGTM. Give it a day or wait for a second review before merging, in case I'm missing anything here.

[tblah]

LGTM too