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

[llvmlistbot at llvm.org]

================
@@ -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) {
+    Block *nextBranchDest = nextBranch.getDest();
+    if (nextBranchDest == successor)
+      return failure();
+    else if (visited.contains(nextBranchDest))
+      break;
----------------
benwu25 wrote:

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., it is only absolutely necessary to bail out when this 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 believe this is because a few canonicalize passes succeed in merging blocks if we choose `break` before the cycle itself is reached.

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

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