[mlir][core] Add an MLIR "pattern catalog" generator

mlir/include/mlir/IR/PatternMatch.h

@@ -475,6 +475,25 @@ class RewriterBase : public OpBuilder {
     RewriterBase::Listener *rewriteListener;
   };
 
+  /// A listener that logs notification events to llvm::dbgs() before
+  /// forwarding to the base listener.
+  struct PatternLoggingListener : public RewriterBase::ForwardingListener {
+    PatternLoggingListener(OpBuilder::Listener *listener, StringRef patternName)
+        : RewriterBase::ForwardingListener(listener), patternName(patternName) {
+    }
+
+    void notifyOperationInserted(Operation *op, InsertPoint previous) override;
+    void notifyOperationModified(Operation *op) override;
+    void notifyOperationReplaced(Operation *op, Operation *newOp) override;
+    void notifyOperationReplaced(Operation *op,
+                                 ValueRange replacement) override;
+    void notifyOperationErased(Operation *op) override;
+    void notifyPatternBegin(const Pattern &pattern, Operation *op) override;
+
+  private:
+    StringRef patternName;
+  };
+
   /// Move the blocks that belong to "region" before the given position in
   /// another region "parent". The two regions must be different. The caller
   /// is responsible for creating or updating the operation transferring flow

mlir/lib/IR/CMakeLists.txt

@@ -29,6 +29,7 @@ add_mlir_library(MLIRIR
   ODSSupport.cpp
   Operation.cpp
   OperationSupport.cpp
+  PatternLoggingListener.cpp
   PatternMatch.cpp
   Region.cpp
   RegionKindInterface.cpp

mlir/lib/IR/PatternLoggingListener.cpp

@@ -0,0 +1,50 @@
+#include "mlir/IR/PatternMatch.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "pattern-logging-listener"
+#define DBGS() (llvm::dbgs() << "[" << DEBUG_TYPE << "] ")
+#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
+
+using namespace mlir;
+
+void RewriterBase::PatternLoggingListener::notifyOperationInserted(
+    Operation *op, InsertPoint previous) {
+  LDBG(patternName << " | notifyOperationInserted"
+                   << " | " << op->getName());
+  ForwardingListener::notifyOperationInserted(op, previous);
+}
+
+void RewriterBase::PatternLoggingListener::notifyOperationModified(
+    Operation *op) {
+  LDBG(patternName << " | notifyOperationModified"
+                   << " | " << op->getName());
+  ForwardingListener::notifyOperationModified(op);
+}
+
+void RewriterBase::PatternLoggingListener::notifyOperationReplaced(
+    Operation *op, Operation *newOp) {
+  LDBG(patternName << " | notifyOperationReplaced (with op)"
+                   << " | " << op->getName() << " | " << newOp->getName());
+  ForwardingListener::notifyOperationReplaced(op, newOp);
+}
+
+void RewriterBase::PatternLoggingListener::notifyOperationReplaced(
+    Operation *op, ValueRange replacement) {
+  LDBG(patternName << " | notifyOperationReplaced (with values)"
+                   << " | " << op->getName());
+  ForwardingListener::notifyOperationReplaced(op, replacement);
+}
+
+void RewriterBase::PatternLoggingListener::notifyOperationErased(
+    Operation *op) {
+  LDBG(patternName << " | notifyOperationErased"
+                   << " | " << op->getName());
+  ForwardingListener::notifyOperationErased(op);
+}
+
+void RewriterBase::PatternLoggingListener::notifyPatternBegin(
+    const Pattern &pattern, Operation *op) {
+  LDBG(patternName << " | notifyPatternBegin"
+                   << " | " << op->getName());
+  ForwardingListener::notifyPatternBegin(pattern, op);
+}

mlir/lib/Rewrite/PatternApplicator.cpp

@@ -15,6 +15,10 @@
 #include "ByteCode.h"
 #include "llvm/Support/Debug.h"
 
+#ifndef NDEBUG
+#include "llvm/ADT/ScopeExit.h"
+#endif
+
 #define DEBUG_TYPE "pattern-application"
 
 using namespace mlir;
@@ -206,11 +210,19 @@ LogicalResult PatternApplicator::matchAndRewrite(
           } else {
             LLVM_DEBUG(llvm::dbgs() << "Trying to match \""
                                     << bestPattern->getDebugName() << "\"\n");
-
             const auto *pattern =
                 static_cast<const RewritePattern *>(bestPattern);
-            result = pattern->matchAndRewrite(op, rewriter);
 
+#ifndef NDEBUG
+            OpBuilder::Listener *oldListener = rewriter.getListener();
+            auto loggingListener =
+                std::make_unique<RewriterBase::PatternLoggingListener>(
+                    oldListener, pattern->getDebugName());
+            rewriter.setListener(loggingListener.get());
+            auto resetListenerCallback = llvm::make_scope_exit(
+                [&] { rewriter.setListener(oldListener); });
+#endif
+            result = pattern->matchAndRewrite(op, rewriter);
             LLVM_DEBUG(llvm::dbgs()
                        << "\"" << bestPattern->getDebugName() << "\" result "
                        << succeeded(result) << "\n");

mlir/test/IR/test-pattern-logging-listener.mlir

@@ -0,0 +1,17 @@
+// RUN: mlir-opt %s --test-walk-pattern-rewrite-driver \
+// RUN:   --allow-unregistered-dialect --debug-only=pattern-logging-listener 2>&1 | FileCheck %s
+
+// Check that when replacing an op with a new op, we get appropriate
+// pattern-logging lines. The regex is because the anonymous namespace is
+// printed differently on different platforms.
+
+// CHECK: [pattern-logging-listener] {{.anonymous.namespace.}}::ReplaceWithNewOp | notifyOperationInserted | test.new_op
+// CHECK: [pattern-logging-listener] {{.anonymous.namespace.}}::ReplaceWithNewOp | notifyOperationReplaced (with values) | test.replace_with_new_op
+// CHECK: [pattern-logging-listener] {{.anonymous.namespace.}}::ReplaceWithNewOp | notifyOperationModified | arith.addi
+// CHECK: [pattern-logging-listener] {{.anonymous.namespace.}}::ReplaceWithNewOp | notifyOperationModified | arith.addi
+// CHECK: [pattern-logging-listener] {{.anonymous.namespace.}}::ReplaceWithNewOp | notifyOperationErased | test.replace_with_new_op
+func.func @replace_with_new_op() -> i32 {
+  %a = "test.replace_with_new_op"() : () -> (i32)
+  %res = arith.addi %a, %a : i32
+  return %res : i32
+}

mlir/test/lit.cfg.py

@@ -301,6 +301,17 @@ def find_real_python_interpreter():
             ToolSubst("mlir-opt", "mlir-opt --verify-roundtrip", unresolved="fatal"),
         ]
     )
+elif "MLIR_GENERATE_PATTERN_CATALOG" in os.environ:
+    tools.extend(
+        [
+            ToolSubst(
+                "mlir-opt",
+                "mlir-opt --debug-only=pattern-logging-listener --mlir-disable-threading",
+                unresolved="fatal",
+            ),
+            ToolSubst("FileCheck", "FileCheck --dump-input=always", unresolved="fatal"),
+        ]
+    )
 else:
     tools.extend(["mlir-opt"])