[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 21:15:57 PDT 2025
https://github.com/benwu25 updated https://github.com/llvm/llvm-project/pull/160783
>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/3] [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/3] [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:
>From 3c8e0e1955990073727dfcc9c439e5ae449bdf4a Mon Sep 17 00:00:00 2001
From: benwu25 <soggysocks206 at gmail.com>
Date: Thu, 25 Sep 2025 21:11:12 -0700
Subject: [PATCH 3/3] [MLIR][CF] Fix canonicalizer regression with 1-block
cycle (#159743)
I caused a regression in flang due to my previous commits, since I
did not consider encountering a single-block cycle. I updated my
check to look for 1-block cycles in addition to any larger cycles
containing the successor block.
---
mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp | 8 +++++---
mlir/test/Transforms/control-flow-cycles.mlir | 12 ++++++++++++
2 files changed, 17 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index 0553fe4f8b3be..ae1c98ed4b4c6 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
@@ -123,14 +123,16 @@ static LogicalResult collapseBranch(Block *&successor,
if (successorDest == successor)
return failure();
// Don't try to collapse branches which participate in a cycle.
- BranchOp nextBranch = dyn_cast<BranchOp>(successorDest->getTerminator());
+ Block *currBlock = successorDest;
+ BranchOp nextBranch = dyn_cast<BranchOp>(currBlock->getTerminator());
while (nextBranch) {
Block *nextBranchDest = nextBranch.getDest();
- if (!nextBranchDest)
+ if (!nextBranchDest || nextBranchDest == currBlock)
break;
else if (nextBranchDest == successor)
return failure();
- nextBranch = dyn_cast<BranchOp>(nextBranchDest->getTerminator());
+ currBlock = nextBranchDest;
+ nextBranch = dyn_cast<BranchOp>(currBlock->getTerminator());
}
// Update the operands to the successor. If the branch parent has no
diff --git a/mlir/test/Transforms/control-flow-cycles.mlir b/mlir/test/Transforms/control-flow-cycles.mlir
index a3f67ccc5d0c7..ebb9275bb5cf5 100644
--- a/mlir/test/Transforms/control-flow-cycles.mlir
+++ b/mlir/test/Transforms/control-flow-cycles.mlir
@@ -79,3 +79,15 @@ func.func @delayed_3_cycle() {
^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
+}
More information about the Mlir-commits
mailing list