[MLIR] Update GreedyRewriter to use the LDBG() debug log mechanism (NFC) (#153961)

Also improve a bit the LDBG() implementation

Author: joker-eph

llvm/include/llvm/Support/DebugLog.h

@@ -71,11 +71,10 @@ namespace llvm {
   for (bool _c =                                                               \
            (::llvm::DebugFlag && ::llvm::isCurrentDebugType(TYPE, LEVEL));     \
        _c; _c = false)                                                         \
-    for (::llvm::impl::RAIINewLineStream NewLineStream{(STREAM)}; _c;          \
-         _c = false)                                                           \
-  ::llvm::impl::raw_ldbg_ostream{                                              \
-      ::llvm::impl::computePrefix(TYPE, FILE, LINE, LEVEL), NewLineStream}     \
-      .asLvalue()
+    for (::llvm::impl::raw_ldbg_ostream LdbgOS{                                \
+             ::llvm::impl::computePrefix(TYPE, FILE, LINE, LEVEL), (STREAM)};  \
+         _c; _c = false)                                                       \
+  ::llvm::impl::RAIINewLineStream{LdbgOS}.asLvalue()
 
 #define DEBUGLOG_WITH_STREAM_TYPE_AND_FILE(STREAM, LEVEL, TYPE, FILE)          \
   DEBUGLOG_WITH_STREAM_TYPE_FILE_AND_LINE(STREAM, LEVEL, TYPE, FILE, __LINE__)
@@ -89,22 +88,22 @@ namespace impl {
 class LLVM_ABI raw_ldbg_ostream final : public raw_ostream {
   std::string Prefix;
   raw_ostream &Os;
-  bool HasPendingNewline;
+  bool ShouldPrefixNextString;
 
   /// Split the line on newlines and insert the prefix before each
   /// newline. Forward everything to the underlying stream.
   void write_impl(const char *Ptr, size_t Size) final {
     auto Str = StringRef(Ptr, Size);
-    // Handle the initial prefix.
-    if (!Str.empty())
-      writeWithPrefix(StringRef());
-
     auto Eol = Str.find('\n');
+    // Handle `\n` occurring in the string, ensure to print the prefix at the
+    // beginning of each line.
     while (Eol != StringRef::npos) {
+      // Take the line up to the newline (including the newline).
       StringRef Line = Str.take_front(Eol + 1);
       if (!Line.empty())
         writeWithPrefix(Line);
-      HasPendingNewline = true;
+      // We printed a newline, record here to print a prefix.
+      ShouldPrefixNextString = true;
       Str = Str.drop_front(Eol + 1);
       Eol = Str.find('\n');
     }
@@ -113,24 +112,21 @@ class LLVM_ABI raw_ldbg_ostream final : public raw_ostream {
   }
   void emitPrefix() { Os.write(Prefix.c_str(), Prefix.size()); }
   void writeWithPrefix(StringRef Str) {
-    flushEol();
+    if (ShouldPrefixNextString) {
+      emitPrefix();
+      ShouldPrefixNextString = false;
+    }
     Os.write(Str.data(), Str.size());
   }
 
 public:
   explicit raw_ldbg_ostream(std::string Prefix, raw_ostream &Os,
-                            bool HasPendingNewline = true)
+                            bool ShouldPrefixNextString = true)
       : Prefix(std::move(Prefix)), Os(Os),
-        HasPendingNewline(HasPendingNewline) {
+        ShouldPrefixNextString(ShouldPrefixNextString) {
     SetUnbuffered();
   }
-  ~raw_ldbg_ostream() final { flushEol(); }
-  void flushEol() {
-    if (HasPendingNewline) {
-      emitPrefix();
-      HasPendingNewline = false;
-    }
-  }
+  ~raw_ldbg_ostream() final {}
 
   /// Forward the current_pos method to the underlying stream.
   uint64_t current_pos() const final { return Os.tell(); }
@@ -149,6 +145,7 @@ class RAIINewLineStream final : public raw_ostream {
   ~RAIINewLineStream() { Os << '\n'; }
   void write_impl(const char *Ptr, size_t Size) final { Os.write(Ptr, Size); }
   uint64_t current_pos() const final { return Os.tell(); }
+  RAIINewLineStream &asLvalue() { return *this; }
 };
 
 /// Remove the path prefix from the file name.

llvm/unittests/Support/DebugLogTest.cpp

@@ -115,8 +115,18 @@ TEST(DebugLogTest, StreamPrefix) {
     ldbg_osA << "5";
     EXPECT_EQ(os.str(), expected);
   }
-  // After destructors, there was a pending newline for stream B.
-  EXPECT_EQ(os.str(), expected + "PrefixB ");
+  EXPECT_EQ(os.str(), expected);
+}
+
+TEST(DebugLogTest, DestructorPrefix) {
+  llvm::DebugFlag = true;
+  std::string str;
+  raw_string_ostream os(str);
+  {
+    llvm::impl::raw_ldbg_ostream ldbg_osB("PrefixB ", os);
+  }
+  // After destructors, nothing should have been printed.
+  EXPECT_EQ(os.str(), "");
 }
 #else
 TEST(DebugLogTest, Basic) {

mlir/lib/Transforms/Utils/DialectConversion.cpp

@@ -1138,8 +1138,8 @@ struct ConversionPatternRewriterImpl : public RewriterBase::Listener {
   SmallPtrSet<Operation *, 1> pendingRootUpdates;
 
   /// A raw output stream used to prefix the debug log.
-  llvm::impl::raw_ldbg_ostream os{(Twine("[") + DEBUG_TYPE + "] ").str(),
-                                  llvm::dbgs(), /*HasPendingNewline=*/false};
+  llvm::impl::raw_ldbg_ostream os{(Twine("[") + DEBUG_TYPE + ":1] ").str(),
+                                  llvm::dbgs()};
 
   /// A logger used to emit diagnostics during the conversion process.
   llvm::ScopedPrinter logger{os};

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

@@ -15,6 +15,8 @@
 #include "mlir/Config/mlir-config.h"
 #include "mlir/IR/Action.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Rewrite/PatternApplicator.h"
@@ -23,7 +25,7 @@
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/ScopeExit.h"
-#include "llvm/Support/Debug.h"
+#include "llvm/Support/DebugLog.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -178,9 +180,8 @@ static Operation *getDumpRootOp(Operation *op) {
   return op;
 }
 static void logSuccessfulFolding(Operation *op) {
-  llvm::dbgs() << "// *** IR Dump After Successful Folding ***\n";
-  op->dump();
-  llvm::dbgs() << "\n\n";
+  LDBG() << "// *** IR Dump After Successful Folding ***\n"
+         << OpWithFlags(op, OpPrintingFlags().elideLargeElementsAttrs());
 }
 #endif // NDEBUG
 
@@ -394,8 +395,12 @@ class GreedyPatternRewriteDriver : public RewriterBase::Listener {
                      function_ref<void(Diagnostic &)> reasonCallback) override;
 
 #ifndef NDEBUG
+  /// A raw output stream used to prefix the debug log.
+
+  llvm::impl::raw_ldbg_ostream os{(Twine("[") + DEBUG_TYPE + ":1] ").str(),
+                                  llvm::dbgs()};
   /// A logger used to emit information during the application process.
-  llvm::ScopedPrinter logger{llvm::dbgs()};
+  llvm::ScopedPrinter logger{os};
 #endif
 
   /// The low-level pattern applicator.
@@ -928,10 +933,9 @@ mlir::applyPatternsGreedily(Region &region,
   RegionPatternRewriteDriver driver(region.getContext(), patterns, config,
                                     region);
   LogicalResult converged = std::move(driver).simplify(changed);
-  LLVM_DEBUG(if (failed(converged)) {
-    llvm::dbgs() << "The pattern rewrite did not converge after scanning "
-                 << config.getMaxIterations() << " times\n";
-  });
+  if (failed(converged))
+    LDBG() << "The pattern rewrite did not converge after scanning "
+           << config.getMaxIterations() << " times";
   return converged;
 }
 
@@ -1063,9 +1067,8 @@ LogicalResult mlir::applyOpPatternsGreedily(
   LogicalResult converged = std::move(driver).simplify(ops, changed);
   if (allErased)
     *allErased = surviving.empty();
-  LLVM_DEBUG(if (failed(converged)) {
-    llvm::dbgs() << "The pattern rewrite did not converge after "
-                 << config.getMaxNumRewrites() << " rewrites";
-  });
+  if (failed(converged))
+    LDBG() << "The pattern rewrite did not converge after "
+           << config.getMaxNumRewrites() << " rewrites";
   return converged;
 }