[mlir][IR] Add notifyBlockRemoved callback to listener
#78306
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir-core @llvm/pr-subscribers-mlir Author: Matthias Springer (matthias-springer) ChangesThere is already a "block inserted" notification (in The purpose of this change is to make the listener API more mature. There is currently a gap between what kind of IR changes can be made and what IR changes can be listened to. At the moment, the only way to inform listeners about "block removal" is to send a manual Full diff: https://github.com/llvm/llvm-project/pull/78306.diff 3 Files Affected:
diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
index 9b4fa65bff49e1..6c5d317a79b7ab 100644
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -402,6 +402,10 @@ class RewriterBase : public OpBuilder {
Listener()
: OpBuilder::Listener(ListenerBase::Kind::RewriterBaseListener) {}
+ /// Notify the listener that the specified block is about to be erased.
+ /// At this point, the block has zero uses.
+ virtual void notifyBlockRemoved(Block *block) {}
+
/// Notify the listener that the specified operation was modified in-place.
virtual void notifyOperationModified(Operation *op) {}
@@ -452,6 +456,10 @@ class RewriterBase : public OpBuilder {
void notifyBlockCreated(Block *block) override {
listener->notifyBlockCreated(block);
}
+ void notifyBlockRemoved(Block *block) override {
+ if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+ rewriteListener->notifyBlockRemoved(block);
+ }
void notifyOperationModified(Operation *op) override {
if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
rewriteListener->notifyOperationModified(op);
diff --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index abb65bc3c38f22..b63baf330c8645 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -1114,7 +1114,7 @@ static scf::ForOp replaceForOpWithNewSignature(RewriterBase &rewriter,
scf::ForOp newLoop = rewriter.create<scf::ForOp>(
loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(), loop.getStep(),
operands);
- newLoop.getBody()->erase();
+ rewriter.eraseBlock(newLoop.getBody());
newLoop.getRegion().getBlocks().splice(
newLoop.getRegion().getBlocks().begin(), loop.getRegion().getBlocks());
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 5e788cdb4897d3..de226f40cb3e9a 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -244,7 +244,7 @@ void RewriterBase::eraseOp(Operation *op) {
for (BlockArgument bbArg : b->getArguments())
bbArg.dropAllUses();
b->dropAllUses();
- b->erase();
+ eraseBlock(b);
}
}
}
@@ -256,10 +256,17 @@ void RewriterBase::eraseOp(Operation *op) {
}
void RewriterBase::eraseBlock(Block *block) {
+ assert(block->use_empty() && "expected 'block' to have no uses");
+
for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block))) {
assert(op.use_empty() && "expected 'op' to have no uses");
eraseOp(&op);
}
+
+ // Notify the listener that the block is about to be removed.
+ if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))
+ rewriteListener->notifyBlockRemoved(block);
+
block->erase();
}
@@ -311,7 +318,7 @@ void RewriterBase::inlineBlockBefore(Block *source, Block *dest,
// Move operations from the source block to the dest block and erase the
// source block.
dest->getOperations().splice(before, source->getOperations());
- source->erase();
+ eraseBlock(source);
}
void RewriterBase::inlineBlockBefore(Block *source, Operation *op,
|
|
@llvm/pr-subscribers-mlir-gpu Author: Matthias Springer (matthias-springer) ChangesThere is already a "block inserted" notification (in The purpose of this change is to make the listener API more mature. There is currently a gap between what kind of IR changes can be made and what IR changes can be listened to. At the moment, the only way to inform listeners about "block removal" is to send a manual Full diff: https://github.com/llvm/llvm-project/pull/78306.diff 3 Files Affected:
diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
index 9b4fa65bff49e1..6c5d317a79b7ab 100644
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -402,6 +402,10 @@ class RewriterBase : public OpBuilder {
Listener()
: OpBuilder::Listener(ListenerBase::Kind::RewriterBaseListener) {}
+ /// Notify the listener that the specified block is about to be erased.
+ /// At this point, the block has zero uses.
+ virtual void notifyBlockRemoved(Block *block) {}
+
/// Notify the listener that the specified operation was modified in-place.
virtual void notifyOperationModified(Operation *op) {}
@@ -452,6 +456,10 @@ class RewriterBase : public OpBuilder {
void notifyBlockCreated(Block *block) override {
listener->notifyBlockCreated(block);
}
+ void notifyBlockRemoved(Block *block) override {
+ if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+ rewriteListener->notifyBlockRemoved(block);
+ }
void notifyOperationModified(Operation *op) override {
if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
rewriteListener->notifyOperationModified(op);
diff --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index abb65bc3c38f22..b63baf330c8645 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -1114,7 +1114,7 @@ static scf::ForOp replaceForOpWithNewSignature(RewriterBase &rewriter,
scf::ForOp newLoop = rewriter.create<scf::ForOp>(
loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(), loop.getStep(),
operands);
- newLoop.getBody()->erase();
+ rewriter.eraseBlock(newLoop.getBody());
newLoop.getRegion().getBlocks().splice(
newLoop.getRegion().getBlocks().begin(), loop.getRegion().getBlocks());
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 5e788cdb4897d3..de226f40cb3e9a 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -244,7 +244,7 @@ void RewriterBase::eraseOp(Operation *op) {
for (BlockArgument bbArg : b->getArguments())
bbArg.dropAllUses();
b->dropAllUses();
- b->erase();
+ eraseBlock(b);
}
}
}
@@ -256,10 +256,17 @@ void RewriterBase::eraseOp(Operation *op) {
}
void RewriterBase::eraseBlock(Block *block) {
+ assert(block->use_empty() && "expected 'block' to have no uses");
+
for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block))) {
assert(op.use_empty() && "expected 'op' to have no uses");
eraseOp(&op);
}
+
+ // Notify the listener that the block is about to be removed.
+ if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))
+ rewriteListener->notifyBlockRemoved(block);
+
block->erase();
}
@@ -311,7 +318,7 @@ void RewriterBase::inlineBlockBefore(Block *source, Block *dest,
// Move operations from the source block to the dest block and erase the
// source block.
dest->getOperations().splice(before, source->getOperations());
- source->erase();
+ eraseBlock(source);
}
void RewriterBase::inlineBlockBefore(Block *source, Operation *op,
|
There is already a "block inserted" notification (in `OpBuilder::Listener`), so there should also be a "block removed" notification. The purpose of this change is to make the listener API more mature. There is currently a gap between what kind of IR changes can be made and what IR changes can be listened to. At the moment, the only way to inform listeners about "block removal" is to send a manual `notifyOperationModified` for the parent op (e.g., by wrapping the `eraseBlock(b)` method call in `updateRootInPlace(b->getParentOp())`). This tells the listener that *something* has changed, but it is somewhat of an API abuse.
matthias-springer
force-pushed
the
notify_block_removed
branch
from
0650628 to
d661ff6
Compare
matthias-springer
added a commit
to matthias-springer/llvm-project
that referenced
this pull request
Jan 16, 2024
The greedy pattern rewrite driver has multiple "expensive checks" to detect invalid rewrite pattern API usage. As part of these checks, it computes fingerprints for every op that is in scope, and compares the fingerprints before and after an attempted pattern application. Until now, each computed fingerprint took into account all nested operations. That is quite expensive because it walks the entire IR subtree. It is also redudant in the expensive checks because we already compute a fingerprint for every op. This commit significantly improves the running time of the "expensive checks" in the greedy pattern rewrite driver. Depends on llvm#78306. Review only the top commit.
matthias-springer
added a commit
to matthias-springer/llvm-project
that referenced
this pull request
Jan 21, 2024
The greedy pattern rewrite driver has multiple "expensive checks" to detect invalid rewrite pattern API usage. As part of these checks, it computes fingerprints for every op that is in scope, and compares the fingerprints before and after an attempted pattern application. Until now, each computed fingerprint took into account all nested operations. That is quite expensive because it walks the entire IR subtree. It is also redudant in the expensive checks because we already compute a fingerprint for every op. This commit significantly improves the running time of the "expensive checks" in the greedy pattern rewrite driver. Depends on llvm#78306. Review only the top commit.
matthias-springer
added a commit
to matthias-springer/llvm-project
that referenced
this pull request
Jan 31, 2024
The greedy pattern rewrite driver has multiple "expensive checks" to detect invalid rewrite pattern API usage. As part of these checks, it computes fingerprints for every op that is in scope, and compares the fingerprints before and after an attempted pattern application. Until now, each computed fingerprint took into account all nested operations. That is quite expensive because it walks the entire IR subtree. It is also redudant in the expensive checks because we already compute a fingerprint for every op. This commit significantly improves the running time of the "expensive checks" in the greedy pattern rewrite driver. Depends on llvm#78306. Review only the top commit.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There is already a "block inserted" notification (in
OpBuilder::Listener), so there should also be a "block removed" notification.The purpose of this change is to make the listener API more mature. There is currently a gap between what kind of IR changes can be made and what IR changes can be listened to. At the moment, the only way to inform listeners about "block removal" is to send a manual
notifyOperationModifiedfor the parent op (e.g., by wrapping theeraseBlock(b)method call inupdateRootInPlace(b->getParentOp())). This tells the listener that something has changed, but it is somewhat of an API abuse.