[MLIR] Support for holistic logging of lowerings in test cases

[frgossen]

Summary

All lowerings that take place during test execution will be recorded in log
files if the macro LOG_LOWERINGS is defined at compile time. This allows to
capture also lowerings that happen in child processes. If the macro remains
undefined, MLIR will behave exactly as before.

[silvas]

Using #ifdef's isn't great. Can we instead have a cl::opt and just use a regular if to control this. Then we can test it too.

[frgossen]

I know, defining this with a macro (instead of cl::opt) feels a bit odd, just like writing all these files instead of printing it to the debug stream.
I considered both, however, there was no way to pass the flag on to all of the pattern applications among the test cases. Especially in the lit tests I would not know how to do this; the debug stream is consumed by, e.g. file-check and the flags are defined in the lit files.

If the aim is to collect all pattern applications, this is still the only way I can think of. If instead this is a debugging feature, we could do it with a simple flag.

mlir/lib/IR/PatternMatch.cpp

 //===- PatternMatch.cpp - Base classes for pattern match ------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Value.h"
+#include "llvm/Support/FileSystem.h"
 
 using namespace mlir;
 
 PatternBenefit::PatternBenefit(unsigned benefit) : representation(benefit) {
   assert(representation == benefit && benefit != ImpossibleToMatchSentinel &&
          "This pattern match benefit is too large to represent");
 }
 
for (RewritePattern *pattern : llvm::concat<RewritePattern *>(
            llvm::make_range(opIt, opE), llvm::make_range(anyIt, anyE))) {
     if (succeeded(matchAndRewrite(op, *pattern, rewriter, canApply, onFailure,
                                   onSuccess)))
       return success();
   }
   return failure();
 }
 
+// Holistic logging of all lowerings that take place during test execution. To
+// capture also lowerings, also in child processes, logs files are written to a
+// target directory that is defined at compile time through the macro
+// `LOWERING_LOG`. If the macro remains undefined, no log files will be
+// generated.
+#ifdef LOWERING_LOG
+namespace {
+struct ListenerWrap : public OpBuilder::Listener {
+  ListenerWrap(OpBuilder::Listener *listener, OperationName rootOpName)
+      : innerListener(listener), rootOpName(rootOpName) {}
+
+  void notifyOperationInserted(Operation *op) override;
+  void notifyBlockCreated(Block *b) override;
+  void publish();
+
+  OpBuilder::Listener *innerListener;
+  OperationName rootOpName;
+  std::vector<OperationName> insertedOpNames;
+};
+
+std::mutex lowering_log_mtx;
+std::unique_ptr<llvm::raw_fd_ostream> lowering_log_os;
+} // namespace
+
+void ListenerWrap::notifyOperationInserted(Operation *op) {
+  insertedOpNames.push_back(op->getName());
+  if (innerListener)
+    innerListener->notifyOperationInserted(op);
+}
+
+void ListenerWrap::notifyBlockCreated(Block *b) {
+  if (innerListener)
+    innerListener->notifyBlockCreated(b);
+}
+
+#define XSTR(X) STR(X)
+#define STR(X) #X
+void ListenerWrap::publish() {
+  std::lock_guard<std::mutex> lock(lowering_log_mtx);
+
+  // Ensure a unique log file for this process.
+  if (!lowering_log_os) {
+    std::error_code ec;
+    auto unique = std::chrono::duration_cast<std::chrono::nanoseconds>(
+                      std::chrono::system_clock::now().time_since_epoch())
+                      .count();
+    std::string fname = std::string(XSTR(LOWERING_LOG)) + "/lowering_" +
+                        std::to_string(unique) + ".log";
+    llvm::dbgs() << fname << "\n";
+    lowering_log_os = std::make_unique<llvm::raw_fd_ostream>(
+        fname, ec, llvm::sys::fs::CD_CreateNew, llvm::sys::fs::FA_Write,
+        llvm::sys::fs::OF_Text);
+    assert(!ec && "failed to open lowering log file");
+  }
+
+  // Log observed lowering.
+  *lowering_log_os << rootOpName << " -> { ";
+  for (auto opName : insertedOpNames)
+    *lowering_log_os << opName << ", ";
+  *lowering_log_os << "}\n";
+}
+#endif
+
 LogicalResult PatternApplicator::matchAndRewrite(
     Operation *op, const RewritePattern &pattern, PatternRewriter &rewriter,
     function_ref<bool(const RewritePattern &)> canApply,
     function_ref<void(const RewritePattern &)> onFailure,
     function_ref<LogicalResult(const RewritePattern &)> onSuccess) {
   // Check that the pattern can be applied.
   if (canApply && !canApply(pattern))
     return failure();
 
   // Try to match and rewrite this pattern. The patterns are sorted by
   // benefit, so if we match we can immediately rewrite.
   rewriter.setInsertionPoint(op);
-  if (succeeded(pattern.matchAndRewrite(op, rewriter)))
+
+#ifdef LOWERING_LOG
+  auto *originalListener = rewriter.getListener();
+  ListenerWrap wrap(originalListener, op->getName());
+  rewriter.setListener(&wrap);
+#endif
+
+  LogicalResult patternApplication = pattern.matchAndRewrite(op, rewriter);
+
+#ifdef LOWERING_LOG
+  wrap.publish();
+  rewriter.setListener(originalListener);
+#endif
+
+  if (succeeded(patternApplication))
     return success(!onSuccess || succeeded(onSuccess(pattern)));
 
   if (onFailure)
     onFailure(pattern);
   return failure();
 }