Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

Conversation

matthias-springer
Copy link
Member

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.

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 llvmbot added mlir:core MLIR Core Infrastructure mlir labels Feb 26, 2024
@llvmbot
Copy link
Member

llvmbot commented Feb 26, 2024

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

Copy link
Contributor

@qedawkins qedawkins left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the quick fix!

@matthias-springer matthias-springer merged commit 45732b6 into main Feb 26, 2024
@matthias-springer matthias-springer deleted the users/matthias-springer/dialect_conv_perf_regression branch February 26, 2024 16:40
qedawkins pushed a commit to iree-org/llvm-project that referenced this pull request Feb 26, 2024
…lvm#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.

llvm#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.
matthias-springer added a commit that referenced this pull request Feb 26, 2024
#83023 fixed a performance regression related to "ignored" ops. This
broke some downstream projects that access ops after they were replaced
(an API violation). This change restores the original behavior before
#83023 (but without the performance regression), to give downstream
users more time to fix their code.
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);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

matthias-springer added a commit to iree-org/llvm-project that referenced this pull request Feb 28, 2024
llvm#83023 fixed a performance regression related to "ignored" ops. This
broke some downstream projects that access ops after they were replaced
(an API violation). This change restores the original behavior before
llvm#83023 (but without the performance regression), to give downstream
users more time to fix their code.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
mlir:core MLIR Core Infrastructure mlir
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants