[Mlir-commits] [mlir] [mlir][Transforms] `GreedyPatternRewriteDriver`: Hash ops separately (PR #78312)

[Matthias Springer]

https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/78312

>From 1d8e71648b31451edda4781dad4d196ee91ef216 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Tue, 16 Jan 2024 13:39:57 +0000
Subject: [PATCH] [mlir][Transforms] `GreedyPatternRewriteDriver`: Hash ops
 separately

The greedy pattern rewrite driver has multiple "expensive checks" to
detect invalid rewrite pattern API usage. As part of these checks, it
computes fingerprints for every op that is in scope, and compares the
fingerprints before and after an attempted pattern application.

Until now, each computed fingerprint took into account all nested
operations. That is quite expensive because it walks the entire IR
subtree. It is also redudant in the expensive checks because we already
compute a fingerprint for every op.

This commit significantly improves the running time of the "expensive
checks" in the greedy pattern rewrite driver.

Depends on #78306. Review only the top commit.
---
 mlir/include/mlir/IR/OperationSupport.h       |  4 +--
 mlir/lib/IR/OperationSupport.cpp              | 15 ++++++++---
 .../Utils/GreedyPatternRewriteDriver.cpp      | 25 ++++++++++++-------
 3 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index 478504b8f76e7ae..f2aa6cee840308e 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -1316,10 +1316,10 @@ LLVM_ENABLE_BITMASK_ENUMS_IN_NAMESPACE();
 //===----------------------------------------------------------------------===//
 
 /// A unique fingerprint for a specific operation, and all of it's internal
-/// operations.
+/// operations (if `includeNested` is set).
 class OperationFingerPrint {
 public:
-  OperationFingerPrint(Operation *topOp);
+  OperationFingerPrint(Operation *topOp, bool includeNested = true);
   OperationFingerPrint(const OperationFingerPrint &) = default;
   OperationFingerPrint &operator=(const OperationFingerPrint &) = default;
 
diff --git a/mlir/lib/IR/OperationSupport.cpp b/mlir/lib/IR/OperationSupport.cpp
index e10cd748e03ba58..a168fe30ba8a486 100644
--- a/mlir/lib/IR/OperationSupport.cpp
+++ b/mlir/lib/IR/OperationSupport.cpp
@@ -911,11 +911,12 @@ static void addDataToHash(llvm::SHA1 &hasher, const T &data) {
       ArrayRef<uint8_t>(reinterpret_cast<const uint8_t *>(&data), sizeof(T)));
 }
 
-OperationFingerPrint::OperationFingerPrint(Operation *topOp) {
+OperationFingerPrint::OperationFingerPrint(Operation *topOp,
+                                           bool includeNested) {
   llvm::SHA1 hasher;
 
-  // Hash each of the operations based upon their mutable bits:
-  topOp->walk([&](Operation *op) {
+  // Helper function that hashes an operation based on its mutable bits:
+  auto addOperationToHash = [&](Operation *op) {
     //   - Operation pointer
     addDataToHash(hasher, op);
     //   - Parent operation pointer (to take into account the nesting structure)
@@ -944,6 +945,12 @@ OperationFingerPrint::OperationFingerPrint(Operation *topOp) {
     //   - Result types
     for (Type t : op->getResultTypes())
       addDataToHash(hasher, t);
-  });
+  };
+
+  if (includeNested)
+    topOp->walk(addOperationToHash);
+  else
+    addOperationToHash(topOp);
+
   hash = hasher.result();
 }
diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index ac73e82bfe92a69..090c833b115b39f 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -60,7 +60,9 @@ struct ExpensiveChecks : public RewriterBase::ForwardingListener {
   void computeFingerPrints(Operation *topLevel) {
     this->topLevel = topLevel;
     this->topLevelFingerPrint.emplace(topLevel);
-    topLevel->walk([&](Operation *op) { fingerprints.try_emplace(op, op); });
+    topLevel->walk([&](Operation *op) {
+      fingerprints.try_emplace(op, op, /*includeNested=*/false);
+    });
   }
 
   /// Clear all finger prints.
@@ -95,7 +97,8 @@ struct ExpensiveChecks : public RewriterBase::ForwardingListener {
       // API.) Finger print computation does may not crash if a new op was
       // created at the same memory location. (But then the finger print should
       // have changed.)
-      if (it.second != OperationFingerPrint(it.first)) {
+      if (it.second !=
+          OperationFingerPrint(it.first, /*includeNested=*/false)) {
         // Note: Run "mlir-opt -debug" to see which pattern is broken.
         llvm::report_fatal_error("operation finger print changed");
       }
@@ -125,14 +128,18 @@ struct ExpensiveChecks : public RewriterBase::ForwardingListener {
 
 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 invalidateFingerPrint(Operation *op) { fingerprints.erase(op); }
 
+  void notifyBlockRemoved(Block *block) override {
+    RewriterBase::ForwardingListener::notifyBlockRemoved(block);
+
+    // The block structure (number of blocks, types of block arguments, etc.)
+    // is part of the fingerprint of the parent op.
+    // TODO: The parent op fingerprint should also be invalidated when modifying
+    // the block arguments of a block, but we do not have a
+    // `notifyBlockModified` callback yet.
+    invalidateFingerPrint(block->getParentOp());
+  }
   void notifyOperationInserted(Operation *op) override {
     RewriterBase::ForwardingListener::notifyOperationInserted(op);
     invalidateFingerPrint(op->getParentOp());