[Mlir-commits] [mlir] bafc4df - [mlir] RewriterBase::Listener: Add notifyOperationModified callback

[Matthias Springer]

Author: Matthias Springer
Date: 2023-02-22T10:41:57+01:00
New Revision: bafc4dfc764587de577a354e459ea351b43fdb43

URL: https://github.com/llvm/llvm-project/commit/bafc4dfc764587de577a354e459ea351b43fdb43
DIFF: https://github.com/llvm/llvm-project/commit/bafc4dfc764587de577a354e459ea351b43fdb43.diff

LOG: [mlir] RewriterBase::Listener: Add notifyOperationModified callback

This callback is triggered by `finalizeRootUpdate`. This allows listeners to listen for in-place op modifications without creating a new RewriterBase subclass.

Differential Revision: https://reviews.llvm.org/D143380

Added: 
    

Modified: 
    mlir/include/mlir/IR/PatternMatch.h
    mlir/lib/IR/PatternMatch.cpp
    mlir/lib/Transforms/Utils/DialectConversion.cpp
    mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
index 7845cb19f612..daade923034a 100644
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -402,6 +402,9 @@ class RewriterBase : public OpBuilder {
     Listener()
         : OpBuilder::Listener(ListenerBase::Kind::RewriterBaseListener) {}
 
+    /// Notify the listener that the specified operation was modified in-place.
+    virtual void notifyOperationModified(Operation *op) {}
+
     /// Notify the listener that the specified operation is about to be replaced
     /// with the set of values potentially produced by new operations. This is
     /// called before the uses of the operation have been changed.
@@ -514,7 +517,7 @@ class RewriterBase : public OpBuilder {
   /// This method is used to signal the end of a root update on the given
   /// operation. This can only be called on operations that were provided to a
   /// call to `startRootUpdate`.
-  virtual void finalizeRootUpdate(Operation *op) {}
+  virtual void finalizeRootUpdate(Operation *op);
 
   /// This method cancels a pending root update. This can only be called on
   /// operations that were provided to a call to `startRootUpdate`.

diff  --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 10baea61d9a4..958d82ae5bcb 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -294,6 +294,12 @@ void RewriterBase::eraseBlock(Block *block) {
   block->erase();
 }
 
+void RewriterBase::finalizeRootUpdate(Operation *op) {
+  // Notify the listener that the operation was modified.
+  if (auto *rewriteListener = dyn_cast_if_present<Listener>(listener))
+    rewriteListener->notifyOperationModified(op);
+}
+
 /// Merge the operations of block 'source' into the end of block 'dest'.
 /// 'source's predecessors must be empty or only contain 'dest`.
 /// 'argValues' is used to replace the block arguments of 'source' after

diff  --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 0d78362da7f0..52ebd54e1bee 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1654,6 +1654,7 @@ void ConversionPatternRewriter::startRootUpdate(Operation *op) {
 }
 
 void ConversionPatternRewriter::finalizeRootUpdate(Operation *op) {
+  PatternRewriter::finalizeRootUpdate(op);
   // There is nothing to do here, we only need to track the operation at the
   // start of the update.
 #ifndef NDEBUG

diff  --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 89088d53400e..15495efca09a 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -54,7 +54,7 @@ class GreedyPatternRewriteDriver : public PatternRewriter,
 
   /// Notify the driver that the specified operation may have been modified
   /// in-place. The operation is added to the worklist.
-  void finalizeRootUpdate(Operation *op) override;
+  void notifyOperationModified(Operation *op) override;
 
   /// Notify the driver that the specified operation was inserted. Update the
   /// worklist as needed: The operation is enqueued depending on scope and
@@ -335,7 +335,7 @@ void GreedyPatternRewriteDriver::notifyOperationInserted(Operation *op) {
   addToWorklist(op);
 }
 
-void GreedyPatternRewriteDriver::finalizeRootUpdate(Operation *op) {
+void GreedyPatternRewriteDriver::notifyOperationModified(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Modified: '" << op->getName() << "'(" << op
                        << ")\n";