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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Sun Sep 28 00:10:51 PDT 2025


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

>From aed78e68e9886abf50c68a0d34559b4746c02c31 Mon Sep 17 00:00:00 2001
From: benwu25 <soggysocks206 at gmail.com>
Date: Thu, 25 Sep 2025 13:46:10 -0700
Subject: [PATCH 1/5] [MLIR][CF] Avoid collapsing blocks which participate in
 cycles (#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.
---
 .../Dialect/ControlFlow/IR/ControlFlowOps.cpp | 10 +++
 mlir/test/Transforms/control-flow-cycles.mlir | 81 +++++++++++++++++++
 2 files changed, 91 insertions(+)
 create mode 100644 mlir/test/Transforms/control-flow-cycles.mlir

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..b1c0111c33668
--- /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-NEXT: cf.br ^bb1
+// CHECK-NEXT: ^bb1:
+// CHECK-NEXT: 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-NEXT: %c1_i32 = arith.constant 1 : i32
+// CHECK-NEXT: return %c1_i32 : 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-NEXT: cf.br ^bb1
+// CHECK-NEXT: ^bb1:
+// CHECK-NEXT: 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-NEXT: %c1_i32 = arith.constant 1 : i32
+// CHECK-NEXT: return %c1_i32 : 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-NEXT: cf.br ^bb1
+// CHECK-NEXT: ^bb1:
+// CHECK-NEXT: 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
+}

