[Mlir-commits] [mlir] 631b89c - [MLIR][CF] Avoid collapsing blocks which participate in cycles (#160783)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Mon Sep 29 01:39:36 PDT 2025


Author: benwu25
Date: 2025-09-29T10:39:31+02:00
New Revision: 631b89cc4766d2b8f0f5a17c6d90fe920436a2c9

URL: https://github.com/llvm/llvm-project/commit/631b89cc4766d2b8f0f5a17c6d90fe920436a2c9
DIFF: https://github.com/llvm/llvm-project/commit/631b89cc4766d2b8f0f5a17c6d90fe920436a2c9.diff

LOG: [MLIR][CF] Avoid collapsing blocks which participate in cycles (#160783)

Previously, collapseBranch did not return failure for successor blocks which were part of a cycle. mlir-opt --canonicalize would run indefinitely for any N-block cycle which is kicked off with an unconditional jump. The simplifyPassThroughBr transform would continue alternating which block was targeted in ^bb0, resulting in an infinite loop.

collapseBranch will not result in any useful transformation on blocks which participate in cycles, since the block is aliased by a different block. To avoid this, we can check for cycles in collapseBranch and abort when one is detected. Simplification of the cycle is left for other transforms.

Fixes #159743.

Added: 
    

Modified: 
    flang/test/Lower/OpenMP/infinite-loop-in-construct.f90
    mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
    mlir/test/Dialect/ControlFlow/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/flang/test/Lower/OpenMP/infinite-loop-in-construct.f90 b/flang/test/Lower/OpenMP/infinite-loop-in-construct.f90
index 16b400a231860..f02d0e5ccc53c 100644
--- a/flang/test/Lower/OpenMP/infinite-loop-in-construct.f90
+++ b/flang/test/Lower/OpenMP/infinite-loop-in-construct.f90
@@ -8,8 +8,10 @@
 ! CHECK:            cf.cond_br %{{[0-9]+}}, ^bb1, ^bb2
 ! CHECK-NEXT:     ^bb1:  // pred: ^bb0
 ! CHECK:            cf.br ^bb2
-! CHECK-NEXT:     ^bb2:  // 3 preds: ^bb0, ^bb1, ^bb2
-! CHECK-NEXT:       cf.br ^bb2
+! CHECK-NEXT:     ^bb2:  // 2 preds: ^bb0, ^bb1
+! CHECK:            cf.br ^bb3
+! CHECK-NEXT:     ^bb3:  // 2 preds: ^bb2, ^bb3
+! CHECK:            cf.br ^bb3
 ! CHECK-NEXT:     }
 
 subroutine sb(ninter, numnod)

diff  --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index 582593adfa5c0..f1da1a125e9ef 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
@@ -122,6 +122,16 @@ static LogicalResult collapseBranch(Block *&successor,
   Block *successorDest = successorBranch.getDest();
   if (successorDest == successor)
     return failure();
+  // Don't try to collapse branches which participate in a cycle.
+  BranchOp nextBranch = dyn_cast<BranchOp>(successorDest->getTerminator());
+  llvm::DenseSet<Block *> visited{successor, successorDest};
+  while (nextBranch) {
+    Block *nextBranchDest = nextBranch.getDest();
+    if (visited.contains(nextBranchDest))
+      return failure();
+    visited.insert(nextBranchDest);
+    nextBranch = dyn_cast<BranchOp>(nextBranchDest->getTerminator());
+  }
 
   // Update the operands to the successor. If the branch parent has no
   // arguments, we can use the branch operands directly.

diff  --git a/mlir/test/Dialect/ControlFlow/canonicalize.mlir b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
index bf69935a00bf0..17f7d28ba59fb 100644
--- a/mlir/test/Dialect/ControlFlow/canonicalize.mlir
+++ b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
@@ -490,3 +490,147 @@ func.func @branchCondProp(%arg0: i1) {
 ^exit:
   return
 }
+
+// -----
+
+/// Test that control-flow cycles are not simplified infinitely.
+
+// CHECK-LABEL:   @cycle_2_blocks
+// CHECK:           cf.br ^bb1
+// CHECK:         ^bb1:
+// CHECK:           cf.br ^bb1
+func.func @cycle_2_blocks() {
+  cf.br ^bb1
+^bb1:
+  cf.br ^bb2
+^bb2:
+  cf.br ^bb1
+}
+
+// CHECK-LABEL:   @no_cycle_2_blocks
+// CHECK:           %[[VAL_0:.*]] = arith.constant 1 : i32
+// CHECK:           return %[[VAL_0]] : i32
+func.func @no_cycle_2_blocks() -> i32 {
+  cf.br ^bb1
+^bb1:
+  cf.br ^bb2
+^bb2:
+  cf.br ^bb3
+^bb3:
+  %ret = arith.constant 1 : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL:   @cycle_4_blocks
+// CHECK:           cf.br ^bb1
+// CHECK:         ^bb1:
+// CHECK:           cf.br ^bb1
+func.func @cycle_4_blocks() {
+  cf.br ^bb1
+^bb1:
+  cf.br ^bb2
+^bb2:
+  cf.br ^bb3
+^bb3:
+  cf.br ^bb4
+^bb4:
+  cf.br ^bb1
+}
+
+// CHECK-LABEL:   @no_cycle_4_blocks
+// CHECK:           %[[VAL_0:.*]] = arith.constant 1 : i32
+// CHECK:           return %[[VAL_0]] : i32
+func.func @no_cycle_4_blocks() -> i32 {
+  cf.br ^bb1
+^bb1:
+  cf.br ^bb2
+^bb2:
+  cf.br ^bb3
+^bb3:
+  cf.br ^bb4
+^bb4:
+  cf.br ^bb5
+^bb5:
+  %ret = arith.constant 1 : i32
+  return %ret : i32
+}
+
+// CHECK-LABEL:   @delayed_3_cycle
+// CHECK:           cf.br ^bb1
+// CHECK:         ^bb1:
+// CHECK:           cf.br ^bb1
+func.func @delayed_3_cycle() {
+  cf.br ^bb1
+^bb1:
+  cf.br ^bb2
+^bb2:
+  cf.br ^bb3
+^bb3:
+  cf.br ^bb4
+^bb4:
+  cf.br ^bb5
+^bb5:
+  cf.br ^bb3
+}
+
+// CHECK-LABEL:   @cycle_1_block
+// CHECK:           cf.br ^bb1
+// CHECK:         ^bb1:
+// CHECK:           cf.br ^bb1
+func.func @cycle_1_block() {
+  cf.br ^bb1
+^bb1:
+  cf.br ^bb2
+^bb2:
+  cf.br ^bb2
+}
+
+// CHECK-LABEL:   @unsimplified_cycle_1
+// CHECK-SAME:      %[[ARG0:.*]]: i1) {
+// CHECK:           cf.cond_br %[[ARG0]], ^bb1, ^bb2
+// CHECK:         ^bb1:
+// CHECK:           cf.br ^bb2
+// CHECK:         ^bb2:
+// CHECK:           cf.br ^bb3
+// CHECK:         ^bb3:
+// CHECK:           cf.br ^bb3
+func.func @unsimplified_cycle_1(%c : i1) {
+  cf.cond_br %c, ^bb1, ^bb2
+^bb1:
+  cf.br ^bb2
+^bb2:
+  cf.br ^bb3
+^bb3:
+  cf.br ^bb4
+^bb4:
+  cf.br ^bb3
+}
+
+// Make sure we terminate when other cf passes can't help us.
+
+// CHECK-LABEL:   @unsimplified_cycle_2
+// CHECK-SAME:      %[[ARG0:.*]]: i1) {
+// CHECK:           cf.cond_br %[[ARG0]], ^bb1, ^bb3
+// CHECK:         ^bb1:
+// CHECK:           cf.br ^bb2 {A}
+// CHECK:         ^bb2:
+// CHECK:           cf.br ^bb2 {E}
+// CHECK:         ^bb3:
+// CHECK:           cf.br ^bb1
+func.func @unsimplified_cycle_2(%c : i1) {
+  cf.cond_br %c, ^bb6, ^bb7
+^bb6:
+  cf.br ^bb5 {F}
+^bb5:
+  cf.br ^bb1 {A}
+^bb1:
+  cf.br ^bb2 {B}
+^bb2:
+  cf.br ^bb3 {C}
+^bb3:
+  cf.br ^bb4 {D}
+^bb4:
+  cf.br ^bb1 {E}
+^bb7:
+  cf.br ^bb6
+}


        


More information about the Mlir-commits mailing list