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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: None (benwu25)

<details>
<summary>Changes</summary>

This PR addresses #<!-- -->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.

Could someone help review and merge? Thanks!

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


2 Files Affected:

- (modified) mlir/lib/Dialect/ControlFlow/IR/ControlFlowOps.cpp (+10) 
- (added) mlir/test/Transforms/control-flow-cycles.mlir (+81) 


``````````diff
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..a3f67ccc5d0c7
--- /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:           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
+}

``````````

</details>


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