[MLIR] Debug log IR after pattern applications

[frgossen]

Summary

Like print-ir-after-all and -before-all, this allows to inspect IR for
debug purposes. While the former allow to inspect only between passes, this
change allows to follow the rewrites that happen within passes.

[jpienaar]

mlir/lib/Rewrite/PatternApplicator.cpp

42

Why not just have this function defined and referenced only in debug? Same below. So rather than hallowing out, you'd have both only in debug and then yes you'd have two more debug blocks at call site (so 1 here, 2 below vs 2 inside) but seems like it would be clearer that it only has effect in debug mode.

[frgossen]

Address comments

[frgossen]

Thanks!

mlir/lib/Rewrite/PatternApplicator.cpp

42

Did it analogous to logImpossibleToMatch. Changed this for both now :)

[rriddle]

mlir/lib/Rewrite/PatternApplicator.cpp

42

This looks really awkward, what are you trying to achieve here? This assumption of Module is invalid, not all IR is anchored under a module and when it is the module isn't necessarily the top-level operation.

[rriddle]

mlir/lib/Rewrite/PatternApplicator.cpp

42

Dumping the entire module also seems a bit like overkill. Is there a more principled approach that we can have here?

[rriddle]

mlir/lib/Rewrite/PatternApplicator.cpp

42

Reading my comments back, they can seem a bit terse. Apologies for that. I'm very +1 on improving the logging here, and dumping IR in-between pattern applications is something that has been discussed for a while. My main reservation is that this seems very stop gap, and this would be a good time to think of something that scales.

[frgossen]

Address comments

[frgossen]

mlir/lib/Rewrite/PatternApplicator.cpp

42

Why the module / top level op?

• Correct me if I'm wrong but in the PatternApplicator there isn't a simple way to know what surrounding ops may be affected by a rewrite. Limiting the IR dump to, e.g., the parent op might unintentionally hide some effects. It would be nice to pick the root op of the currently running pass here but propagating that information would need changes in many places, I think - probably not worth it.
• I would like to print the IR from the same root op throughout the entire rewrite process. This makes IR dumps comparable and I found it much easier to follow the changes. Again, the current pass' root op would be a good choice.

Something that scales

• I am not entirely sure what you mean here. The debug log in the dialect conversion, e.g., is implemented in a very similar way.
• Instead of the LLVM_DEBUG mechanism, we could use MLIRContext flags here. This would also allow to make the IR dumps more configurable, e.g., dump only the parent op, dump the top level op, etc. I'd be happy to implement this if it is what you mean with "something that scales". Otherwise, can you elaborate?
• Either way, this change is very local and easy enough to replace once something more advanced is in place.

[frgossen]

Friendly ping.

[mehdi_amini]

mlir/lib/Rewrite/PatternApplicator.cpp

42

Correct me if I'm wrong but in the PatternApplicator there isn't a simple way to know what surrounding ops may be affected by a rewrite. Limiting the IR dump to, e.g., the parent op might unintentionally hide some effects.

You can limit yourself to the next IsolatedFromAbove enclosing operation: passes (and so multi-threading) is limited by this.

56

(I think the else should never be needed here if you update the previous function to look for the first IsolatedFromAbove parent op / or the top-level).

191

Add a comment about why you're getting this here (i.e. op may be invalidated below). This got me to think a little...

[rriddle]

mlir/lib/Rewrite/PatternApplicator.cpp

42

By something that scales, I do in fact mean something that is more configurable. My main reservation is that dumping huge chunks of IR while debugging only fits one particular debugging use case, and doesn't really scale to general debugging (especially if you are in a medium or larger module). If I wanted to use my normal debug flow, I'd have to come in here and comment this out (which is relatively the same as what you have to do now when you want to dump it). There has been a TODO for a while (though I've been dragged in many other directions) to have a general "debug config" that gets piped from passes into the different pattern drivers. This config would have different options to aid debugging/testing/etc. For example, disabling/enabling specific patterns by name. I see this kind of IR dumping as fitting in that kind of scheme, i.e. control IR dumping granularity in a similar way to what happens for passes (though these configs would be pass specific).

As for finding a specific root operation to dump from, this could be something provided to the PatternApplicator constructor by the pattern driver. An alternative is to go with @mehdi_amini 's suggestion, which wouldn't necessarily give you the same operation every time, but would give something at a finer grain than the top-level module.

At any rate, maybe my concern isn't that warranted anyways. This debug logging is under a different DEBUG_TYPE than say dialect-conversion, which already kind of gives it that "special flag to enable" kind of feel. That may work for now.

[frgossen]

Address comments

[frgossen]

Thanks!

[rriddle]

Given that nothing else is really logged with this DEBUG_TYPE, I'm alright with this until we have something better.

mlir/lib/Rewrite/PatternApplicator.cpp

 // user defined cost model.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Rewrite/PatternApplicator.h"
 #include "ByteCode.h"
 #include "llvm/Support/Debug.h"
 
-#define DEBUG_TYPE "pattern-match"
+#define DEBUG_TYPE "pattern-application"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 PatternApplicator::PatternApplicator(
     const FrozenRewritePatternSet &frozenPatternList)
     : frozenPatternList(frozenPatternList) {
   if (const PDLByteCode *bytecode = frozenPatternList.getPDLByteCode()) {
     mutableByteCodeState = std::make_unique<PDLByteCodeMutableState>();
     bytecode->initializeMutableState(*mutableByteCodeState);
   }
 }
 PatternApplicator::~PatternApplicator() {}
 
