[mlir] GreedyPatternRewriteDriver: Support optional Listener

[springerm]

Summary

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

Depends On: D143339

[springerm]

rebase

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

public:
   /// * GreedyRewriteStrictness::AnyOp: No ops are excluded.
   /// * GreedyRewriteStrictness::ExistingAndNewOps: Only pre-existing ops (that
   ///   were on the worklist at the very beginning) and newly created ops are
   ///   enqueued. All other ops are excluded.
   /// * GreedyRewriteStrictness::ExistingOps: Only pre-existing ops (that were
   ///   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;
 };
 
 //===----------------------------------------------------------------------===//
 // applyPatternsGreedily
 //===----------------------------------------------------------------------===//
 
 /// Rewrite ops in the given region, which must be isolated from above, by
 /// repeatedly applying the highest benefit patterns in a greedy work-list

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

private:
   /// be re-added to the worklist. This function should be called when an
   /// operation is modified or removed, as it may trigger further
   /// simplifications.
   void addOperandsToWorklist(ValueRange operands);
 
   /// 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,
                      function_ref<void(Diagnostic &)> reasonCallback) override;
 
   /// If the specified operation is in the worklist, remove it.
   void removeFromWorklist(Operation *op);
 
void GreedyPatternRewriteDriver::removeFromWorklist(Operation *op) {
   auto it = worklistMap.find(op);
   if (it != worklistMap.end()) {
     assert(worklist[it->second] == op && "malformed worklist data structure");
     worklist[it->second] = nullptr;
     worklistMap.erase(it);
   }
 }
 
+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);
 }
 
 void GreedyPatternRewriteDriver::finalizeRootUpdate(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Modified: '" << op->getName() << "'(" << op
void GreedyPatternRewriteDriver::addOperandsToWorklist(ValueRange operands) {
   }
 }
 
 void GreedyPatternRewriteDriver::notifyOperationRemoved(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Erase   : '" << op->getName() << "'(" << op
                        << ")\n";
   });
+  if (config.listener)
+    config.listener->notifyOperationRemoved(op);
 
   addOperandsToWorklist(op->getOperands());
   op->walk([this](Operation *operation) {
     removeFromWorklist(operation);
     folder.notifyRemoval(operation);
   });
 
   if (config.strictMode != GreedyRewriteStrictness::AnyOp)
     strictModeFilteredOps.erase(op);
 }
 
 void GreedyPatternRewriteDriver::notifyOperationReplaced(
     Operation *op, ValueRange replacement) {
   LLVM_DEBUG({
     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);
 }
 
 LogicalResult GreedyPatternRewriteDriver::notifyMatchFailure(
     Location loc, function_ref<void(Diagnostic &)> reasonCallback) {
   LLVM_DEBUG({
     Diagnostic diag(loc, DiagnosticSeverity::Remark);
     reasonCallback(diag);
     logger.startLine() << "** Failure : " << diag.str() << "\n";
   });
+  if (config.listener)
+    return config.listener->notifyMatchFailure(loc, reasonCallback);
   return failure();
 }
 
 //===----------------------------------------------------------------------===//
 // RegionPatternRewriteDriver
 //===----------------------------------------------------------------------===//
 
 namespace {