>From a7c8f61d2142e5f49a3b7fb5b83f05bfdfe0d85c Mon Sep 17 00:00:00 2001
From: benwu25 <soggysocks206 at gmail.com>
Date: Thu, 25 Sep 2025 14:39:57 -0700
Subject: [PATCH 2/5] [MLIR][CF] Update control-flow-cycles test (#159743)

---
 mlir/test/Transforms/control-flow-cycles.mlir | 36 +++++++++----------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/mlir/test/Transforms/control-flow-cycles.mlir b/mlir/test/Transforms/control-flow-cycles.mlir
index b1c0111c33668..a3f67ccc5d0c7 100644
--- a/mlir/test/Transforms/control-flow-cycles.mlir
+++ b/mlir/test/Transforms/control-flow-cycles.mlir
@@ -2,10 +2,10 @@
 
 // Test that control-flow cycles are not simplified infinitely.
 
-// CHECK-LABEL: @cycle_2_blocks
-// CHECK-NEXT: cf.br ^bb1
-// CHECK-NEXT: ^bb1:
-// CHECK-NEXT: cf.br ^bb1
+// CHECK-LABEL:   @cycle_2_blocks
+// CHECK:           cf.br ^bb1
+// CHECK:         ^bb1:
+// CHECK:           cf.br ^bb1
 func.func @cycle_2_blocks() {
   cf.br ^bb1
   ^bb1:
@@ -14,9 +14,9 @@ func.func @cycle_2_blocks() {
     cf.br ^bb1
 }
 
-// CHECK-LABEL: @no_cycle_2_blocks
-// CHECK-NEXT: %c1_i32 = arith.constant 1 : i32
-// CHECK-NEXT: return %c1_i32 : i32
+// 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:
@@ -28,10 +28,10 @@ func.func @no_cycle_2_blocks() -> i32 {
     return %ret : i32
 }
 
-// CHECK-LABEL: @cycle_4_blocks
-// CHECK-NEXT: cf.br ^bb1
-// CHECK-NEXT: ^bb1:
-// CHECK-NEXT: cf.br ^bb1
+// CHECK-LABEL:   @cycle_4_blocks
+// CHECK:           cf.br ^bb1
+// CHECK:         ^bb1:
+// CHECK:           cf.br ^bb1
 func.func @cycle_4_blocks() {
   cf.br ^bb1
   ^bb1:
@@ -44,9 +44,9 @@ func.func @cycle_4_blocks() {
     cf.br ^bb1
 }
 
-// CHECK-LABEL: @no_cycle_4_blocks
-// CHECK-NEXT: %c1_i32 = arith.constant 1 : i32
-// CHECK-NEXT: return %c1_i32 : i32
+// 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:
@@ -62,10 +62,10 @@ func.func @no_cycle_4_blocks() -> i32 {
     return %ret : i32
 }
 
-// CHECK-LABEL: @delayed_3_cycle
-// CHECK-NEXT: cf.br ^bb1
-// CHECK-NEXT: ^bb1:
-// CHECK-NEXT: cf.br ^bb1
+// CHECK-LABEL:   @delayed_3_cycle
+// CHECK:           cf.br ^bb1
+// CHECK:         ^bb1:
+// CHECK:           cf.br ^bb1
 func.func @delayed_3_cycle() {
   cf.br ^bb1
   ^bb1:

>From 3c8e0e1955990073727dfcc9c439e5ae449bdf4a Mon Sep 17 00:00:00 2001
From: benwu25 <soggysocks206 at gmail.com>
Date: Thu, 25 Sep 2025 21:11:12 -0700
Subject: [PATCH 3/5] [MLIR][CF] Fix canonicalizer regression with 1-block
 cycle (#159743)

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.
---
 mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp |  8 +++++---
 mlir/test/Transforms/control-flow-cycles.mlir      | 12 ++++++++++++
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index 0553fe4f8b3be..ae1c98ed4b4c6 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
@@ -123,14 +123,16 @@ static LogicalResult collapseBranch(Block *&successor,
   if (successorDest == successor)
     return failure();
   // Don't try to collapse branches which participate in a cycle.
-  BranchOp nextBranch = dyn_cast<BranchOp>(successorDest->getTerminator());
+  Block *currBlock = successorDest;
+  BranchOp nextBranch = dyn_cast<BranchOp>(currBlock->getTerminator());
   while (nextBranch) {
     Block *nextBranchDest = nextBranch.getDest();
-    if (!nextBranchDest)
+    if (!nextBranchDest || nextBranchDest == currBlock)
       break;
     else if (nextBranchDest == successor)
       return failure();
-    nextBranch = dyn_cast<BranchOp>(nextBranchDest->getTerminator());
+    currBlock = nextBranchDest;
+    nextBranch = dyn_cast<BranchOp>(currBlock->getTerminator());
   }
 
   // Update the operands to the successor. If the branch parent has no
diff --git a/mlir/test/Transforms/control-flow-cycles.mlir b/mlir/test/Transforms/control-flow-cycles.mlir
index a3f67ccc5d0c7..ebb9275bb5cf5 100644
--- a/mlir/test/Transforms/control-flow-cycles.mlir
+++ b/mlir/test/Transforms/control-flow-cycles.mlir
@@ -79,3 +79,15 @@ func.func @delayed_3_cycle() {
   ^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
+}

>From 4250c1da73375b6b548f8946134b6f918a61c150 Mon Sep 17 00:00:00 2001
From: benwu25 <soggysocks206 at gmail.com>
Date: Fri, 26 Sep 2025 13:33:50 -0700
Subject: [PATCH 4/5] [MLIR][CF] Use a DenseSet to detect more cycles (#159743)

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.
---
 .../Dialect/ControlFlow/IR/ControlFlowOps.cpp |  14 +-
 .../Dialect/ControlFlow/canonicalize.mlir     | 136 ++++++++++++++++++
 mlir/test/Transforms/control-flow-cycles.mlir |  93 ------------
 3 files changed, 143 insertions(+), 100 deletions(-)
 delete mode 100644 mlir/test/Transforms/control-flow-cycles.mlir

diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index ae1c98ed4b4c6..dfa83e30ec40a 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
+++ b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
@@ -123,16 +123,16 @@ static LogicalResult collapseBranch(Block *&successor,
   if (successorDest == successor)
     return failure();
   // Don't try to collapse branches which participate in a cycle.
-  Block *currBlock = successorDest;
-  BranchOp nextBranch = dyn_cast<BranchOp>(currBlock->getTerminator());
+  BranchOp nextBranch = dyn_cast<BranchOp>(successorDest->getTerminator());
+  llvm::DenseSet<Block *> visited{successorDest};
   while (nextBranch) {
     Block *nextBranchDest = nextBranch.getDest();
-    if (!nextBranchDest || nextBranchDest == currBlock)
-      break;
-    else if (nextBranchDest == successor)
+    if (nextBranchDest == successor)
       return failure();
-    currBlock = nextBranchDest;
-    nextBranch = dyn_cast<BranchOp>(currBlock->getTerminator());
+    else if (visited.contains(nextBranchDest))
+      break;
+    visited.insert(nextBranchDest);
+    nextBranch = dyn_cast<BranchOp>(nextBranchDest->getTerminator());
   }
 
   // Update the operands to the successor. If the branch parent has no
diff --git a/mlir/test/Dialect/ControlFlow/canonicalize.mlir b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
index bf69935a00bf0..bdae567075a56 100644
--- a/mlir/test/Dialect/ControlFlow/canonicalize.mlir
+++ b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
@@ -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
+}
diff --git a/mlir/test/Transforms/control-flow-cycles.mlir b/mlir/test/Transforms/control-flow-cycles.mlir
deleted file mode 100644
index ebb9275bb5cf5..0000000000000
--- a/mlir/test/Transforms/control-flow-cycles.mlir
+++ /dev/null
@@ -1,93 +0,0 @@
-// 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
-}
-
-// 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
-}

>From 63576148476c2f1a3b73c0c2d9c8d5b799672ee5 Mon Sep 17 00:00:00 2001
From: benwu25 <soggysocks206 at gmail.com>
Date: Sun, 28 Sep 2025 00:04:56 -0700
Subject: [PATCH 5/5] [MLIR][CF] Bail out of collapseBranch when any cycle is
 detected (#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.
---
 .../Dialect/ControlFlow/IR/ControlFlowOps.cpp |   6 +-
 .../Dialect/ControlFlow/canonicalize.mlir     | 144 +++++++++---------
 2 files changed, 78 insertions(+), 72 deletions(-)

diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
index dfa83e30ec40a..f1da1a125e9ef 100644
--- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp
+++ b/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());
   }
diff --git a/mlir/test/Dialect/ControlFlow/canonicalize.mlir b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
index bdae567075a56..17f7d28ba59fb 100644
--- a/mlir/test/Dialect/ControlFlow/canonicalize.mlir
+++ b/mlir/test/Dialect/ControlFlow/canonicalize.mlir
@@ -501,10 +501,10 @@ func.func @branchCondProp(%arg0: i1) {
 // CHECK:           cf.br ^bb1
 func.func @cycle_2_blocks() {
   cf.br ^bb1
-  ^bb1:
-    cf.br ^bb2
-  ^bb2:
-    cf.br ^bb1
+^bb1:
+  cf.br ^bb2
+^bb2:
+  cf.br ^bb1
 }
 
 // CHECK-LABEL:   @no_cycle_2_blocks
@@ -512,13 +512,13 @@ func.func @cycle_2_blocks() {
 // 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
+^bb1:
+  cf.br ^bb2
+^bb2:
+  cf.br ^bb3
+^bb3:
+  %ret = arith.constant 1 : i32
+  return %ret : i32
 }
 
 // CHECK-LABEL:   @cycle_4_blocks
@@ -527,14 +527,14 @@ func.func @no_cycle_2_blocks() -> i32 {
 // 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
+^bb1:
+  cf.br ^bb2
+^bb2:
+  cf.br ^bb3
+^bb3:
+  cf.br ^bb4
+^bb4:
+  cf.br ^bb1
 }
 
 // CHECK-LABEL:   @no_cycle_4_blocks
@@ -542,17 +542,17 @@ func.func @cycle_4_blocks() {
 // 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
+^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
@@ -561,16 +561,16 @@ func.func @no_cycle_4_blocks() -> i32 {
 // 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
+^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
@@ -579,50 +579,58 @@ func.func @delayed_3_cycle() {
 // CHECK:           cf.br ^bb1
 func.func @cycle_1_block() {
   cf.br ^bb1
-  ^bb1:
-    cf.br ^bb2
-  ^bb2:
-    cf.br ^bb2
+^bb1:
+  cf.br ^bb2
+^bb2:
+  cf.br ^bb2
 }
 
 // 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:
-    cf.br ^bb2
-  ^bb2:
-    cf.br ^bb3
-  ^bb3:
-    cf.br ^bb4
-  ^bb4:
-    cf.br ^bb3
+^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:           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:
+^bb6:
   cf.br ^bb5 {F}
-  ^bb5:
+^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
+^bb1:
+  cf.br ^bb2 {B}
+^bb2:
+  cf.br ^bb3 {C}
+^bb3:
+  cf.br ^bb4 {D}
+^bb4:
+  cf.br ^bb1 {E}
+^bb7:
+  cf.br ^bb6
 }



More information about the Mlir-commits mailing list