[mlir][Transforms] Greedy pattern rewriter: fix infinite folding loop

[matthias-springer]

Fix an infinite folding loop when the result of a folding is identical to the folded operation. (This can be the case when the folder is faulty or the input IR is invalid.)

Fixes one issue mentioned in #159675 (comment).

[llvmbot]

@llvm/pr-subscribers-mlir-arith

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

Fix an infinite folding loop when the result of a folding is identical to the folded operation.

---

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

2 Files Affected:

• (modified) mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp (+2-1)
• (modified) mlir/test/Transforms/canonicalize.mlir (+9)

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 74e4a822b4fd7..93468dd79808f 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -555,7 +555,8 @@ bool GreedyPatternRewriteDriver::processWorklist() {
           replacements.push_back(constOp->getResult(0));
         }
 
-        if (materializationSucceeded) {
+        if (materializationSucceeded &&
+            !llvm::equal(replacements, op->getResults())) {
           rewriter.replaceOp(op, replacements);
           changed = true;
           LLVM_DEBUG(logSuccessfulFolding(dumpRootOp));
diff --git a/mlir/test/Transforms/canonicalize.mlir b/mlir/test/Transforms/canonicalize.mlir
index 8e02c06a0a293..ed987e555926c 100644
--- a/mlir/test/Transforms/canonicalize.mlir
+++ b/mlir/test/Transforms/canonicalize.mlir
@@ -1248,3 +1248,12 @@ func.func @test_materialize_failure() -> i64 {
   %u = index.castu %const : index to i64
   return %u: i64
 }
+
+// -----
+
+// Make sure that the canonicalizer does not fold infinitely.
+
+// CHECK: %[[c0:.*]] = arith.constant 0 : index
+%c0 = arith.constant 0 : index
+// CHECK: %[[add:.*]] = arith.addi %[[c0]], %[[add]] : index
+%0 = arith.addi %c0, %0 : index

[llvmbot]

@llvm/pr-subscribers-mlir-core

Author: Matthias Springer (matthias-springer)

Changes

Fix an infinite folding loop when the result of a folding is identical to the folded operation.

---

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

2 Files Affected:

• (modified) mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp (+2-1)
• (modified) mlir/test/Transforms/canonicalize.mlir (+9)

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 74e4a822b4fd7..93468dd79808f 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -555,7 +555,8 @@ bool GreedyPatternRewriteDriver::processWorklist() {
           replacements.push_back(constOp->getResult(0));
         }
 
-        if (materializationSucceeded) {
+        if (materializationSucceeded &&
+            !llvm::equal(replacements, op->getResults())) {
           rewriter.replaceOp(op, replacements);
           changed = true;
           LLVM_DEBUG(logSuccessfulFolding(dumpRootOp));
diff --git a/mlir/test/Transforms/canonicalize.mlir b/mlir/test/Transforms/canonicalize.mlir
index 8e02c06a0a293..ed987e555926c 100644
--- a/mlir/test/Transforms/canonicalize.mlir
+++ b/mlir/test/Transforms/canonicalize.mlir
@@ -1248,3 +1248,12 @@ func.func @test_materialize_failure() -> i64 {
   %u = index.castu %const : index to i64
   return %u: i64
 }
+
+// -----
+
+// Make sure that the canonicalizer does not fold infinitely.
+
+// CHECK: %[[c0:.*]] = arith.constant 0 : index
+%c0 = arith.constant 0 : index
+// CHECK: %[[add:.*]] = arith.addi %[[c0]], %[[add]] : index
+%0 = arith.addi %c0, %0 : index

[matthias-springer]

Note: There is still an infinite loop when canonicalizing this IR:

%c1 = arith.constant 1 : index
%0 = arith.addi %c1, %0 : index

[matthias-springer]

This PR does not fix the problem yet. I'm seeing an infinite loop again, not sure why the CI has a different error.

The problem is here:

OpFoldResult arith::AddIOp::fold(FoldAdaptor adaptor) {
  // addi(x, 0) -> x
  if (matchPattern(adaptor.getRhs(), m_Zero())) {

    // THIS CHECK IS NEEDED:
    if (getLhs() == *this) return {};
    
    return getLhs();
  }

Without the above check, we return lhs, which is the addi op itself. Then, in foldSingleResultHook, this is interpreted as "operation folded in-place", triggering a notifyOperationModified and, therefore, infinite looping. (We could check in the greedy pattern rewrite if the op was actually modified, but that's costly.) Note: The multi-result folding function does not have this problem, because it explicitly encodes "op folded in place" by "success result + empty OpFoldResults".

Fixing the folder itself it not enough, the next pattern that causes infinite looping is AddIAddConstant.

I'm not sure what's the right fix here. We likely have many folders/patterns that run into infinite loops in such cases. Should we fix all these? Are they even broken? I am using the addi op in a graph region, and in a way that maybe should be disallowed. (What does %0 = addi %0, %c1 even mean?) Maybe such IR should be disallowed by extra verification?

[joker-eph]

What about in an unreachable block in a CFG? DialectConversion does not visit there, but they could happen in the greedy rewriter.

[matthias-springer]

What about in an unreachable block in a CFG? DialectConversion does not visit there, but they could happen in the greedy rewriter.

I'm not sure how this comment relates to this PR.

[joker-eph]

I was following up on:

I am using the addi op in a graph region, and in a way that maybe should be disallowed. (What does %0 = addi %0, %c1 even mean?) Maybe such IR should be disallowed by extra verification?

%0 = addi %0, %c1 is a valid IR construct in a CFG region in an unreachable block, it's not specific to graph regions.

[matthias-springer]

Ok, then it cannot be fixed in the verifier.

Then we have to fix the folder. The API of the single-result folder is: if the folder returns the op result of the same op, the op was folded in-place. This is violated here. (I updated the PR. But only in one place. There are many places like that.)

Single-result folder called here:

  template <typename ConcreteOpT>
  static LogicalResult
  foldSingleResultHook(Operation *op, ArrayRef<Attribute> operands,
                       SmallVectorImpl<OpFoldResult> &results) {
...
    // If the fold failed or was in-place, try to fold the traits of the
    // operation.
    if (!result ||
        llvm::dyn_cast_if_present<Value>(result) == op->getResult(0)) {

Also, it doesn't fix the issue yet, because now this canonicalization pattern gets into an infinite loop. Unfortunately, it's written in TableGen and there may be no way to fix this without rewriting it in C++:

// addi(addi(x, c0), c1) -> addi(x, c0 + c1)
def AddIAddConstant :
    Pat<(Arith_AddIOp:$res
          (Arith_AddIOp $x, (ConstantLikeMatcher APIntAttr:$c0), $ovf1),
          (ConstantLikeMatcher APIntAttr:$c1), $ovf2),
        (Arith_AddIOp $x, (Arith_ConstantOp (AddIntAttrs $res, $c0, $c1)),
            (MergeOverflow $ovf1, $ovf2))>;

[kuhar]

Could we detect this in: Operation::fold (or similar) such that we never enter Op::fold on operations that feed into themselves?

[joker-eph]

The issue is that this isn't necessarily about the immediate operands:

%sub = subi %add, %cst
%add = addi %div , %cst

A folder on addi here could replace itself with itself.