[mlir][NFC] Refactor DialectConversion debug logging

[rriddle]

Summary

This revision beefs up the debug logging within dialect conversion. Given the nature of multi-level legalization, and legalization in general, it is one of the harder pieces of infrastructure to debug. This revision adds nice formatting to make the output log easier to parse:

Legalizing operation : 'std.constant'(0x608000002420) {
  * Fold {
  } -> FAILURE : unable to fold

  * Pattern : 'std.constant -> ()' {
  } -> FAILURE : pattern failed to match

  * Pattern : 'std.constant -> ()' {
  } -> FAILURE : pattern failed to match

  * Pattern : 'std.constant -> (spv.constant)' {
    ** Insert  : 'spv.constant'(0x608000002c20)
    ** Replace : 'std.constant'(0x608000002420)

    //===-------------------------------------------===//
    Legalizing operation : 'spv.constant'(0x608000002c20) {
    } -> SUCCESS : operation marked legal by the target
    //===-------------------------------------------===//
  } -> SUCCESS : pattern applied successfully
} -> SUCCESS

[nicolasvasilache]

Very cool, thanks a lot for this!

[merge_guards_bot]

Unit tests: pass. 62303 tests passed, 0 failed and 838 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

mlir/lib/Transforms/DialectConversion.cpp

 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/Transforms/Utils.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/ScopedPrinter.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 #define DEBUG_TYPE "dialect-conversion"
 
 /// Recursively collect all of the operations to convert from within 'region'.
 /// If 'target' is nonnull, operations that are recursively legal have their
computeConversionSet(iterator_range<Region::iterator> region,
 
   // Check that all blocks in the region were visited.
   if (llvm::any_of(llvm::drop_begin(region, 1),
                    [&](Block &block) { return !visitedBlocks.count(&block); }))
     return emitError(regionLoc, "unreachable blocks were not converted");
   return success();
 }
 
+#ifndef NDEBUG
+/// A utility function to log a successful result for the given reason.
+template <typename... Args>
+static void logSuccess(llvm::ScopedPrinter &os, StringRef fmt,
+                       Args &&... args) {
+  LLVM_DEBUG({
+    os.unindent();
+    os.startLine() << "} -> SUCCESS";
+    if (!fmt.empty())
+      os.getOStream() << " : "
+                      << llvm::formatv(fmt.data(), std::forward<Args>(args)...);
+    os.getOStream() << "\n";
+  });
+}
+
+/// A utility function to log a failure result for the given reason.
+template <typename... Args>
+static void logFailure(llvm::ScopedPrinter &os, StringRef fmt,
+                       Args &&... args) {
+  LLVM_DEBUG({
+    os.unindent();
+    os.startLine() << "} -> FAILURE : "
+                   << llvm::formatv(fmt.data(), std::forward<Args>(args)...)
+                   << "\n";
+  });
+}
+#endif
+
 //===----------------------------------------------------------------------===//
 // Multi-Level Value Mapper
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class wraps a BlockAndValueMapping to provide recursive lookup
 /// functionality, i.e. we will traverse if the mapped value also has a mapping.
 struct ConversionValueMapping {
struct ConversionPatternRewriterImpl {
   /// A transaction state for each of operations that were updated in-place.
   SmallVector<OperationTransactionState, 4> rootUpdates;
 
 #ifndef NDEBUG
   /// A set of operations that have pending updates. This tracking isn't
   /// strictly necessary, and is thus only active during debug builds for extra
   /// verification.
   SmallPtrSet<Operation *, 1> pendingRootUpdates;
+
+  /// A logger used to emit diagnostics during the conversion process.
+  llvm::ScopedPrinter logger{llvm::dbgs()};
 #endif
 };
 } // end namespace detail
 } // end namespace mlir
 
 RewriterState ConversionPatternRewriterImpl::getCurrentState() {
   return RewriterState(createdOps.size(), replacements.size(),
                        blockActions.size(), ignoredOps.size(),

 
 ConversionPatternRewriter::ConversionPatternRewriter(MLIRContext *ctx,
                                                      TypeConverter *converter)
     : PatternRewriter(ctx),
       impl(new detail::ConversionPatternRewriterImpl(*this, converter)) {}
 ConversionPatternRewriter::~ConversionPatternRewriter() {}
 
 /// PatternRewriter hook for replacing the results of an operation.
-void ConversionPatternRewriter::replaceOp(Operation *op,
-                                          ValueRange newValues) {
-  LLVM_DEBUG(llvm::dbgs() << "** Replacing operation : " << op->getName()
-                          << "\n");
+void ConversionPatternRewriter::replaceOp(Operation *op, ValueRange newValues) {
+  LLVM_DEBUG({
+    impl->logger.startLine()
+        << "** Replace : '" << op->getName() << "'(" << op << ")\n";
+  });
   impl->replaceOp(op, newValues);
 }
 
 /// PatternRewriter hook for erasing a dead operation. The uses of this
 /// operation *must* be made dead by the end of the conversion process,
 /// otherwise an assert will be issued.
 void ConversionPatternRewriter::eraseOp(Operation *op) {
-  LLVM_DEBUG(llvm::dbgs() << "** Erasing operation : " << op->getName()
-                          << "\n");
+  LLVM_DEBUG({
+    impl->logger.startLine()
+        << "** Erase   : '" << op->getName() << "'(" << op << ")\n";
+  });
   SmallVector<Value, 1> nullRepls(op->getNumResults(), nullptr);
   impl->replaceOp(op, nullRepls);
 }
 
 /// Apply a signature conversion to the entry block of the given region.
 Block *ConversionPatternRewriter::applySignatureConversion(
     Region *region, TypeConverter::SignatureConversion &conversion) {
   return impl->applySignatureConversion(region, conversion);
void ConversionPatternRewriter::cloneRegionBefore(
   // Collect the range of the cloned blocks.
   auto clonedBeginIt = mapping.lookup(&region.front())->getIterator();
   auto clonedBlocks = llvm::make_range(clonedBeginIt, before);
   impl->notifyRegionWasClonedBefore(clonedBlocks, region.getLoc());
 }
 
 /// PatternRewriter hook for creating a new operation.
 Operation *ConversionPatternRewriter::insert(Operation *op) {
-  LLVM_DEBUG(llvm::dbgs() << "** Inserting operation : " << op->getName()
-                          << "\n");
+  LLVM_DEBUG({
+    impl->logger.startLine()
+        << "** Insert  : '" << op->getName() << "'(" << op << ")\n";
+  });
   impl->createdOps.push_back(op);
   return OpBuilder::insert(op);
 }
 
 /// PatternRewriter hook for updating the root operation in-place.
 void ConversionPatternRewriter::startRootUpdate(Operation *op) {
 #ifndef NDEBUG
   impl->pendingRootUpdates.insert(op);

 bool OperationLegalizer::isIllegal(Operation *op) const {
   // Check if the target explicitly marked this operation as illegal.
   return target.getOpAction(op->getName()) == LegalizationAction::Illegal;
 }
 
 LogicalResult
 OperationLegalizer::legalize(Operation *op,
                              ConversionPatternRewriter &rewriter) {
-  LLVM_DEBUG(llvm::dbgs() << "Legalizing operation : " << op->getName()
-                          << "\n");
+  const char *logLineComment =
+      "//===-------------------------------------------===//\n";
+
+  auto &rewriterImpl = rewriter.getImpl();
+  LLVM_DEBUG({
+    auto &os = rewriterImpl.logger;
+    os.getOStream() << "\n";
+    os.startLine() << logLineComment;
+    os.startLine() << "Legalizing operation : '" << op->getName() << "'(" << op
+                   << ") {\n";
+    os.indent();
+  });
 
   // Check if this operation is legal on the target.
   if (auto legalityInfo = target.isLegal(op)) {
-    LLVM_DEBUG(llvm::dbgs()
-               << "-- Success : Operation marked legal by the target\n");
+    logSuccess(
+        rewriterImpl.logger, "operation marked legal by the target{0}",
+        legalityInfo->isRecursivelyLegal
+            ? "; NOTE: operation is recursively legal; skipping internals"
+            : "");
+    LLVM_DEBUG(rewriterImpl.logger.startLine() << logLineComment);
+
     // If this operation is recursively legal, mark its children as ignored so
     // that we don't consider them for legalization.
-    if (legalityInfo->isRecursivelyLegal) {
-      LLVM_DEBUG(llvm::dbgs() << "-- Success : Operation is recursively legal; "
-                                 "Skipping internals\n");
+    if (legalityInfo->isRecursivelyLegal)
       rewriter.getImpl().markNestedOpsIgnored(op);
-    }
     return success();
   }
 
   // Check to see if the operation is ignored and doesn't need to be converted.
   if (rewriter.getImpl().isOpIgnored(op)) {
-    LLVM_DEBUG(llvm::dbgs()
-               << "-- Success : Operation marked ignored during conversion\n");
+    logSuccess(rewriterImpl.logger,
+               "operation marked 'ignored' during conversion");
+    LLVM_DEBUG(rewriterImpl.logger.startLine() << logLineComment);
     return success();
   }
 
   // If the operation isn't legal, try to fold it in-place.
   // TODO(riverriddle) Should we always try to do this, even if the op is
   // already legal?
   if (succeeded(legalizeWithFold(op, rewriter))) {
-    LLVM_DEBUG(llvm::dbgs() << "-- Success : Operation was folded\n");
+    logSuccess(rewriterImpl.logger, "operation was folded");
+    LLVM_DEBUG(rewriterImpl.logger.startLine() << logLineComment);
     return success();
   }
 
   // Otherwise, we need to apply a legalization pattern to this operation.
   auto it = legalizerPatterns.find(op->getName());
   if (it == legalizerPatterns.end()) {
-    LLVM_DEBUG(llvm::dbgs() << "-- FAIL : no known legalization path.\n");
+    logFailure(rewriterImpl.logger, "no known legalization path");
+    LLVM_DEBUG(rewriterImpl.logger.startLine() << logLineComment);
     return failure();
   }
 
   // The patterns are sorted by expected benefit, so try to apply each in-order.
-  for (auto *pattern : it->second)
-    if (succeeded(legalizePattern(op, pattern, rewriter)))
+  for (auto *pattern : it->second) {
+    if (succeeded(legalizePattern(op, pattern, rewriter))) {
+      logSuccess(rewriterImpl.logger, "");
+      LLVM_DEBUG(rewriterImpl.logger.startLine() << logLineComment);
       return success();
+    }
+  }
 
-  LLVM_DEBUG(llvm::dbgs() << "-- FAIL : no matched legalization pattern.\n");
+  logFailure(rewriterImpl.logger, "no matched legalization pattern");
+  LLVM_DEBUG(rewriterImpl.logger.startLine() << logLineComment);
   return failure();
 }
 
 LogicalResult
 OperationLegalizer::legalizeWithFold(Operation *op,
                                      ConversionPatternRewriter &rewriter) {
   auto &rewriterImpl = rewriter.getImpl();
   RewriterState curState = rewriterImpl.getCurrentState();
 
+  LLVM_DEBUG({
+    rewriterImpl.logger.startLine() << "* Fold {\n";
+    rewriterImpl.logger.indent();
+  });
+
   // Try to fold the operation.
   SmallVector<Value, 2> replacementValues;
   rewriter.setInsertionPoint(op);
-  if (failed(rewriter.tryFold(op, replacementValues)))
+  if (failed(rewriter.tryFold(op, replacementValues))) {
+    logFailure(rewriterImpl.logger, "unable to fold");
     return failure();
+  }
 
   // Insert a replacement for 'op' with the folded replacement values.
   rewriter.replaceOp(op, replacementValues);
 
   // Recursively legalize any new constant operations.
   for (unsigned i = curState.numCreatedOps, e = rewriterImpl.createdOps.size();
        i != e; ++i) {
     Operation *cstOp = rewriterImpl.createdOps[i];
     if (failed(legalize(cstOp, rewriter))) {
-      LLVM_DEBUG(llvm::dbgs() << "-- FAIL: Generated folding constant '"
-                              << cstOp->getName() << "' was illegal.\n");
+      logFailure(rewriterImpl.logger, "generated constant '{0}' was illegal",
+                 cstOp->getName());
       rewriterImpl.resetState(curState);
       return failure();
     }
   }
+
+  logSuccess(rewriterImpl.logger, "");
   return success();
 }
 
 LogicalResult
 OperationLegalizer::legalizePattern(Operation *op, RewritePattern *pattern,
                                     ConversionPatternRewriter &rewriter) {
+  auto &rewriterImpl = rewriter.getImpl();
   LLVM_DEBUG({
-    llvm::dbgs() << "-* Applying rewrite pattern '" << op->getName() << " -> (";
+    auto &os = rewriterImpl.logger;
+    os.getOStream() << "\n";
+    os.startLine() << "* Pattern : '" << pattern->getRootKind() << " -> (";
     interleaveComma(pattern->getGeneratedOps(), llvm::dbgs());
-    llvm::dbgs() << ")'.\n";
+    os.getOStream() << ")' {\n";
+    os.indent();
   });
 
   // Ensure that we don't cycle by not allowing the same pattern to be
   // applied twice in the same recursion stack.
   // TODO(riverriddle) We could eventually converge, but that requires more
   // complicated analysis.
   if (!appliedPatterns.insert(pattern).second) {
-    LLVM_DEBUG(llvm::dbgs() << "-- FAIL: Pattern was already applied.\n");
+    logFailure(rewriterImpl.logger, "pattern was already applied");
     return failure();
   }
 
-  auto &rewriterImpl = rewriter.getImpl();
   RewriterState curState = rewriterImpl.getCurrentState();
   auto cleanupFailure = [&] {
     // Reset the rewriter state and pop this pattern.
     rewriterImpl.resetState(curState);
     appliedPatterns.erase(pattern);
     return failure();
   };
 
   // Try to rewrite with the given pattern.
   rewriter.setInsertionPoint(op);
   auto matchedPattern = pattern->matchAndRewrite(op, rewriter);
 #ifndef NDEBUG
   assert(rewriterImpl.pendingRootUpdates.empty() && "dangling root updates");
 #endif
 
   if (!matchedPattern) {
-    LLVM_DEBUG(llvm::dbgs() << "-- FAIL: Pattern failed to match.\n");
+    logFailure(rewriterImpl.logger, "pattern failed to match");
     return cleanupFailure();
   }
 
   // If the pattern moved or created any blocks, try to legalize their types.
   // This ensures that the types of the block arguments are legal for the region
   // they were moved into.
   for (unsigned i = curState.numBlockActions,
                 e = rewriterImpl.blockActions.size();
        i != e; ++i) {
     auto &action = rewriterImpl.blockActions[i];
     if (action.kind ==
         ConversionPatternRewriterImpl::BlockActionKind::TypeConversion)
       continue;
 
     // Convert the block signature.
     if (failed(rewriterImpl.convertBlockSignature(action.block))) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "-- FAIL: failed to convert types of moved block.\n");
+      logFailure(rewriterImpl.logger, "failed to convert types of moved block");
       return cleanupFailure();
     }
   }
 
   // Check all of the replacements to ensure that the pattern actually replaced
   // the root operation. We also mark any other replaced ops as 'dead' so that
   // we don't try to legalize them later.
   bool replacedRoot = false;
assert((replacedRoot || updatedRootInPlace()) &&
          "expected pattern to replace the root operation");
 
   // Recursively legalize each of the operations updated in place.
   for (unsigned i = curState.numRootUpdates,
                 e = rewriterImpl.rootUpdates.size();
        i != e; ++i) {
     auto &state = rewriterImpl.rootUpdates[i];
     if (failed(legalize(state.getOperation(), rewriter))) {
-      LLVM_DEBUG(llvm::dbgs() << "-- FAIL: Operation updated in-place '"
-                              << op->getName() << "' was illegal.\n");
+      logFailure(rewriterImpl.logger,
+                 "operation updated in-place '{0}' was illegal", op->getName());
       return cleanupFailure();
     }
   }
 
   // Recursively legalize each of the new operations.
   for (unsigned i = curState.numCreatedOps, e = rewriterImpl.createdOps.size();
        i != e; ++i) {
     Operation *op = rewriterImpl.createdOps[i];
     if (failed(legalize(op, rewriter))) {
-      LLVM_DEBUG(llvm::dbgs() << "-- FAIL: Generated operation '"
-                              << op->getName() << "' was illegal.\n");
+      logFailure(rewriterImpl.logger,
+                 "generated operation '{0}'({1}) was illegal", op->getName(),
+                 op);
       return cleanupFailure();
     }
   }
 
+  logSuccess(rewriterImpl.logger, "pattern applied successfully");
   appliedPatterns.erase(pattern);
   return success();
 }
 
 void OperationLegalizer::buildLegalizationGraph(
     const OwningRewritePatternList &patterns) {
   // A mapping between an operation and a set of operations that can be used to
   // generate it.