[MLIR][CF] Avoid collapsing blocks which participate in cycles #160783

benwu25 · 2025-09-25T21:54:58Z

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!

…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.

llvmbot · 2025-09-25T21:55:31Z

@llvm/pr-subscribers-mlir

Author: None (benwu25)

Changes

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!

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

2 Files Affected:

(modified) mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp (+10)
(added) mlir/test/Transforms/control-flow-cycles.mlir (+81)

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..a3f67ccc5d0c7
--- /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:           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
+}

llvmbot · 2025-09-25T21:55:32Z

@llvm/pr-subscribers-mlir-cf

Author: None (benwu25)

Changes

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!

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

2 Files Affected:

(modified) mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp (+10)
(added) mlir/test/Transforms/control-flow-cycles.mlir (+81)

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..a3f67ccc5d0c7
--- /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:           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
+}

benwu25 · 2025-09-25T22:45:11Z

Oops, looks like I missed at least one case:

func.func @test() {
    cf.br ^bb1
  ^bb1:
    cf.br ^bb2
  ^bb2:
    cf.br ^bb2
}

Will push a fix soon...

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.

benwu25 · 2025-09-26T05:44:54Z

@matthias-springer @zero9178

matthias-springer · 2025-09-26T07:15:34Z

mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp

+  // Don't try to collapse branches which participate in a cycle.
+  Block *currBlock = successorDest;
+  BranchOp nextBranch = dyn_cast<BranchOp>(currBlock->getTerminator());
+  while (nextBranch) {


Is it possible that this while loop runs into an endless loop? E.g.:

A -> B -> C <- | \ v \ D -> E

Assuming that we call collapseBranch(A, ...).

I feel like nextBranchDest == successor is not sufficient and we need a DenseSet<Block *> visited instead?

I see your argument, and weirdly enough, I have already written this testcase (@delayed_3_cycle). I think that A -> B -> C is first canonicalized as just C by a few calls to simplifyBrToBlockWithSinglePred, so then when we collapse on C it somehow works out.

However, if B has another predecessor, maybe the loop will be broken. Will look into that soon.

Here's an example that gets into an endless loop.

func.func @cycle_4_blocks(%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 }

Yes, I think a DenseSet will be the right tool here. I have added changes to employ one for cycle detection and new tests which require it.

matthias-springer · 2025-09-26T07:16:26Z

mlir/test/Transforms/control-flow-cycles.mlir

@@ -0,0 +1,93 @@
+// RUN: mlir-opt --canonicalize %s | FileCheck %s
+
+// Test that control-flow cycles are not simplified infinitely.


All these test cases should be in test/Dialect/ControlFlow/canonicalize.mlir.

We need to use a DenseSet in case we encounter cycles further in the chain of control flow which have not been simplified by other cf transforms.

matthias-springer · 2025-09-27T10:19:08Z

mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp

+    if (nextBranchDest == successor)
+      return failure();
+    else if (visited.contains(nextBranchDest))
+      break;


Should this be break or return failure()?

Also, is the above if (nextBranchDest == successor) still necessary if you initialize the visited set with successor?

I chose break to allow for more merges. The only truly bad case seems to be when successor itself is part of the cycle. I.e., I think it is only necessary to bail out when that is the case. Consider the testcase from above:

func.func @cycle_4_blocks(%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 }

With the prior implementation, this testcase gets simplified to:

func.func @cycle_4_blocks(%c : i1) { cf.br ^bb1 ^bb1: cf.br ^bb1 {E} }

But if we return failure(), the result is:

func.func @cycle_4_blocks(%c : i1) { cf.cond_br %c, ^bb1, ^bb3 ^bb1: // 2 preds: ^bb0, ^bb3 cf.br ^bb2 {A} ^bb2: // 2 preds: ^bb1, ^bb2 cf.br ^bb2 {E} ^bb3: // pred: ^bb0 cf.br ^bb1 }

I had thought that a few canonicalize passes were able to merge blocks if we chose break before the cycle itself is reached. Here it seems that blocks were still merged successfully, but the cycle itself is not simplified to a 1-block cycle for some reason.

However, it doesn't matter much if infinite loops are fully optimized, and it is the most simple and correct to just bail out, so I will add that change.

Update:

After adding this change to bail out early, I caused a regression in flang:
flang/test/Lower/OpenMP/infinite-loop-in-construct.f90.

I broke this test with my initial PR because there was an infinite loop. Now, it seems like the infinite loop is less optimized than expected for this testcase. I suppose I could fix this by updating this testcase to expect a different infinite loop? Or, if we use break, the loop is more optimized and the test passes. I am not sure which is preferred.

BTW, we would need to update the flang test to check for this cycle instead of what it currently checks if we want to bail out early.

! CHECK-LABEL: func.func @_QPsb ! CHECK: omp.parallel ! CHECK: cf.cond_br %{{[0-9]+}}, ^bb1, ^bb2 ! CHECK-NEXT: ^bb1: ! CHECK: cf.br ^bb2 ! CHECK-NEXT: ^bb2: ! CHECK-NEXT: cf.br ^bb3 ! CHECK-NEXT: ^bb3: ! CHECK-NEXT: cf.br ^bb3 ! CHECK-NEXT: }

The reason why I suggested failure instead of break: I believe the break will still perform some canonicalization, so the function will return success. This will make the canonicalization succeed and another iteration of the underlying greedy pattern rewriter will be scheduled. That's problematic because collapseBranch can be applied over and over without making any progress.

I would update the expected IR in the Flang test case.

Ok, I have pushed again the changes to bail out early, as well as updated the flang test. @matthias-springer could you merge these changes for me if they are now acceptable? Thanks!

matthias-springer · 2025-09-27T10:21:46Z

mlir/test/Dialect/ControlFlow/canonicalize.mlir

+// CHECK:           cf.br ^bb1
+func.func @unsimplified_cycle_1(%c : i1) {
+  cf.cond_br %c, ^bb1, ^bb2
+  ^bb1:


nit: Basic blocks should not be indented. (See other test cases in this file.)

…m#159743) It is safest to bail out as soon as we detect a cycle in collapseBranch. We need not continue with optimizations on infinite loops.

llvm#159743) Bailing out as soon as a cycle is detected causes a regression in flang, where the test expects more optimizations than we allowed on an infinite loop. We may be able to adjust this testcase.

I forgot to update this test with my previous commit. (e8135aa)

We can bail out of collapseBranch upon detecting any cycle, but we must update a test in flang to prevent a regression.

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.

…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 llvm#159743.

benwu25 added 2 commits September 25, 2025 13:46

[MLIR][CF] Update control-flow-cycles test (llvm#159743)

a7c8f61

llvmbot added mlir mlir:cf labels Sep 25, 2025

Merge branch 'main' into main

35ac930

benwu25 added 2 commits September 25, 2025 21:11

Merge branch 'main' into main

e051038

matthias-springer reviewed Sep 26, 2025

View reviewed changes

benwu25 added 2 commits September 26, 2025 13:37

[MLIR][CF] Use a DenseSet to detect more cycles (llvm#159743)

4250c1d

We need to use a DenseSet in case we encounter cycles further in the chain of control flow which have not been simplified by other cf transforms.

Merge branch 'main' into main

d6f7526

benwu25 requested a review from matthias-springer September 26, 2025 22:36

matthias-springer reviewed Sep 27, 2025

View reviewed changes

benwu25 added 6 commits September 28, 2025 00:04

[MLIR][CF] Bail out of collapseBranch when any cycle is detected (llv…

6357614

…m#159743) It is safest to bail out as soon as we detect a cycle in collapseBranch. We need not continue with optimizations on infinite loops.

Merge branch 'main' into main

44d52d3

[MLIR][CF] Do not bail out early in collapseBranch to avoid regression (

e8135aa

llvm#159743) Bailing out as soon as a cycle is detected causes a regression in flang, where the test expects more optimizations than we allowed on an infinite loop. We may be able to adjust this testcase.

Merge branch 'main' into main

63e77da

[MLIR][CF] Update broken test (llvm#159743)

9c27d37

I forgot to update this test with my previous commit. (e8135aa)

Merge branch 'main' into main

a4ae338

benwu25 requested a review from matthias-springer September 28, 2025 08:29

matthias-springer approved these changes Sep 29, 2025

View reviewed changes

benwu25 added 2 commits September 29, 2025 00:58

[MLIR][CF] Abort collapseBranch for any cycle (llvm#159743)

13dfab5

We can bail out of collapseBranch upon detecting any cycle, but we must update a test in flang to prevent a regression.

Merge branch 'main' into main

d2a47c6

matthias-springer merged commit 631b89c into llvm:main Sep 29, 2025
9 checks passed

		@@ -0,0 +1,93 @@
		// RUN: mlir-opt --canonicalize %s \| FileCheck %s

		// Test that control-flow cycles are not simplified infinitely.

[MLIR][CF] Avoid collapsing blocks which participate in cycles #160783

[MLIR][CF] Avoid collapsing blocks which participate in cycles #160783

Conversation

benwu25 commented Sep 25, 2025

Uh oh!

llvmbot commented Sep 25, 2025

Uh oh!

llvmbot commented Sep 25, 2025

Uh oh!

benwu25 commented Sep 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

benwu25 commented Sep 26, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benwu25 Sep 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benwu25 Sep 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benwu25 Sep 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benwu25 Sep 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benwu25 Sep 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

benwu25 commented Sep 25, 2025 •

edited

Loading

benwu25 Sep 26, 2025 •

edited

Loading

benwu25 Sep 26, 2025 •

edited

Loading

benwu25 Sep 28, 2025 •

edited

Loading

benwu25 Sep 28, 2025 •

edited

Loading

benwu25 Sep 29, 2025 •

edited

Loading