[Mlir-commits] [mlir] [mlir][Transforms] Fix compile time regression in dialect conversion (PR #83023)

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/83023

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.

>From ed57433cfd3d151c7c4e57d5712aa21ea6a9373d Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Mon, 26 Feb 2024 15:54:11 +0000
Subject: [PATCH] [mlir][Transforms] Fix compile time regression in dialect
 conversion

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.
---
 mlir/lib/Transforms/Utils/DialectConversion.cpp | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

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);
 }