[mlir][Transforms] Fix compile time regression in dialect conversion

[matthias-springer]

The dialect conversion does not directly erase ops that are replaced/erased with a rewriter. Instead, the op stays in place and is erased at the end if the dialect conversion succeeds. However, ops that were replaced/erased are ignored from that point on.

#81757 introduced a compile time regression that made the check whether an op is ignored or not more expensive. Whether an op is ignored or not is queried many times throughout a dialect conversion, so the check must be fast.

After this change, replaced ops are stored in the ignoredOps set. This also simplifies the dialect conversion a bit.

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

The dialect conversion does not directly erase ops that are replaced/erased with a rewriter. Instead, the op stays in place and is erased at the end if the dialect conversion succeeds. However, ops that were replaced/erased are ignored from that point on.

#81757 introduced a compile time regression that made the check whether an op is ignored or not more expensive. Whether an op is ignored or not is queried many times throughout a dialect conversion, so the check must be fast.

After this change, replaced ops are stored in the ignoredOps set. This also simplifies the dialect conversion a bit.

---

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

1 Files Affected:

• (modified) mlir/lib/Transforms/Utils/DialectConversion.cpp (+4-9)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 857b601acbc35a..4165e0a52428f9 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1229,9 +1229,8 @@ LogicalResult ConversionPatternRewriterImpl::remapValues(
 }
 
 bool ConversionPatternRewriterImpl::isOpIgnored(Operation *op) const {
-  // Check to see if this operation was replaced or its parent ignored.
-  return ignoredOps.count(op->getParentOp()) ||
-         hasRewrite<ReplaceOperationRewrite>(rewrites, op);
+  // Check to see if this operation or the parent operation is ignored.
+  return ignoredOps.count(op->getParentOp()) || ignoredOps.count(op);
 }
 
 void ConversionPatternRewriterImpl::markNestedOpsIgnored(Operation *op) {
@@ -1480,12 +1479,7 @@ void ConversionPatternRewriterImpl::notifyOperationInserted(
 void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,
                                                      ValueRange newValues) {
   assert(newValues.size() == op->getNumResults());
-#ifndef NDEBUG
-  for (auto &rewrite : rewrites)
-    if (auto *opReplacement = dyn_cast<ReplaceOperationRewrite>(rewrite.get()))
-      assert(opReplacement->getOperation() != op &&
-             "operation was already replaced");
-#endif // NDEBUG
+  assert(!ignoredOps.contains(op) && "operation was already replaced");
 
   // Track if any of the results changed, e.g. erased and replaced with null.
   bool resultChanged = false;
@@ -1506,6 +1500,7 @@ void ConversionPatternRewriterImpl::notifyOpReplaced(Operation *op,
 
   // Mark this operation as recursively ignored so that we don't need to
   // convert any nested operations.
+  ignoredOps.insert(op);
   markNestedOpsIgnored(op);
 }
 

[qedawkins]

Thanks for the quick fix!

[joker-eph]

Is the parentOp more likely to be ignored? Otherwise why not check the op first before dereferencing the 2 pointers to get the parent?

[matthias-springer]

I had to partially undo this change (#83051) because of downstream code that uses the rewriter API incorrectly: accessing nested regions of an op after replacing the op; the dialect conversion does not crash when doing that because ops are not really erased the very end. But this change caused those nested ops to be ignored. I'm going to fix that code, then I should be able to change this to simply ignoredOps.count(op).