+#ifndef NDEBUG
 /// Log a message for a pattern that is impossible to match.
 static void logImpossibleToMatch(const Pattern &pattern) {
-  LLVM_DEBUG({
     llvm::dbgs() << "Ignoring pattern '" << pattern.getRootKind()
                  << "' because it is impossible to match or cannot lead "
                     "to legal IR (by cost model)\n";
-  });
 }
 
+/// Log IR after pattern application.
+static Operation *getDumpRootOp(Operation *op) {
+  return op->getParentWithTrait<mlir::OpTrait::IsIsolatedFromAbove>();
+}
+static void logSucessfulPatternApplication(Operation *op) {
+  llvm::dbgs() << "// *** IR Dump After Pattern Application ***\n";
+  op->dump();
+  llvm::dbgs() << "\n\n";
+}
+#endif
+
 void PatternApplicator::applyCostModel(CostModel model) {
   // Apply the cost model to the bytecode patterns first, and then the native
   // patterns.
   if (const PDLByteCode *bytecode = frozenPatternList.getPDLByteCode()) {
     for (auto it : llvm::enumerate(bytecode->getPatterns()))
       mutableByteCodeState->updatePatternBenefit(it.index(), model(it.value()));
   }
 
   // Copy over the patterns so that we can sort by benefit based on the cost
   // model. Patterns that are already impossible to match are ignored.
   patterns.clear();
   for (const auto &it : frozenPatternList.getOpSpecificNativePatterns()) {
     for (const RewritePattern *pattern : it.second) {
       if (pattern->getBenefit().isImpossibleToMatch())
-        logImpossibleToMatch(*pattern);
+        LLVM_DEBUG(logImpossibleToMatch(*pattern));
       else
         patterns[it.first].push_back(pattern);
     }
   }
   anyOpPatterns.clear();
   for (const RewritePattern &pattern :
        frozenPatternList.getMatchAnyOpNativePatterns()) {
     if (pattern.getBenefit().isImpossibleToMatch())
-      logImpossibleToMatch(pattern);
+      LLVM_DEBUG(logImpossibleToMatch(pattern));
     else
       anyOpPatterns.push_back(&pattern);
   }
 
   // Sort the patterns using the provided cost model.
   llvm::SmallDenseMap<const Pattern *, PatternBenefit> benefits;
   auto cmp = [&benefits](const Pattern *lhs, const Pattern *rhs) {
     return benefits[lhs] > benefits[rhs];
   };
   auto processPatternList = [&](SmallVectorImpl<const RewritePattern *> &list) {
     // Special case for one pattern in the list, which is the most common case.
     if (list.size() == 1) {
       if (model(*list.front()).isImpossibleToMatch()) {
-        logImpossibleToMatch(*list.front());
+        LLVM_DEBUG(logImpossibleToMatch(*list.front()));
         list.clear();
       }
       return;
     }
 
     // Collect the dynamic benefits for the current pattern list.
     benefits.clear();
     for (const Pattern *pat : list)
       benefits.try_emplace(pat, model(*pat));
 
     // Sort patterns with highest benefit first, and remove those that are
     // impossible to match.
     std::stable_sort(list.begin(), list.end(), cmp);
-    while (!list.empty() && benefits[list.back()].isImpossibleToMatch())
-      logImpossibleToMatch(*list.pop_back_val());
+    while (!list.empty() && benefits[list.back()].isImpossibleToMatch()) {
+      LLVM_DEBUG(logImpossibleToMatch(*list.back()));
+      list.pop_back();
+    }
   };
   for (auto &it : patterns)
     processPatternList(it.second);
   processPatternList(anyOpPatterns);
 }
 
 void PatternApplicator::walkAllPatterns(
     function_ref<void(const Pattern &)> walk) {
do {
     // Check that the pattern can be applied.
     if (canApply && !canApply(*bestPattern))
       continue;
 
     // Try to match and rewrite this pattern. The patterns are sorted by
     // benefit, so if we match we can immediately rewrite. For PDL patterns, the
     // match has already been performed, we just need to rewrite.
     rewriter.setInsertionPoint(op);
+#ifndef NDEBUG
+    // Operation `op` may be invalidated after applying the rewrite pattern.
+    Operation *dumpRootOp = getDumpRootOp(op);
+#endif
     if (pdlMatch) {
       bytecode->rewrite(rewriter, *pdlMatch, *mutableByteCodeState);
       result = success(!onSuccess || succeeded(onSuccess(*bestPattern)));
-
     } else {
       const auto *pattern = static_cast<const RewritePattern *>(bestPattern);
       result = pattern->matchAndRewrite(op, rewriter);
       if (succeeded(result) && onSuccess && failed(onSuccess(*pattern)))
         result = failure();
     }
-    if (succeeded(result))
+    if (succeeded(result)) {
+      LLVM_DEBUG(logSucessfulPatternApplication(dumpRootOp));
       break;
+    }
 
     // Perform any necessary cleanups.
     if (onFailure)
       onFailure(*bestPattern);
   } while (true);
 
   if (mutableByteCodeState)
     mutableByteCodeState->cleanupAfterMatchAndRewrite();
   return result;
 }