[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

260

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

383

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

[mehdi_amini]

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

272

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

260

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

260

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

17 ↗

(On Diff #524679)

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

17 ↗

(On Diff #524679)

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

17 ↗

(On Diff #524679)

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

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

72

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/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

protected:
   /// Notify the driver that the specified operation was removed. Update the
   /// worklist as needed: The operation and its children are removed from the
   /// worklist.
   void notifyOperationRemoved(Operation *op) override;
 
   /// Notify the driver that the specified operation was replaced. Update the
   /// worklist as needed: New users are added enqueued.
   void notifyOperationReplaced(Operation *op, ValueRange replacement) override;
 
   /// Process ops until the worklist is empty or `config.maxNumRewrites` is
   /// reached. Return `true` if any IR was changed.
   bool processWorklist();
 
   /// The worklist for this transformation keeps track of the operations that
   /// need to be revisited, plus their index in the worklist.  This allows us to
   /// efficiently remove operations from the worklist when they are erased, even
   /// if they aren't the root of a pattern.

 
 #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;
+
+#ifndef NDEBUG
+  DenseMap<Operation *, OperationFingerPrint> fingerprints;
+
+  Operation *fingerprintTopLevel = nullptr;
+#endif // NDEBUG
 };
 } // namespace
 
 GreedyPatternRewriteDriver::GreedyPatternRewriteDriver(
     MLIRContext *ctx, const FrozenRewritePatternSet &patterns,
     const GreedyRewriteConfig &config)
     : PatternRewriter(ctx), folder(ctx), config(config), matcher(patterns) {
   worklist.reserve(64);
#ifndef NDEBUG
       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();
       };
 
+      // Compute finger print to detect faulty rewrite patterns.
+      auto clearFingerprints =
+          llvm::make_scope_exit([&]() { fingerprints.clear(); });
+      fingerprintTopLevel = config.scope ? config.scope->getParentOp() : op;
+      fingerprintTopLevel->walk(
+          [&](Operation *op) { fingerprints.try_emplace(op, op); });
+      OperationFingerPrint &beforeFingerPrint =
+          fingerprints.find(fingerprintTopLevel)->second;
+
       LogicalResult matchResult =
           matcher.matchAndRewrite(op, *this, canApply, onFailure, onSuccess);
-      if (succeeded(matchResult))
-        LLVM_DEBUG(logResultWithLine("success", "pattern matched"));
-      else
-        LLVM_DEBUG(logResultWithLine("failure", "pattern failed to match"));
+      OperationFingerPrint afterFingerPrint(fingerprintTopLevel);
+      if (succeeded(matchResult)) {
+      assert(beforeFingerPrint != afterFingerPrint &&
+             "pattern returned success but IR did not change");
+      fingerprintTopLevel->walk([&](Operation *op) {
+        if (op == fingerprintTopLevel)
+          return;
+        auto it = fingerprints.find(op);
+        if (it != fingerprints.end()) {
+          assert(it->second == OperationFingerPrint(op) &&
+                 "operation finger print changed");
+        }
+      });
+      } else {
+      assert(beforeFingerPrint == afterFingerPrint &&
+             "pattern returned failure but IR did change");
+      }
 #else
       LogicalResult matchResult = matcher.matchAndRewrite(op, *this);
 #endif
 
       if (succeeded(matchResult)) {
         changed = true;
         ++numRewrites;
       }
   }
 
   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.
if (config.listener)
     config.listener->notifyBlockCreated(block);
 }
 
 void GreedyPatternRewriteDriver::notifyOperationInserted(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Insert  : '" << op->getName() << "'(" << op
                        << ")\n";
   });
+#ifndef NDEBUG
+  Operation *fpOp = op->getParentOp();
+  while (fpOp && fpOp != fingerprintTopLevel) {
+    fingerprints.erase(fpOp);
+    fpOp = fpOp->getParentOp();
+  }
+#endif // NDEBUG
   if (config.listener)
     config.listener->notifyOperationInserted(op);
   if (config.strictMode == GreedyRewriteStrictness::ExistingAndNewOps)
     strictModeFilteredOps.insert(op);
   addToWorklist(op);
 }
 
 void GreedyPatternRewriteDriver::notifyOperationModified(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Modified: '" << op->getName() << "'(" << op
                        << ")\n";
   });
+#ifndef NDEBUG
+  Operation *fpOp = op;
+  while (fpOp && fpOp != fingerprintTopLevel) {
+    fingerprints.erase(fpOp);
+    fpOp = fpOp->getParentOp();
+  }
+#endif // NDEBUG
   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
void GreedyPatternRewriteDriver::notifyOperationRemoved(Operation *op) {
   });
   if (config.listener)
     config.listener->notifyOperationRemoved(op);
 
   addOperandsToWorklist(op->getOperands());
   op->walk([this](Operation *operation) {
     removeFromWorklist(operation);
     folder.notifyRemoval(operation);
+#ifndef NDEBUG
+    while (operation && operation != fingerprintTopLevel) {
+      fingerprints.erase(operation);
+      operation = operation->getParentOp();
+    }
+#endif // NDEBUG
   });
 
   if (config.strictMode != GreedyRewriteStrictness::AnyOp)
     strictModeFilteredOps.erase(op);
 }
 
 void GreedyPatternRewriteDriver::notifyOperationReplaced(
     Operation *op, ValueRange replacement) {