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
matthias-springer merged 15 commits into from
Sep 29, 2025
Merged

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

Show file tree
Hide file tree
Changes from 13 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
12 changes: 12 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,18 @@ 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{successorDest};
while (nextBranch) {
Copy link
Member

@matthias-springer matthias-springer Sep 26, 2025

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

@matthias-springer matthias-springer Sep 26, 2025

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 (nextBranchDest == successor)
return failure();
else if (visited.contains(nextBranchDest))
break;
Copy link
Member

@matthias-springer matthias-springer Sep 27, 2025

Choose a reason for hiding this comment

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

Should this be break or return failure()?

Copy link
Member

@matthias-springer matthias-springer Sep 27, 2025

Choose a reason for hiding this comment

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

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

Copy link
Contributor Author

@benwu25 benwu25 Sep 28, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

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.

Copy link
Contributor Author

@benwu25 benwu25 Sep 28, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

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.

Copy link
Contributor Author

@benwu25 benwu25 Sep 29, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

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:     }

Copy link
Member

@matthias-springer matthias-springer Sep 29, 2025

Choose a reason for hiding this comment

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

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.

Copy link
Contributor Author

@benwu25 benwu25 Sep 29, 2025

Choose a reason for hiding this comment

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

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!

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
136 changes: 136 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,139 @@ 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.br ^bb1
// CHECK: ^bb1:
// CHECK: cf.br ^bb1
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.br ^bb1
// CHECK: ^bb1:
// CHECK: cf.br ^bb1 {E}
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
}