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

[Matthias Springer]

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

Similar to successful pattern applications, dump the rewritten IR after each successful folding when running with `-debug`.

>From 9787cb932ed0e6a332081419b5430c66ec2d47f5 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         | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 67c2d9d59f4c92..211756c568d79f 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -136,6 +136,17 @@ struct DebugFingerPrints : public RewriterBase::ForwardingListener {
 };
 #endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
 
+#ifndef NDEBUG
+static Operation *getDumpRootOp(Operation *op) {
+  return op->getParentWithTrait<mlir::OpTrait::IsIsolatedFromAbove>();
+}
+static void logSuccessfulFolding(Operation *op) {
+  llvm::dbgs() << "// *** IR Dump After Successful Folding ***\n";
+  op->dump();
+  llvm::dbgs() << "\n\n";
+}
+#endif // NDEBUG
+
 //===----------------------------------------------------------------------===//
 // Worklist
 //===----------------------------------------------------------------------===//
@@ -434,10 +445,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 +507,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");