[Mlir-commits] [mlir] 279c1d2 - [mlir] GreedyPatternRewriteDriver: Support optional Listener

[Matthias Springer]

Author: Matthias Springer
Date: 2023-02-22T10:32:16+01:00
New Revision: 279c1d2ba7fab2b7f84bbe58b542806a21596603

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

LOG: [mlir] GreedyPatternRewriteDriver: Support optional Listener

Allow an optional `RewriterBase::Listener` to be attached to greedy pattern rewrites, so that clients can listen for IR modifications.

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h b/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
index 423221dd80da..2c5c63e0ac4d 100644
--- a/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
+++ b/mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h
@@ -76,6 +76,9 @@ class GreedyRewriteConfig {
   ///   were on the worklist at the very beginning) enqueued. All other ops are
   ///   excluded.
   GreedyRewriteStrictness strictMode = GreedyRewriteStrictness::AnyOp;
+
+  /// An optional listener that should be notified about IR modifications.
+  RewriterBase::Listener *listener = nullptr;
 };
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index adf8b5121ab9..89088d53400e 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -103,6 +103,9 @@ class GreedyPatternRewriteDriver : public PatternRewriter,
   /// Pop the next operation from the worklist.
   Operation *popFromWorklist();
 
+  /// Notify the driver that the given block was created.
+  void notifyBlockCreated(Block *block) override;
+
   /// For debugging only: Notify the driver of a pattern match failure.
   LogicalResult
   notifyMatchFailure(Location loc,
@@ -315,11 +318,18 @@ void GreedyPatternRewriteDriver::removeFromWorklist(Operation *op) {
   }
 }
 
+void GreedyPatternRewriteDriver::notifyBlockCreated(Block *block) {
+  if (config.listener)
+    config.listener->notifyBlockCreated(block);
+}
+
 void GreedyPatternRewriteDriver::notifyOperationInserted(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Insert  : '" << op->getName() << "'(" << op
                        << ")\n";
   });
+  if (config.listener)
+    config.listener->notifyOperationInserted(op);
   if (config.strictMode == GreedyRewriteStrictness::ExistingAndNewOps)
     strictModeFilteredOps.insert(op);
   addToWorklist(op);
@@ -352,6 +362,8 @@ void GreedyPatternRewriteDriver::notifyOperationRemoved(Operation *op) {
     logger.startLine() << "** Erase   : '" << op->getName() << "'(" << op
                        << ")\n";
   });
+  if (config.listener)
+    config.listener->notifyOperationRemoved(op);
 
   addOperandsToWorklist(op->getOperands());
   op->walk([this](Operation *operation) {
@@ -369,6 +381,8 @@ void GreedyPatternRewriteDriver::notifyOperationReplaced(
     logger.startLine() << "** Replace : '" << op->getName() << "'(" << op
                        << ")\n";
   });
+  if (config.listener)
+    config.listener->notifyOperationReplaced(op, replacement);
   for (auto result : op->getResults())
     for (auto *user : result.getUsers())
       addToWorklist(user);
@@ -381,6 +395,8 @@ LogicalResult GreedyPatternRewriteDriver::notifyMatchFailure(
     reasonCallback(diag);
     logger.startLine() << "** Failure : " << diag.str() << "\n";
   });
+  if (config.listener)
+    return config.listener->notifyMatchFailure(loc, reasonCallback);
   return failure();
 }