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

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

Author: benwu25

flang/test/Lower/OpenMP/infinite-loop-in-construct.f90

@@ -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)

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

@@ -124,13 +124,11 @@ static LogicalResult collapseBranch(Block *&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};
+  llvm::DenseSet<Block *> visited{successor, successorDest};
   while (nextBranch) {
     Block *nextBranchDest = nextBranch.getDest();
-    if (nextBranchDest == successor)
+    if (visited.contains(nextBranchDest))
       return failure();
-    else if (visited.contains(nextBranchDest))
-      break;
     visited.insert(nextBranchDest);
     nextBranch = dyn_cast<BranchOp>(nextBranchDest->getTerminator());
   }

mlir/test/Dialect/ControlFlow/canonicalize.mlir

@@ -587,9 +587,13 @@ func.func @cycle_1_block() {
 
 // CHECK-LABEL:   @unsimplified_cycle_1
 // CHECK-SAME:      %[[ARG0:.*]]: i1) {
-// CHECK:           cf.br ^bb1
+// CHECK:           cf.cond_br %[[ARG0]], ^bb1, ^bb2
 // CHECK:         ^bb1:
-// CHECK:           cf.br ^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:
@@ -606,9 +610,13 @@ func.func @unsimplified_cycle_1(%c : i1) {
 
 // CHECK-LABEL:   @unsimplified_cycle_2
 // CHECK-SAME:      %[[ARG0:.*]]: i1) {
-// CHECK:           cf.br ^bb1
+// CHECK:           cf.cond_br %[[ARG0]], ^bb1, ^bb3
 // CHECK:         ^bb1:
-// CHECK:           cf.br ^bb1 {E}
+// 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: