[mlir:GreedyPatternRewriter] Add debug logging for pattern rewriter actions

[rriddle]

Summary

This effectively mirrors the logging in dialect conversion, which has proven
very useful for understanding the pattern application process.

[jpienaar]

Nice, thanks!

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

177

Could we print ops with regions just with regions elided? (I thought we had such an option)

[rriddle]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

177

We've talked about it before, but haven't actually added the option yet. I did some sneaky hacky thing in the lsp-server, but that technique doesn't apply well generally. I can look at adding a real option to skip regions in a followup (given that it is a common desire in many printing cases).

[silvas]

nice!

[mehdi_amini]

mlir/docs/PatternRewriter.md

353

It's probably worth mentioning that this is enabled only in an assert build

[bondhugula]

This is really useful - thanks!

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

177

It would be great to have such an option and it could perhaps even be the default here. I've often wanted to dump funcs, fors, and ifs without their regions, and this would be useful for debugging in general. I can't see any such option on OpPrintingFlags at the moment.

mlir/docs/PatternRewriter.md

 in a pre-order over the region tree.  This is generally more efficient in
 compile time.  2) the default is a "bottom up" traversal, which builds the
 initial worklist with a postorder traversal of the region tree.  This may
 match larger patterns with ambiguous pattern sets.
 
 Note: This driver is the one used by the [canonicalization](Canonicalization.md)
 [pass](Passes.md/#-canonicalize-canonicalize-operations) in MLIR.
 
+### Debugging
+
+To debug the execution of the greedy pattern rewrite driver,
+`-debug-only=greedy-rewriter` may be used. This command line flag activates
+LLVM's debug logging infrastructure solely for the greedy pattern rewriter. The
+output is formatted as a tree structure, mirroring the structure of the pattern
+application process. This output contains all of the actions performed by the
+rewriter, how operations get processed and patterns are applied, and why they
+fail.
+
+Example output is shown below:
+
+```
+//===-------------------------------------------===//
+Processing operation : 'std.cond_br'(0x60f000001120) {
+  "std.cond_br"(%arg0)[^bb2, ^bb2] {operand_segment_sizes = dense<[1, 0, 0]> : vector<3xi32>} : (i1) -> ()
+
+  * Pattern SimplifyConstCondBranchPred : 'std.cond_br -> ()' {
+  } -> failure : pattern failed to match
+
+  * Pattern SimplifyCondBranchIdenticalSuccessors : 'std.cond_br -> ()' {
+    ** Insert  : 'std.br'(0x60b000003690)
+    ** Replace : 'std.cond_br'(0x60f000001120)
+  } -> success : pattern applied successfully
+} -> success : pattern matched
+//===-------------------------------------------===//
+```
+
+This output is describing the processing of a `std.cond_br` operation. We first
+try to apply the `SimplifyConstCondBranchPred`, which fails. From there, another
+pattern (`SimplifyCondBranchIdenticalSuccessors`) is applied that matches the
+`std.cond_br` and replaces it with a `std.br`.
+
 ## Debugging
 
 ### Pattern Filtering
 
 To simplify test case definition and reduction, the `FrozenRewritePatternSet`
 class provides built-in support for filtering which patterns should be provided
 to the pattern driver for application. Filtering behavior is specified by
 providing a `disabledPatterns` and `enabledPatterns` list when constructing the

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Rewrite/PatternApplicator.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/ScopedPrinter.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 
-#define DEBUG_TYPE "pattern-matcher"
+#define DEBUG_TYPE "greedy-rewriter"
 
 //===----------------------------------------------------------------------===//
 // GreedyPatternRewriteDriver
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This is a worklist-driven driver for the PatternMatcher, which repeatedly
 /// applies the locally optimal patterns in a roughly "bottom up" way.
protected:
   // worklist anymore because we'd get dangling references to it.
   void notifyOperationRemoved(Operation *op) override;
 
   // When the root of a pattern is about to be replaced, it can trigger
   // simplifications to its users - make sure to add them to the worklist
   // before the root is changed.
   void notifyRootReplaced(Operation *op) override;
 
+  /// PatternRewriter hook for erasing a dead operation.
+  void eraseOp(Operation *op) override;
+
+  /// PatternRewriter hook for notifying match failure reasons.
+  LogicalResult
+  notifyMatchFailure(Operation *op,
+                     function_ref<void(Diagnostic &)> reasonCallback) override;
+
   /// The low-level pattern applicator.
   PatternApplicator matcher;
 
   /// The worklist for this transformation keeps track of the operations that
   /// need to be revisited, plus their index in the worklist.  This allows us to
   /// efficiently remove operations from the worklist when they are erased, even
   /// if they aren't the root of a pattern.
   std::vector<Operation *> worklist;
   DenseMap<Operation *, unsigned> worklistMap;
 
   /// Non-pattern based folder for operations.
   OperationFolder folder;
 
 private:
   /// Configuration information for how to simplify.
   GreedyRewriteConfig config;
+
+#ifndef NDEBUG
+  /// A logger used to emit information during the application process.
+  llvm::ScopedPrinter logger{llvm::dbgs()};
+#endif
 };
 } // end anonymous namespace
 
 GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
     MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
     const GreedyRewriteConfig &config)
     : PatternRewriter(ctx), matcher(patterns), folder(ctx), config(config) {
   worklist.reserve(64);
 
   // Apply a simple cost model based solely on pattern benefit.
   matcher.applyDefaultCostModel();
 }
 
 bool GreedyPatternRewriteDriver::simplify(MutableArrayRef<Region> regions) {
+#ifndef NDEBUG
+  const char *logLineComment =
+      "//===-------------------------------------------===//\n";
+
+  /// A utility function to log a process result for the given reason.
+  auto logResult = [&](StringRef result, const llvm::Twine &msg = {}) {
+    logger.unindent();
+    logger.startLine() << "} -> " << result;
+    if (!msg.isTriviallyEmpty())
+      logger.getOStream() << " : " << msg;
+    logger.getOStream() << "\n";
+  };
+  auto logResultWithLine = [&](StringRef result, const llvm::Twine &msg = {}) {
+    logResult(result, msg);
+    logger.startLine() << logLineComment;
+  };
+#endif
+
   bool changed = false;
   unsigned iteration = 0;
   do {
     worklist.clear();
     worklistMap.clear();
 
     if (!config.useTopDownTraversal) {
       // Add operations to the worklist in postorder.
do {
     while (!worklist.empty()) {
       auto *op = popFromWorklist();
 
       // Nulls get added to the worklist when operations are removed, ignore
       // them.
       if (op == nullptr)
         continue;
 
+      LLVM_DEBUG({
+        logger.getOStream() << "\n";
+        logger.startLine() << logLineComment;
+        logger.startLine() << "Processing operation : '" << op->getName()
+                           << "'(" << op << ") {\n";
+        logger.indent();
+
+        // If the operation has no regions, just print it here.
+        if (op->getNumRegions() == 0) {
+          op->print(
+              logger.startLine(),
+              OpPrintingFlags().printGenericOpForm().elideLargeElementsAttrs());
+          logger.getOStream() << "\n\n";
+        }
+      });
+
       // If the operation is trivially dead - remove it.
       if (isOpTriviallyDead(op)) {
         notifyOperationRemoved(op);
         op->erase();
         changed = true;
+
+        LLVM_DEBUG(logResultWithLine("success", "operation is trivially dead"));
         continue;
       }
 
       // Collects all the operands and result uses of the given `op` into work
       // list. Also remove `op` and nested ops from worklist.
       originalOperands.assign(op->operand_begin(), op->operand_end());
       auto preReplaceAction = [&](Operation *op) {
         // Add the operands to the worklist for visitation.
while (!worklist.empty()) {
 
       // Add the given operation to the worklist.
       auto collectOps = [this](Operation *op) { addToWorklist(op); };
 
       // Try to fold this op.
       bool inPlaceUpdate;
       if ((succeeded(folder.tryToFold(op, collectOps, preReplaceAction,
                                       &inPlaceUpdate)))) {
+        LLVM_DEBUG(logResultWithLine("success", "operation was folded"));
+
         changed = true;
         if (!inPlaceUpdate)
           continue;
       }
 
       // Try to match one of the patterns. The rewriter is automatically
       // notified of any necessary changes, so there is nothing else to do
       // here.
-      changed |= succeeded(matcher.matchAndRewrite(op, *this));
+#ifndef NDEBUG
+      auto canApply = [&](const Pattern &pattern) {
+        LLVM_DEBUG({
+          logger.getOStream() << "\n";
+          logger.startLine() << "* Pattern " << pattern.getDebugName() << " : '"
+                             << op->getName() << " -> (";
+          llvm::interleaveComma(pattern.getGeneratedOps(), logger.getOStream());
+          logger.getOStream() << ")' {\n";
+          logger.indent();
+        });
+        return true;
+      };
+      auto onFailure = [&](const Pattern &pattern) {
+        LLVM_DEBUG(logResult("failure", "pattern failed to match"));
+      };
+      auto onSuccess = [&](const Pattern &pattern) {
+        LLVM_DEBUG(logResult("success", "pattern applied successfully"));
+        return success();
+      };
+
+      LogicalResult matchResult =
+          matcher.matchAndRewrite(op, *this, canApply, onFailure, onSuccess);
+      if (succeeded(matchResult))
+        LLVM_DEBUG(logResultWithLine("success", "pattern matched"));
+      else
+        LLVM_DEBUG(logResultWithLine("failure", "pattern failed to match"));
+#else
+      LogicalResult matchResult = matcher.matchAndRewrite(op, *this);
+#endif
+
+
+#ifndef NDEBUG
+#endif
+
+      changed |= succeeded(matchResult);
     }
 
     // After applying patterns, make sure that the CFG of each of the regions
     // is kept up to date.
     if (config.enableRegionSimplification)
       changed |= succeeded(simplifyRegions(*this, regions));
   } while (changed &&
            (++iteration < config.maxIterations ||
void GreedyPatternRewriteDriver::removeFromWorklist(Operation *op) {
   if (it != worklistMap.end()) {
     assert(worklist[it->second] == op && "malformed worklist data structure");
     worklist[it->second] = nullptr;
     worklistMap.erase(it);
   }
 }
 
 void GreedyPatternRewriteDriver::notifyOperationInserted(Operation *op) {
+  LLVM_DEBUG({
+    logger.startLine() << "** Insert  : '" << op->getName() << "'(" << op
+                       << ")\n";
+  });
   addToWorklist(op);
 }
 
 template <typename Operands>
 void GreedyPatternRewriteDriver::addToWorklist(Operands &&operands) {
   for (Value operand : operands) {
     // If the use count of this operand is now < 2, we re-add the defining
     // operation to the worklist.
void GreedyPatternRewriteDriver::notifyOperationRemoved(Operation *op) {
   addToWorklist(op->getOperands());
   op->walk([this](Operation *operation) {
     removeFromWorklist(operation);
     folder.notifyRemoval(operation);
   });
 }
 
 void GreedyPatternRewriteDriver::notifyRootReplaced(Operation *op) {
+  LLVM_DEBUG({
+    logger.startLine() << "** Replace : '" << op->getName() << "'(" << op
+                       << ")\n";
+  });
   for (auto result : op->getResults())
     for (auto *user : result.getUsers())
       addToWorklist(user);
 }
 
+void GreedyPatternRewriteDriver::eraseOp(Operation *op) {
+  LLVM_DEBUG({
+    logger.startLine() << "** Erase   : '" << op->getName() << "'(" << op
+                       << ")\n";
+  });
+  PatternRewriter::eraseOp(op);
+}
+
+LogicalResult GreedyPatternRewriteDriver::notifyMatchFailure(
+    Operation *op, function_ref<void(Diagnostic &)> reasonCallback) {
+  LLVM_DEBUG({
+    Diagnostic diag(op->getLoc(), DiagnosticSeverity::Remark);
+    reasonCallback(diag);
+    logger.startLine() << "** Failure : " << diag.str() << "\n";
+  });
+  return failure();
+}
+
 /// Rewrite the regions of the specified operation, which must be isolated from
 /// above, by repeatedly applying the highest benefit patterns in a greedy
 /// work-list driven manner. Return success if no more patterns can be matched
 /// in the result operation regions. Note: This does not apply patterns to the
 /// top-level operation itself.
 ///
 LogicalResult
 mlir::applyPatternsAndFoldGreedily(MutableArrayRef<Region> regions,