[flang-commits] [flang] [flang][OpenMP] Handle unstructured CF in compound loop constructs (PR #111111)

[Kareem Ergawy via flang-commits]

https://github.com/ergawy created https://github.com/llvm/llvm-project/pull/111111

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.

>From fba8a524a5fb2e376bf493b2244b7590ba3c20da Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Fri, 4 Oct 2024 02:36:26 -0500
Subject: [PATCH] [flang][OpenMP] Handle unstructured CF in compound loop
 constructs

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.
---
 flang/lib/Lower/OpenMP/OpenMP.cpp         |  7 +++--
 flang/test/Lower/OpenMP/loop-compound.f90 | 36 ++++++++++++++++++++++-
 2 files changed, 40 insertions(+), 3 deletions(-)

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