Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 15 commits into from
Sep 29, 2025
+158 −2
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
aed78e6
[MLIR][CF] Avoid collapsing blocks which participate in cycles (#159743)
benwu25 Sep 25, 2025
a7c8f61
[MLIR][CF] Update control-flow-cycles test (#159743)
benwu25 Sep 25, 2025
35ac930
Merge branch 'main' into main
benwu25 Sep 25, 2025
3c8e0e1
[MLIR][CF] Fix canonicalizer regression with 1-block cycle (#159743)
benwu25 Sep 26, 2025
e051038
Merge branch 'main' into main
benwu25 Sep 26, 2025
4250c1d
[MLIR][CF] Use a DenseSet to detect more cycles (#159743)
benwu25 Sep 26, 2025
d6f7526
Merge branch 'main' into main
benwu25 Sep 26, 2025
6357614
[MLIR][CF] Bail out of collapseBranch when any cycle is detected (#15…
benwu25 Sep 28, 2025
44d52d3
Merge branch 'main' into main
benwu25 Sep 28, 2025
e8135aa
[MLIR][CF] Do not bail out early in collapseBranch to avoid regressio…
benwu25 Sep 28, 2025
63e77da
Merge branch 'main' into main
benwu25 Sep 28, 2025
9c27d37
[MLIR][CF] Update broken test (#159743)
benwu25 Sep 28, 2025
a4ae338
Merge branch 'main' into main
benwu25 Sep 28, 2025
13dfab5
[MLIR][CF] Abort collapseBranch for any cycle (#159743)
benwu25 Sep 29, 2025
d2a47c6
Merge branch 'main' into main
benwu25 Sep 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions flang/test/Lower/OpenMP/infinite-loop-in-construct.f90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
10 changes: 10 additions & 0 deletions mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Contributor Author

@benwu25 benwu25 Sep 26, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
}

Copy link
Contributor Author

@benwu25 benwu25 Sep 26, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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.
Expand Down
144 changes: 144 additions & 0 deletions mlir/test/Dialect/ControlFlow/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}