[Mlir-commits] [mlir] [mlir][cf] Fix non-termination in simplifyBrToBlockWithSinglePred (PR #221037)

[Vijay Kandiah]

================
@@ -687,6 +687,33 @@ func.func @no_merge_self_arg_loop(%step: i1) -> i1 {
   return %result : i1
 }
 
+// -----
+
+// Verify that a block is not merged into its sole predecessor when a branch
+// operand only derives from a block argument of the destination: substituting
+// %next for %iv would define %next in terms of itself. Folding the constant
+// compare is what leaves ^latch as ^header's only predecessor.
+
+// CHECK-LABEL: @no_merge_derived_arg_loop
+//  CHECK-NEXT:   return
+func.func @no_merge_derived_arg_loop() {
+  %c0 = arith.constant 0 : i32
+  %c1 = arith.constant 1 : i32
+  %c2 = arith.constant 2 : i32
+  %enter = arith.cmpi sgt, %c2, %c0 : i32
+  cf.cond_br %enter, ^exit, ^header(%c2 : i32)
+^header(%iv: i32):
+  %done = arith.cmpi eq, %iv, %c2 : i32
+  cf.cond_br %done, ^exit, ^latch
+^latch:
+  %next = arith.addi %iv, %c1 : i32
+  cf.br ^header(%next : i32)
----------------
VijayKandiah wrote:

Thanks for the review! I missed that MLIR does not enforce SSA dominance in unreachable code. So the merge happening here does produce valid IR and we shouldn't need the restriction in this change. I agree this is a greedy driver issue rather than a `cf` pattern issue.

I opened https://github.com/llvm/llvm-project/pull/221152 with a greedy driver fix. Closing this one.

For this test case, once the first `cf.cond_br` is constant folded, `^latch` becomes `^header`'s only predecessor and `simplifyBrToBlockWithSinglePred` merges them, substituting `%next` for `%iv`. Since `%next = arith.addi %iv, %c1`, the result is `%0 = arith.addi %0, ...`, which the driver then keeps folding without ever converging.

That is the infinite loop mentioned in the greedy driver at https://github.com/llvm/llvm-project/blob/85081add2d8b401e5f5ee67cf4af11531560f3bf/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp#L892-L899 `eraseUnreachableBlocks` runs at the top of each iteration, so it catches blocks that are unreachable in the input, but not one that becomes unreachable during `processWorklist()`, which never returns. The fix I shared in https://github.com/llvm/llvm-project/pull/221152 skips operations that use their own result, leaving them to hit `eraseUnreachableBlocks` on the next iteration.

I made the test case for this bug a bit smaller in the new PR. However, the first `cond_br` (which is constant folded) has to stay in the input. Pre-running the canonicalizer like you suggested leaves the blocks unreachable on entry, which is the case that already works, so the bug doesn't reproduce. A block must become unreachable during the greedy rewriter's `processWorklist()` for this bug to be reproduced.

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