[mlir][Transforms] GreedyPatternRewriteDriver: log successful folding

[matthias-springer]

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

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/77796.diff

1 Files Affected:

• (modified) mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp (+16)

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");

[matthias-springer]

Note: This change was useful for debugging an issue with applyOpPatternsAndFold in iree-org/iree#16073.

[hanhanW]

Thank you very very much!! I was struggle about printing debug message, and I still don't know how to do it. This could help a lot!

[joker-eph]

May be a good tradeoff for now. We can always adjust later.