[Mlir-commits] [mlir] [mlir] Fix non-termination when a block becomes unreachable during rewriting (PR #221152)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Vijay Kandiah (VijayKandiah)

<details>
<summary>Changes</summary>

`GreedyPatternRewriteDriver` erases unreachable blocks at the top of each iteration here https://github.com/llvm/llvm-project/blob/85081add2d8b401e5f5ee67cf4af11531560f3bf/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp#L892-L899


This handles blocks that are unreachable in the input, but not a block that becomes unreachable *during* `processWorklist()`, which then never returns to give `eraseUnreachableBlocks` another turn. 

Here's a reproducer:

```mlir
// test.mlir
func.func @<!-- -->loop(%cond: i1, %init: i32) {
  %true = arith.constant true
  %c1 = arith.constant 1 : i32
  cf.cond_br %true, ^exit, ^header(%init : i32)
^header(%iv: i32):
  cf.cond_br %cond, ^exit, ^latch
^latch:
  %next = arith.addi %iv, %c1 : i32
  cf.br ^header(%next : i32)
^exit:
  return
}
```

`mlir-opt test.mlir -canonicalize` never terminates: folding `%true` removes the only entry edge into `^header`, so the loop becomes unreachable while the worklist is still running.

This PR skips operations that use their own result, leaving them to hit `eraseUnreachableBlocks` on the next iteration. Such operations are legal in unreachable code.

Assisted-by: Claude Opus 5 (Cursor)

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


2 Files Affected:

- (modified) mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp (+9) 
- (modified) mlir/test/Transforms/canonicalize-dce.mlir (+20) 


``````````diff
diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index b71194e783ba0..8c5c39992d512 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -482,6 +482,15 @@ bool GreedyPatternRewriteDriver::processWorklist() {
       }
     });
 
+    // Folding an operation that uses its own result can loop forever, so
+    // skip it. Such operations are legal in unreachable code.
+    if (llvm::any_of(op->getOperands(), [op](Value operand) {
+          return operand.getDefiningOp() == op;
+        })) {
+      LLVM_DEBUG(logResultWithLine("failure", "operation uses its own result"));
+      continue;
+    }
+
     // If the operation is trivially dead - remove it.
     if (isOpTriviallyDead(op)) {
       rewriter.eraseOp(op);
diff --git a/mlir/test/Transforms/canonicalize-dce.mlir b/mlir/test/Transforms/canonicalize-dce.mlir
index 84631947970de..57b5832c71dbd 100644
--- a/mlir/test/Transforms/canonicalize-dce.mlir
+++ b/mlir/test/Transforms/canonicalize-dce.mlir
@@ -187,3 +187,23 @@ func.func @f() {
   %0 = "test.test_effects_result"() : () -> i32
   return
 }
+
+// -----
+
+// Test case: Terminate when a block becomes unreachable while patterns run.
+
+// CHECK:      func @f
+// CHECK-NEXT:   return
+
+func.func @f(%cond: i1, %init: i32) {
+  %true = arith.constant true
+  %c1 = arith.constant 1 : i32
+  cf.cond_br %true, ^exit, ^header(%init : i32)
+^header(%iv: i32):
+  cf.cond_br %cond, ^exit, ^latch
+^latch:
+  %next = arith.addi %iv, %c1 : i32
+  cf.br ^header(%next : i32)
+^exit:
+  return
+}

``````````

</details>


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