[mlir][Transforms] GreedyPatternRewriteDriver debugging: Detect faulty patterns

[springerm]

Summary

Compute operation finger prints to detect incorrect API usage in RewritePatterns. Does not work for dialect conversion patterns.

Detect patterns that:

• Returned failure but changed the IR.
• Returned success but did not change the IR.
• Inserted/removed/modified ops, bypassing the rewriter. Not all cases are detected.

These new checks are quite expensive, so they are only enabled with -DMLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON.

Depends On: D151207

[springerm]

add config flag

[springerm]

update

[springerm]

update

[springerm]

rebase

[springerm]

Note: There are two more tests that I have to fix before landing this: IR/greedy-pattern-rewriter-driver.mlir and Transforms/test-strict-pattern-driver.mlir.

[jpienaar]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

316

This feels like it could utilize the Action framework work here.

443

Couldn't this be in a listener rather than interspersed here?

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

317

I'm not fond of the #if / #else here, and the actual code "doing the work" is now duplicated three times because of nested #if !

I would try to:

1. outline the debug code (keep the working code "small").
2. Not duplicate the code that is executed in both debug and non-debug

[springerm]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

316

I'm not sure how to use actions here. I thought about encapsulating all fingerprint-related logic in an action, but that doesn't work because fingerprints must be accessed after a pattern was applied. Also, various listener callbacks (e.g., notifyOperationRemoved) must delete fingerprints.

[springerm]

address comments

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

316

The thing that prevents from using Action I think is that we can't intercept the result of the pattern application I believe, which is important to know if a change in the fingerprint is expected or not.

[mehdi_amini]

mlir/include/mlir/Config/mlir-config.h.cmake

18

I didn't find the NDEBUG requirement by looking at the code?

Also why is ASAN related to this check?

[springerm]

address comments

mlir/include/mlir/Config/mlir-config.h.cmake

18

NDEBUG is not required indeed. ASAN is recommended because some failures manifest as invalid memory accesses. ASAN shows where the memory was deallocated, which is typically something like op->erase, which should have been rewriter.erase(op). Without ASAN, you just see a crash when the deallocated memory was accessed. I added some extra comments.

[mehdi_amini]

mlir/include/mlir/Config/mlir-config.h.cmake

18

Ah! It's during the fingerprinting: makes sense.

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

74

Actually NDEBUG is required here, but you could just do

if(it.second != OperationFingerPrint(it.first))
   llvm_report_fatal("operation finger print changed");

(same for the other asserts)

[springerm]

address comments

mlir/CMakeLists.txt

 set(MLIR_INSTALL_AGGREGATE_OBJECTS 1 CACHE BOOL
     "Installs object files needed for add_mlir_aggregate to work out of \
     tree. Package maintainers can disable this to exclude these assets if \
     not desired. Enabling this will result in object files being written \
     under lib/objects-{CMAKE_BUILD_TYPE}.")
 
 set(MLIR_BUILD_MLIR_C_DYLIB 0 CACHE BOOL "Builds libMLIR-C shared library.")
 
+configure_file(
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Config/mlir-config.h.cmake
+  ${MLIR_INCLUDE_DIR}/mlir/Config/mlir-config.h)
+
 #-------------------------------------------------------------------------------
 # Python Bindings Configuration
 # Requires:
 #   The pybind11 library can be found (set with -DPYBIND_DIR=...)
 #   The python executable is correct (set with -DPython3_EXECUTABLE=...)
 #-------------------------------------------------------------------------------
 
 set(MLIR_ENABLE_BINDINGS_PYTHON 0 CACHE BOOL

mlir/include/mlir/Config/mlir-config.h.cmake

+//===- mlir-config.h - MLIR configuration ------------------------*- C -*-===*//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+/* This file enumerates variables from the MLIR configuration so that they
+   can be in exported headers and won't override package specific directives.
+   This is a C header that can be included in the mlir-c headers. */
+
+#ifndef MLIR_CONFIG_H
+#define MLIR_CONFIG_H
+
+/* Enable expensive checks to detect invalid pattern API usage. Failed checks
+   manifest as failed assertions or invalid memory accesses (e.g., accessing
+   deallocated memory) that cause a crash. Running with ASAN is recommended for
+   easier debugging. */
+#cmakedefine01 MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+
+#endif

mlir/include/mlir/IR/PatternMatch.h

struct Listener : public OpBuilder::Listener {
     notifyMatchFailure(Location loc,
                        function_ref<void(Diagnostic &)> reasonCallback) {
       return failure();
     }
 
     static bool classof(const OpBuilder::Listener *base);
   };
 
+  /// A listener that forwards all notifications to another listener. This
+  /// struct can be used as a base to create listener chains, so that multiple
+  /// listeners can be notified of IR changes.
+  struct ForwardingListener : public RewriterBase::Listener {
+    ForwardingListener(Listener *listener) : listener(listener) {}
+
+    void notifyOperationInserted(Operation *op) override {
+      listener->notifyOperationInserted(op);
+    }
+    void notifyBlockCreated(Block *block) override {
+      listener->notifyBlockCreated(block);
+    }
+    void notifyOperationModified(Operation *op) override {
+      listener->notifyOperationModified(op);
+    }
+    void notifyOperationReplaced(Operation *op,
+                                 ValueRange replacement) override {
+      listener->notifyOperationReplaced(op, replacement);
+    }
+    void notifyOperationRemoved(Operation *op) override {
+      listener->notifyOperationRemoved(op);
+    }
+    LogicalResult notifyMatchFailure(
+        Location loc,
+        function_ref<void(Diagnostic &)> reasonCallback) override {
+      return listener->notifyMatchFailure(loc, reasonCallback);
+    }
+
+  private:
+    Listener *listener;
+  };
+
   /// Move the blocks that belong to "region" before the given position in
   /// another region "parent". The two regions must be different. The caller
   /// is responsible for creating or updating the operation transferring flow
   /// of control to the region and passing it the correct block arguments.
   virtual void inlineRegionBefore(Region &region, Region &parent,
                                   Region::iterator before);
   void inlineRegionBefore(Region &region, Block *before);
 

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 //===- GreedyPatternRewriteDriver.cpp - A greedy rewriter -----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements mlir::applyPatternsAndFoldGreedily.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Config/mlir-config.h"
 #include "mlir/IR/Action.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Rewrite/PatternApplicator.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/ScopeExit.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 "greedy-rewriter"
 
 //===----------------------------------------------------------------------===//
-// GreedyPatternRewriteDriver
+// Debugging Infrastructure
 //===----------------------------------------------------------------------===//
 
 namespace {
+#ifdef MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+/// A helper struct that stores finger prints of ops in order to detect broken
+/// RewritePatterns. A rewrite pattern is broken if it modifies IR without
+/// using the rewriter API or if it returns an inconsistent return value.
+struct DebugFingerPrints : public RewriterBase::ForwardingListener {
+  DebugFingerPrints(RewriterBase::Listener *driver)
+      : RewriterBase::ForwardingListener(driver) {}
+
+  /// Compute finger prints of the given op and its nested ops.
+  void computeFingerPrints(Operation *topLevel) {
+    this->topLevel = topLevel;
+    this->topLevelFingerPrint.emplace(topLevel);
+    topLevel->walk([&](Operation *op) { fingerprints.try_emplace(op, op); });
+  }
+
+  /// Clear all finger prints.
+  void clear() {
+    topLevel = nullptr;
+    topLevelFingerPrint.reset();
+    fingerprints.clear();
+  }
+
+  void notifyRewriteSuccess() {
+    // Pattern application success => IR must have changed.
+    OperationFingerPrint afterFingerPrint(topLevel);
+    assert(*topLevelFingerPrint != afterFingerPrint &&
+           "pattern returned success but IR did not change");
+    for (const auto &it : fingerprints) {
+      // Skip top-level op, its finger print is never invalidated.
+      if (it.first == topLevel)
+        continue;
+      // Note: Finger print computation may crash when an op was erased
+      // without notifying the rewriter. (Run with ASAN.) It does not crash
+      // if a new op was created at the same memory location. (But then the
+      // finger print should have changed.)
+      assert(it.second == OperationFingerPrint(it.first) &&
+             "operation finger print changed");
+    }
+  }
+
+  void notifyRewriteFailure() {
+    // Pattern application failure => IR must not have changed.
+    OperationFingerPrint afterFingerPrint(topLevel);
+    assert(*topLevelFingerPrint == afterFingerPrint &&
+           "pattern returned failure but IR did change");
+  }
+
+protected:
+  /// Invalidate the finger print of the given op, i.e., remove it from the map.
+  void invalidateFingerPrint(Operation *op) {
+    // Invalidate all finger prints until the top level.
+    while (op && op != topLevel) {
+      fingerprints.erase(op);
+      op = op->getParentOp();
+    }
+  }
+
+  void notifyOperationInserted(Operation *op) override {
+    RewriterBase::ForwardingListener::notifyOperationInserted(op);
+    invalidateFingerPrint(op->getParentOp());
+  }
+
+  void notifyOperationModified(Operation *op) override {
+    RewriterBase::ForwardingListener::notifyOperationModified(op);
+    invalidateFingerPrint(op);
+  }
+
+  void notifyOperationRemoved(Operation *op) override {
+    RewriterBase::ForwardingListener::notifyOperationRemoved(op);
+    op->walk([this](Operation *op) { invalidateFingerPrint(op); });
+  }
+
+  /// Operation finger prints to detect invalid pattern API usage. IR is checked
+  /// against these finger prints after pattern application to detect cases
+  /// where IR was modified directly, bypassing the rewriter API.
+  DenseMap<Operation *, OperationFingerPrint> fingerprints;
+
+  /// Top-level operation of the current greedy rewrite.
+  Operation *topLevel = nullptr;
+
+  /// Finger print of the top-level operation.
+  std::optional<OperationFingerPrint> topLevelFingerPrint;
+};
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+
+//===----------------------------------------------------------------------===//
+// GreedyPatternRewriteDriver
+//===----------------------------------------------------------------------===//
+
 /// This is a worklist-driven driver for the PatternMatcher, which repeatedly
 /// applies the locally optimal patterns.
 ///
 /// This abstract class manages the worklist and contains helper methods for
 /// rewriting ops on the worklist. Derived classes specify how ops are added
 /// to the worklist in the beginning.
 class GreedyPatternRewriteDriver : public PatternRewriter,
                                    public RewriterBase::Listener {

 
 #ifndef NDEBUG
   /// A logger used to emit information during the application process.
   llvm::ScopedPrinter logger{llvm::dbgs()};
 #endif
 
   /// The low-level pattern applicator.
   PatternApplicator matcher;
+
+#ifdef MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+  DebugFingerPrints debugFingerPrints;
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
 };
 } // namespace
 
 GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
     MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
     const GreedyRewriteConfig &config)
-    : PatternRewriter(ctx), folder(ctx, this), config(config),
-      matcher(patterns) {
+    : PatternRewriter(ctx), folder(ctx, this), config(config), matcher(patterns)
+#ifdef MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+      // clang-format off
+      , debugFingerPrints(this)
+// clang-format on
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+{
   worklist.reserve(64);
 
   // Apply a simple cost model based solely on pattern benefit.
   matcher.applyDefaultCostModel();
 
   // Set up listener.
+#ifdef MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+  // Send IR notifications to the debug handler. This handler will then forward
+  // all notifications to this GreedyPatternRewriteDriver.
+  setListener(&debugFingerPrints);
+#else
   setListener(this);
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
 }
 
 bool GreedyPatternRewriteDriver::processWorklist() {
 #ifndef NDEBUG
   const char *logLineComment =
       "//===-------------------------------------------===//\n";
 
   /// A utility function to log a process result for the given reason.
if (succeeded(folder.tryToFold(op))) {
       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.
 #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();
     };
 #else
     function_ref<bool(const Pattern &)> canApply = {};
     function_ref<void(const Pattern &)> onFailure = {};
     function_ref<LogicalResult(const Pattern &)> onSuccess = {};
 #endif
 
+#ifdef MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+    debugFingerPrints.computeFingerPrints(
+        /*topLevel=*/config.scope ? config.scope->getParentOp() : op);
+    auto clearFingerprints =
+        llvm::make_scope_exit([&]() { debugFingerPrints.clear(); });
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+
     LogicalResult matchResult =
         matcher.matchAndRewrite(op, *this, canApply, onFailure, onSuccess);
 
     if (succeeded(matchResult)) {
       LLVM_DEBUG(logResultWithLine("success", "pattern matched"));
+#ifdef MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+      debugFingerPrints.notifyRewriteSuccess();
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
       changed = true;
       ++numRewrites;
     } else {
       LLVM_DEBUG(logResultWithLine("failure", "pattern failed to match"));
+#ifdef MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
+      debugFingerPrints.notifyRewriteFailure();
+#endif // MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS
     }
   }
 
   return changed;
 }
 
 void GreedyPatternRewriteDriver::addToWorklist(Operation *op) {
+  assert(op && "expected valid op");
   // Gather potential ancestors while looking for a "scope" parent region.
   SmallVector<Operation *, 8> ancestors;
   Region *region = nullptr;
   do {
     ancestors.push_back(op);
     region = op->getParentRegion();
     if (config.scope == region) {
       // Scope (can be `nullptr`) was reached. Stop traveral and enqueue ops.
void GreedyPatternRewriteDriver::notifyOperationModified(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Modified: '" << op->getName() << "'(" << op
                        << ")\n";
   });
   if (config.listener)
     config.listener->notifyOperationModified(op);
   addToWorklist(op);
 }
 
 void GreedyPatternRewriteDriver::addOperandsToWorklist(ValueRange operands) {
   for (Value operand : operands) {
     // If the use count of this operand is now < 2, we re-add the defining
     // operation to the worklist.
     // TODO: This is based on the fact that zero use operations
     // may be deleted, and that single use values often have more
     // canonicalization opportunities.
     if (!operand || (!operand.use_empty() && !operand.hasOneUse()))