remove mutex/file writes in favor of toolsubst+llvm::dbgs

Author: j2kun

mlir/CMakeLists.txt

@@ -202,22 +202,6 @@ if(MLIR_ENABLE_BINDINGS_PYTHON)
   mlir_configure_python_dev_packages()
 endif()
 
-#-------------------------------------------------------------------------------
-# MLIR Pattern Catalog Generator Configuration
-#
-# When enabled, causes all rewriter patterns to dump their type names and the
-# names of affected operations, which can be used to build a search index
-# mapping operations to patterns.
-#-------------------------------------------------------------------------------
-
-set(MLIR_ENABLE_CATALOG_GENERATOR 0 CACHE BOOL
-       "Enables construction of a catalog of rewrite patterns.")
-
-if (MLIR_ENABLE_CATALOG_GENERATOR)
-  message(STATUS "Enabling MLIR pattern catalog generator")
-  add_definitions(-DMLIR_ENABLE_CATALOG_GENERATOR)
-endif()
-
 set(CMAKE_INCLUDE_CURRENT_DIR ON)
 
 include_directories(BEFORE
@@ -338,4 +322,3 @@ endif()
 if(MLIR_STANDALONE_BUILD)
   llvm_distribution_add_targets()
 endif()
-

mlir/cmake/modules/MLIRConfig.cmake.in

@@ -16,7 +16,6 @@ set(MLIR_IRDL_TO_CPP_EXE "@MLIR_CONFIG_IRDL_TO_CPP_EXE@")
 set(MLIR_INSTALL_AGGREGATE_OBJECTS "@MLIR_INSTALL_AGGREGATE_OBJECTS@")
 set(MLIR_ENABLE_BINDINGS_PYTHON "@MLIR_ENABLE_BINDINGS_PYTHON@")
 set(MLIR_ENABLE_EXECUTION_ENGINE "@MLIR_ENABLE_EXECUTION_ENGINE@")
-set(MLIR_ENABLE_CATALOG_GENERATOR "@MLIR_ENABLE_CATALOG_GENERATOR@")
 
 set_property(GLOBAL PROPERTY MLIR_ALL_LIBS "@MLIR_ALL_LIBS@")
 set_property(GLOBAL PROPERTY MLIR_DIALECT_LIBS "@MLIR_DIALECT_LIBS@")

mlir/include/mlir/IR/PatternMatch.h

@@ -476,76 +476,20 @@ class RewriterBase : public OpBuilder {
   };
 
   struct CatalogingListener : public RewriterBase::ForwardingListener {
-    CatalogingListener(OpBuilder::Listener *listener, StringRef patternName,
-                       raw_ostream &os, std::mutex &writeMutex)
-        : RewriterBase::ForwardingListener(listener), patternName(patternName),
-          os(os), writeMutex(writeMutex) {}
-
-    void notifyOperationInserted(Operation *op, InsertPoint previous) override {
-      {
-        std::lock_guard<std::mutex> lock(writeMutex);
-        os << patternName << " | notifyOperationInserted"
-           << " | " << op->getName() << "\n";
-        os.flush();
-      }
-      ForwardingListener::notifyOperationInserted(op, previous);
-    }
-
-    void notifyOperationModified(Operation *op) override {
-      {
-        std::lock_guard<std::mutex> lock(writeMutex);
-        os << patternName << " | notifyOperationModified"
-           << " | " << op->getName() << "\n";
-        os.flush();
-      }
-      ForwardingListener::notifyOperationModified(op);
-    }
-
-    void notifyOperationReplaced(Operation *op, Operation *newOp) override {
-      {
-        std::lock_guard<std::mutex> lock(writeMutex);
-        os << patternName << " | notifyOperationReplaced (with op)"
-           << " | " << op->getName() << " | " << newOp->getName() << "\n";
-        os.flush();
-      }
-      ForwardingListener::notifyOperationReplaced(op, newOp);
+    CatalogingListener(OpBuilder::Listener *listener, StringRef patternName)
+        : RewriterBase::ForwardingListener(listener), patternName(patternName) {
     }
 
+    void notifyOperationInserted(Operation *op, InsertPoint previous) override;
+    void notifyOperationModified(Operation *op) override;
+    void notifyOperationReplaced(Operation *op, Operation *newOp) override;
     void notifyOperationReplaced(Operation *op,
-                                 ValueRange replacement) override {
-      {
-        std::lock_guard<std::mutex> lock(writeMutex);
-        os << patternName << " | notifyOperationReplaced (with values)"
-           << " | " << op->getName() << "\n";
-        os.flush();
-      }
-      ForwardingListener::notifyOperationReplaced(op, replacement);
-    }
-
-    void notifyOperationErased(Operation *op) override {
-      {
-        std::lock_guard<std::mutex> lock(writeMutex);
-        os << patternName << " | notifyOperationErased"
-           << " | " << op->getName() << "\n";
-        os.flush();
-      }
-      ForwardingListener::notifyOperationErased(op);
-    }
-
-    void notifyPatternBegin(const Pattern &pattern, Operation *op) override {
-      {
-        std::lock_guard<std::mutex> lock(writeMutex);
-        os << patternName << " | notifyPatternBegin"
-           << " | " << op->getName() << "\n";
-        os.flush();
-      }
-      ForwardingListener::notifyPatternBegin(pattern, op);
-    }
+                                 ValueRange replacement) override;
+    void notifyOperationErased(Operation *op) override;
+    void notifyPatternBegin(const Pattern &pattern, Operation *op) override;
 
   private:
     StringRef patternName;
-    raw_ostream &os;
-    std::mutex &writeMutex;
   };
 
   /// Move the blocks that belong to "region" before the given position in

mlir/lib/IR/CMakeLists.txt

@@ -18,6 +18,7 @@ add_mlir_library(MLIRIR
   BuiltinDialectBytecode.cpp
   BuiltinTypes.cpp
   BuiltinTypeInterfaces.cpp
+  CatalogingListener.cpp
   Diagnostics.cpp
   Dialect.cpp
   DialectResourceBlobManager.cpp

mlir/lib/IR/CatalogingListener.cpp

@@ -0,0 +1,49 @@
+#include "mlir/IR/PatternMatch.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "generate-pattern-catalog"
+
+using namespace mlir;
+
+void RewriterBase::CatalogingListener::notifyOperationInserted(
+    Operation *op, InsertPoint previous) {
+  LLVM_DEBUG(llvm::dbgs() << patternName << " | notifyOperationInserted"
+                          << " | " << op->getName() << "\n");
+  ForwardingListener::notifyOperationInserted(op, previous);
+}
+
+void RewriterBase::CatalogingListener::notifyOperationModified(Operation *op) {
+  LLVM_DEBUG(llvm::dbgs() << patternName << " | notifyOperationModified"
+                          << " | " << op->getName() << "\n");
+  ForwardingListener::notifyOperationModified(op);
+}
+
+void RewriterBase::CatalogingListener::notifyOperationReplaced(
+    Operation *op, Operation *newOp) {
+  LLVM_DEBUG(llvm::dbgs() << patternName
+                          << " | notifyOperationReplaced (with op)"
+                          << " | " << op->getName() << " | " << newOp->getName()
+                          << "\n");
+  ForwardingListener::notifyOperationReplaced(op, newOp);
+}
+
+void RewriterBase::CatalogingListener::notifyOperationReplaced(
+    Operation *op, ValueRange replacement) {
+  LLVM_DEBUG(llvm::dbgs() << patternName
+                          << " | notifyOperationReplaced (with values)"
+                          << " | " << op->getName() << "\n");
+  ForwardingListener::notifyOperationReplaced(op, replacement);
+}
+
+void RewriterBase::CatalogingListener::notifyOperationErased(Operation *op) {
+  LLVM_DEBUG(llvm::dbgs() << patternName << " | notifyOperationErased"
+                          << " | " << op->getName() << "\n");
+  ForwardingListener::notifyOperationErased(op);
+}
+
+void RewriterBase::CatalogingListener::notifyPatternBegin(
+    const Pattern &pattern, Operation *op) {
+  LLVM_DEBUG(llvm::dbgs() << patternName << " | notifyPatternBegin"
+                          << " | " << op->getName() << "\n");
+  ForwardingListener::notifyPatternBegin(pattern, op);
+}

mlir/lib/Rewrite/PatternApplicator.cpp

@@ -15,18 +15,8 @@
 #include "ByteCode.h"
 #include "llvm/Support/Debug.h"
 
-#ifdef MLIR_ENABLE_CATALOG_GENERATOR
-#include "llvm/Support/FileSystem.h"
-#include "llvm/Support/raw_ostream.h"
-#include <mutex>
-#endif
-
 #define DEBUG_TYPE "pattern-application"
 
-#ifdef MLIR_ENABLE_CATALOG_GENERATOR
-static std::mutex catalogWriteMutex;
-#endif
-
 using namespace mlir;
 using namespace mlir::detail;
 
@@ -162,16 +152,6 @@ LogicalResult PatternApplicator::matchAndRewrite(
   unsigned anyIt = 0, anyE = anyOpPatterns.size();
   unsigned pdlIt = 0, pdlE = pdlMatches.size();
   LogicalResult result = failure();
-#ifdef MLIR_ENABLE_CATALOG_GENERATOR
-  std::error_code ec;
-  llvm::raw_fd_ostream catalogOs("pattern_catalog.txt", ec,
-                                 llvm::sys::fs::OF_Append);
-  if (ec) {
-    op->emitError("Failed to open pattern catalog file: " + ec.message());
-    return failure();
-  }
-#endif
-
   do {
     // Find the next pattern with the highest benefit.
     const Pattern *bestPattern = nullptr;
@@ -229,19 +209,18 @@ LogicalResult PatternApplicator::matchAndRewrite(
             const auto *pattern =
                 static_cast<const RewritePattern *>(bestPattern);
 
-#ifdef MLIR_ENABLE_CATALOG_GENERATOR
+#ifndef NDEBUG
             OpBuilder::Listener *oldListener = rewriter.getListener();
             RewriterBase::CatalogingListener *catalogingListener =
-                new RewriterBase::CatalogingListener(
-                    oldListener, pattern->getDebugName(), catalogOs,
-                    catalogWriteMutex);
+                new RewriterBase::CatalogingListener(oldListener,
+                                                     pattern->getDebugName());
             rewriter.setListener(catalogingListener);
 #endif
             result = pattern->matchAndRewrite(op, rewriter);
             LLVM_DEBUG(llvm::dbgs()
                        << "\"" << bestPattern->getDebugName() << "\" result "
                        << succeeded(result) << "\n");
-#ifdef MLIR_ENABLE_CATALOG_GENERATOR
+#ifndef NDEBUG
             rewriter.setListener(oldListener);
             delete catalogingListener;
 #endif

mlir/test/lit.cfg.py

@@ -301,6 +301,19 @@ def find_real_python_interpreter():
             ToolSubst("mlir-opt", "mlir-opt --verify-roundtrip", unresolved="fatal"),
         ]
     )
+elif "MLIR_GENERATE_PATTERN_CATALOG" in os.environ:
+    tools.extend([
+        ToolSubst(
+            "mlir-opt",
+            "mlir-opt --debug-only=generate-pattern-catalog --mlir-disable-threading",
+            unresolved="fatal"
+        ),
+        ToolSubst(
+            "FileCheck",
+            "FileCheck --dump-input=always",
+            unresolved="fatal"
+        ),
+    ])
 else:
     tools.extend(["mlir-opt"])
 

utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel

@@ -44,7 +44,6 @@ expand_template(
         "@MLIR_ENABLE_SPIRV_CPU_RUNNER@": "0",
         "@MLIR_ENABLE_VULKAN_RUNNER@": "0",
         "@MLIR_ENABLE_BINDINGS_PYTHON@": "0",
-        "@MLIR_ENABLE_CATALOG_GENERATOR@": "0",
         "@MLIR_RUN_AMX_TESTS@": "0",
         "@MLIR_RUN_ARM_SVE_TESTS@": "0",
         "@MLIR_RUN_ARM_SME_TESTS@": "0",