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

[Vijay Kandiah]

================
@@ -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.
----------------
VijayKandiah wrote:

I tried patching both the folder and the pattern for `arith.addi`:

```
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -431,6 +431,9 @@ Value mlir::arith::getZeroConstant(OpBuilder &builder, Location loc,
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::AddIOp::fold(FoldAdaptor adaptor) {
+  if (getLhs() == getResult() || getRhs() == getResult())
+    return {};
+
   // addi(x, 0) -> x
   if (matchPattern(adaptor.getRhs(), m_Zero()))
     return getLhs();
```

```
--- a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
+++ b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
@@ -57,6 +57,9 @@ def TypesMatch : Constraint<CPred<"$0.getType() == $1.getType()">>;
 def IntAttrTypeMatchesResult :
     Constraint<CPred<"$0.getType() == cast<IntegerAttr>($1).getType()">>;
 
+// Check that two values are not the same.
+def AreDistinctValues : Constraint<CPred<"$0 != $1">>;
+
 // Check that a CmpIPredicateAttr is eq or ne.
 def IsEqualityCmpIPredicate :
     Constraint<CPred<"$0.getValue() == arith::CmpIPredicate::eq || "
@@ -72,10 +75,11 @@ def IsEqualityCmpIPredicate :
 // addi(addi(x, c0), c1) -> addi(x, c0 + c1)
 def AddIAddConstant :
     Pat<(Arith_AddIOp:$res
-          (Arith_AddIOp $x, (ConstantLikeMatcher APIntAttr:$c0), $ovf1),
+          (Arith_AddIOp:$inner $x, (ConstantLikeMatcher APIntAttr:$c0), $ovf1),
           (ConstantLikeMatcher APIntAttr:$c1), $ovf2),
         (Arith_AddIOp $x, (Arith_ConstantOp (AddIntAttrs $res, $c0, $c1)),
-            (MergeOverflow $ovf1, $ovf2))>;
+            (MergeOverflow $ovf1, $ovf2)),
+        [(AreDistinctValues $res, $inner)]>;
 
 // addi(subi(x, c0), c1) -> addi(x, c1 - c0)
 def AddISubConstantRHS :
```

This stops the infinite loop at its source and fixes the test case in this PR. However, the problem is that we would have to do the same for other ops like `arith.subi`, `arith.andi`, `arith.ori` and `arith.xori` in `arith`, and ones from other dialects as well like `index.add` and `builtin.unrealized_conversion_cast` (the assertion in #194824). I'm not sure we should go down this route.

As for the cost of a listener-based approach: tracking reachability could very well be expensive and I did not try implementing that. I tried simply setting a flag when the control flow changes in the driver and scanning the operands only once this is set. That's cheap but it is not a reachability check, though it should cover all the unreachable cases. A block can only become unreachable mid-run through a control flow change, which is how the test in this PR gets there. Blocks that are already unreachable are erased before the run starts, so those do not hang anyway.


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