From aed78e68e9886abf50c68a0d34559b4746c02c31 Mon Sep 17 00:00:00 2001 From: benwu25 Date: Thu, 25 Sep 2025 13:46:10 -0700 Subject: [PATCH 1/8] [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(successorDest->getTerminator()); + while (nextBranch) { + Block *nextBranchDest = nextBranch.getDest(); + if (!nextBranchDest) + break; + else if (nextBranchDest == successor) + return failure(); + nextBranch = dyn_cast(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 Date: Thu, 25 Sep 2025 14:39:57 -0700 Subject: [PATCH 2/8] [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 Date: Thu, 25 Sep 2025 21:11:12 -0700 Subject: [PATCH 3/8] [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(successorDest->getTerminator()); + Block *currBlock = successorDest; + BranchOp nextBranch = dyn_cast(currBlock->getTerminator()); while (nextBranch) { Block *nextBranchDest = nextBranch.getDest(); - if (!nextBranchDest) + if (!nextBranchDest || nextBranchDest == currBlock) break; else if (nextBranchDest == successor) return failure(); - nextBranch = dyn_cast(nextBranchDest->getTerminator()); + currBlock = nextBranchDest; + nextBranch = dyn_cast(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 Date: Fri, 26 Sep 2025 13:33:50 -0700 Subject: [PATCH 4/8] [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(currBlock->getTerminator()); + BranchOp nextBranch = dyn_cast(successorDest->getTerminator()); + llvm::DenseSet 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(currBlock->getTerminator()); + else if (visited.contains(nextBranchDest)) + break; + visited.insert(nextBranchDest); + nextBranch = dyn_cast(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 Date: Sun, 28 Sep 2025 00:04:56 -0700 Subject: [PATCH 5/8] [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(successorDest->getTerminator()); - llvm::DenseSet visited{successorDest}; + llvm::DenseSet 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(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 } From e8135aabd5fe3a909819af29c655bd65fc3c1873 Mon Sep 17 00:00:00 2001 From: benwu25 Date: Sun, 28 Sep 2025 00:48:05 -0700 Subject: [PATCH 6/8] [MLIR][CF] Do not bail out early in collapseBranch to avoid regression (#159743) Bailing out as soon as a cycle is detected causes a regression in flang, where the test expects more optimizations than we allowed on an infinite loop. We may be able to adjust this testcase. --- mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp index f1da1a125e9ef..dfa83e30ec40a 100644 --- a/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp +++ b/mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp @@ -124,11 +124,13 @@ static LogicalResult collapseBranch(Block *&successor, return failure(); // Don't try to collapse branches which participate in a cycle. BranchOp nextBranch = dyn_cast(successorDest->getTerminator()); - llvm::DenseSet visited{successor, successorDest}; + llvm::DenseSet visited{successorDest}; while (nextBranch) { Block *nextBranchDest = nextBranch.getDest(); - if (visited.contains(nextBranchDest)) + if (nextBranchDest == successor) return failure(); + else if (visited.contains(nextBranchDest)) + break; visited.insert(nextBranchDest); nextBranch = dyn_cast(nextBranchDest->getTerminator()); } From 9c27d37cacbd5058c6162cf357292d2be02b6990 Mon Sep 17 00:00:00 2001 From: benwu25 Date: Sun, 28 Sep 2025 01:06:43 -0700 Subject: [PATCH 7/8] [MLIR][CF] Update broken test (#159743) I forgot to update this test with my previous commit. (e8135aabd5fe3a909819af29c655bd65fc3c1873) --- mlir/test/Dialect/ControlFlow/canonicalize.mlir | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/mlir/test/Dialect/ControlFlow/canonicalize.mlir b/mlir/test/Dialect/ControlFlow/canonicalize.mlir index 17f7d28ba59fb..774fec65620fb 100644 --- a/mlir/test/Dialect/ControlFlow/canonicalize.mlir +++ b/mlir/test/Dialect/ControlFlow/canonicalize.mlir @@ -587,13 +587,9 @@ func.func @cycle_1_block() { // CHECK-LABEL: @unsimplified_cycle_1 // CHECK-SAME: %[[ARG0:.*]]: i1) { -// CHECK: cf.cond_br %[[ARG0]], ^bb1, ^bb2 +// CHECK: cf.br ^bb1 // CHECK: ^bb1: -// CHECK: cf.br ^bb2 -// CHECK: ^bb2: -// CHECK: cf.br ^bb3 -// CHECK: ^bb3: -// CHECK: cf.br ^bb3 +// CHECK: cf.br ^bb1 func.func @unsimplified_cycle_1(%c : i1) { cf.cond_br %c, ^bb1, ^bb2 ^bb1: @@ -610,13 +606,9 @@ func.func @unsimplified_cycle_1(%c : i1) { // 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 +// CHECK: ^bb1: +// CHECK: cf.br ^bb1 {E} func.func @unsimplified_cycle_2(%c : i1) { cf.cond_br %c, ^bb6, ^bb7 ^bb6: From 13dfab586a40172ed78104ed4b097b0aca0239d4 Mon Sep 17 00:00:00 2001 From: benwu25 Date: Mon, 29 Sep 2025 00:58:45 -0700 Subject: [PATCH 8/8] [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. --- .../Lower/OpenMP/infinite-loop-in-construct.f90 | 6 ++++-- .../Dialect/ControlFlow/IR/ControlFlowOps.cpp | 6 ++---- mlir/test/Dialect/ControlFlow/canonicalize.mlir | 16 ++++++++++++---- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/flang/test/Lower/OpenMP/infinite-loop-in-construct.f90 b/flang/test/Lower/OpenMP/infinite-loop-in-construct.f90 index 16b400a231860..f02d0e5ccc53c 100644 --- a/flang/test/Lower/OpenMP/infinite-loop-in-construct.f90 +++ b/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) 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(successorDest->getTerminator()); - llvm::DenseSet visited{successorDest}; + llvm::DenseSet 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(nextBranchDest->getTerminator()); } diff --git a/mlir/test/Dialect/ControlFlow/canonicalize.mlir b/mlir/test/Dialect/ControlFlow/canonicalize.mlir index 774fec65620fb..17f7d28ba59fb 100644 --- a/mlir/test/Dialect/ControlFlow/canonicalize.mlir +++ b/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: