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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Thu Sep 25 14:55:05 PDT 2025


https://github.com/benwu25 created https://github.com/llvm/llvm-project/pull/160783

This PR addresses #159743.

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.

Could someone help review and merge? Thanks!

>From aed78e68e9886abf50c68a0d34559b4746c02c31 Mon Sep 17 00:00:00 2001
From: benwu25 <soggysocks206 at gmail.com>
Date: Thu, 25 Sep 2025 13:46:10 -0700
Subject: [PATCH 1/2] [MLIR][CF] Avoid collapsing blocks which participate in
 cycles (#159743)

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.
---
 .../Dialect/ControlFlow/IR/ControlFlowOps.cpp | 10 +++
 mlir/test/Transforms/control-flow-cycles.mlir | 81 +++++++++++++++++++
 2 files changed, 91 insertions(+)
 create mode 100644 mlir/test/Transforms/control-flow-cycles.mlir

diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index 582593adfa5c0..0553fe4f8b3be 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());
+  while (nextBranch) {
+    Block *nextBranchDest = nextBranch.getDest();
+    if (!nextBranchDest)
+      break;
+    else if (nextBranchDest == successor)
+      return failure();
+    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/Transforms/control-flow-cycles.mlir b/mlir/test/Transforms/control-flow-cycles.mlir
new file mode 100644
index 0000000000000..b1c0111c33668
--- /dev/null
+++ b/mlir/test/Transforms/control-flow-cycles.mlir
@@ -0,0 +1,81 @@
+// RUN: mlir-opt --canonicalize %s | FileCheck %s
+
+// Test that control-flow cycles are not simplified infinitely.
+
+// CHECK-LABEL: @cycle_2_blocks
+// CHECK-NEXT: cf.br ^bb1
+// CHECK-NEXT: ^bb1:
+// CHECK-NEXT: 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-NEXT: %c1_i32 = arith.constant 1 : i32
+// CHECK-NEXT: return %c1_i32 : 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-NEXT: cf.br ^bb1
+// CHECK-NEXT: ^bb1:
+// CHECK-NEXT: 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-NEXT: %c1_i32 = arith.constant 1 : i32
+// CHECK-NEXT: return %c1_i32 : 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-NEXT: cf.br ^bb1
+// CHECK-NEXT: ^bb1:
+// CHECK-NEXT: 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
+}

>From a7c8f61d2142e5f49a3b7fb5b83f05bfdfe0d85c Mon Sep 17 00:00:00 2001
From: benwu25 <soggysocks206 at gmail.com>
Date: Thu, 25 Sep 2025 14:39:57 -0700
Subject: [PATCH 2/2] [MLIR][CF] Update control-flow-cycles test (#159743)

---
 mlir/test/Transforms/control-flow-cycles.mlir | 36 +++++++++----------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/mlir/test/Transforms/control-flow-cycles.mlir b/mlir/test/Transforms/control-flow-cycles.mlir
index b1c0111c33668..a3f67ccc5d0c7 100644
--- a/mlir/test/Transforms/control-flow-cycles.mlir
+++ b/mlir/test/Transforms/control-flow-cycles.mlir
@@ -2,10 +2,10 @@
 
 // Test that control-flow cycles are not simplified infinitely.
 
-// CHECK-LABEL: @cycle_2_blocks
-// CHECK-NEXT: cf.br ^bb1
-// CHECK-NEXT: ^bb1:
-// CHECK-NEXT: cf.br ^bb1
+// CHECK-LABEL:   @cycle_2_blocks
+// CHECK:           cf.br ^bb1
+// CHECK:         ^bb1:
+// CHECK:           cf.br ^bb1
 func.func @cycle_2_blocks() {
   cf.br ^bb1
   ^bb1:
@@ -14,9 +14,9 @@ func.func @cycle_2_blocks() {
     cf.br ^bb1
 }
 
-// CHECK-LABEL: @no_cycle_2_blocks
-// CHECK-NEXT: %c1_i32 = arith.constant 1 : i32
-// CHECK-NEXT: return %c1_i32 : i32
+// 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:
@@ -28,10 +28,10 @@ func.func @no_cycle_2_blocks() -> i32 {
     return %ret : i32
 }
 
-// CHECK-LABEL: @cycle_4_blocks
-// CHECK-NEXT: cf.br ^bb1
-// CHECK-NEXT: ^bb1:
-// CHECK-NEXT: cf.br ^bb1
+// CHECK-LABEL:   @cycle_4_blocks
+// CHECK:           cf.br ^bb1
+// CHECK:         ^bb1:
+// CHECK:           cf.br ^bb1
 func.func @cycle_4_blocks() {
   cf.br ^bb1
   ^bb1:
@@ -44,9 +44,9 @@ func.func @cycle_4_blocks() {
     cf.br ^bb1
 }
 
-// CHECK-LABEL: @no_cycle_4_blocks
-// CHECK-NEXT: %c1_i32 = arith.constant 1 : i32
-// CHECK-NEXT: return %c1_i32 : i32
+// 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:
@@ -62,10 +62,10 @@ func.func @no_cycle_4_blocks() -> i32 {
     return %ret : i32
 }
 
-// CHECK-LABEL: @delayed_3_cycle
-// CHECK-NEXT: cf.br ^bb1
-// CHECK-NEXT: ^bb1:
-// CHECK-NEXT: cf.br ^bb1
+// CHECK-LABEL:   @delayed_3_cycle
+// CHECK:           cf.br ^bb1
+// CHECK:         ^bb1:
+// CHECK:           cf.br ^bb1
 func.func @delayed_3_cycle() {
   cf.br ^bb1
   ^bb1:



More information about the Mlir-commits mailing list