[Mlir-commits] [mlir] [mlir][Transforms] `GreedyPatternRewriteDriver`: log successful folding (PR #77796)

[Matthias Springer]

https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/77796

>From edc4bd3b133eaf9fafea128a77a408abc6263060 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Thu, 11 Jan 2024 16:36:19 +0000
Subject: [PATCH] [mlir][Transforms] `GreedyPatternRewriteDriver`: log
 successful folding

Similar to successful pattern applications, dump the rewritten IR after each successful folding when running with `-debug`.
---
 .../Utils/GreedyPatternRewriteDriver.cpp      | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 67c2d9d59f4c92b..d31408d240ebd57 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -136,6 +136,21 @@ struct DebugFingerPrints : public RewriterBase::ForwardingListener {
 };
 #endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
 
+#ifndef NDEBUG
+static Operation *getDumpRootOp(Operation *op) {
+  // Dump the parent op so that materialized constants are visible. If the op
+  // is a top-level op, dump it directly.
+  if (Operation *parentOp = op->getParentOp())
+    return parentOp;
+  return op;
+}
+static void logSuccessfulFolding(Operation *op) {
+  llvm::dbgs() << "// *** IR Dump After Successful Folding ***\n";
+  op->dump();
+  llvm::dbgs() << "\n\n";
+}
+#endif // NDEBUG
+
 //===----------------------------------------------------------------------===//
 // Worklist
 //===----------------------------------------------------------------------===//
@@ -434,10 +449,14 @@ bool GreedyPatternRewriteDriver::processWorklist() {
       SmallVector<OpFoldResult> foldResults;
       if (succeeded(op->fold(foldResults))) {
         LLVM_DEBUG(logResultWithLine("success", "operation was folded"));
+#ifndef NDEBUG
+        Operation *dumpRootOp = getDumpRootOp(op);
+#endif // NDEBUG
         if (foldResults.empty()) {
           // Op was modified in-place.
           notifyOperationModified(op);
           changed = true;
+          LLVM_DEBUG(logSuccessfulFolding(dumpRootOp));
 #if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
           if (config.scope && failed(verify(config.scope->getParentOp())))
             llvm::report_fatal_error("IR failed to verify after folding");
@@ -492,6 +511,7 @@ bool GreedyPatternRewriteDriver::processWorklist() {
         if (materializationSucceeded) {
           replaceOp(op, replacements);
           changed = true;
+          LLVM_DEBUG(logSuccessfulFolding(dumpRootOp));
 #if MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
           if (config.scope && failed(verify(config.scope->getParentOp())))
             llvm::report_fatal_error("IR failed to verify after folding");