[flang][OpenMP] Skip multi-block teams regions when processing loop directives

[ergawy]

Fixes a regression when the generic loop directive conversion pass encounters a multi-block teams region. At the moment, we skip such regions.

[llvmbot]

@llvm/pr-subscribers-flang-openmp

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

Author: Kareem Ergawy (ergawy)

Changes

Fixes a regression when the generic loop directive conversion pass encounters a multi-block teams region. At the moment, we skip such regions.

---

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

2 Files Affected:

• (modified) flang/lib/Optimizer/OpenMP/GenericLoopConversion.cpp (+2-2)
• (modified) flang/test/Lower/OpenMP/loop-directive.f90 (+21)

diff --git a/flang/lib/Optimizer/OpenMP/GenericLoopConversion.cpp b/flang/lib/Optimizer/OpenMP/GenericLoopConversion.cpp
index 2ca1de1e1cbd6..74ad6330b11a7 100644
--- a/flang/lib/Optimizer/OpenMP/GenericLoopConversion.cpp
+++ b/flang/lib/Optimizer/OpenMP/GenericLoopConversion.cpp
@@ -398,8 +398,8 @@ class ReductionsHoistingPattern
 
   static mlir::omp::LoopOp
   tryToFindNestedLoopWithReduction(mlir::omp::TeamsOp teamsOp) {
-    assert(!teamsOp.getRegion().empty() &&
-           teamsOp.getRegion().getBlocks().size() == 1);
+    if (teamsOp.getRegion().getBlocks().size() != 1)
+      return nullptr;
 
     mlir::Block &teamsBlock = *teamsOp.getRegion().begin();
     auto loopOpIter = llvm::find_if(teamsBlock, [](mlir::Operation &op) {
diff --git a/flang/test/Lower/OpenMP/loop-directive.f90 b/flang/test/Lower/OpenMP/loop-directive.f90
index 8cf2e98332d7d..954985e2d64f1 100644
--- a/flang/test/Lower/OpenMP/loop-directive.f90
+++ b/flang/test/Lower/OpenMP/loop-directive.f90
@@ -337,3 +337,24 @@ subroutine loop_teams_loop_reduction
     x = x + i
   end do
 end subroutine
+
+
+! Tests a regression when the pass encounters a multi-block `teams` region.
+subroutine multi_block_teams
+  implicit none
+  integer :: i
+
+  ! CHECK: omp.target {{.*}} {
+  ! CHECK:   omp.teams {
+  ! CHECK:   ^bb1:
+  ! CHECK:     cf.br ^bb2
+  ! CHECK:   ^bb2:
+  ! CHECK:     omp.terminator
+  ! CHECK:   }
+  ! CHECK: }
+  !$omp target teams
+  select case (i)
+  case(1)
+  end select
+  !$omp end target teams
+end subroutine