From e33cf8cbd302099114a4d88af81519387d591a05 Mon Sep 17 00:00:00 2001
From: Will Dietz <w@wdtz.org>
Date: Fri, 9 Aug 2024 09:34:36 -0500
Subject: [PATCH] [mlir] Verifier: steal bit to track seen instead of set.

Tracking a set containing every block and operation visited
can become very expensive and is unnecessary.

This also scales well to small number of operations.
---
 mlir/lib/IR/Verifier.cpp | 27 ++++++++++++++++-----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/mlir/lib/IR/Verifier.cpp b/mlir/lib/IR/Verifier.cpp
index 5d81e9b99814ed..90ff8ef3b497fe 100644
--- a/mlir/lib/IR/Verifier.cpp
+++ b/mlir/lib/IR/Verifier.cpp
@@ -32,6 +32,7 @@
 #include "mlir/IR/RegionKindInterface.h"
 #include "mlir/IR/Threading.h"
 #include "llvm/ADT/DenseMapInfoVariant.h"
+#include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/PrettyStackTrace.h"
@@ -55,6 +56,7 @@ class OperationVerifier {
 
 private:
   using WorkItem = llvm::PointerUnion<Operation *, Block *>;
+  using WorkItemEntry = llvm::PointerIntPair<WorkItem, 1, bool>;
 
   /// This verifier uses a DFS of the tree of operations/blocks. The method
   /// verifyOnEntrance is invoked when we visit a node for the first time, i.e.
@@ -267,10 +269,9 @@ LogicalResult OperationVerifier::verifyOnExit(Operation &op) {
 /// Such ops are collected separately and verified inside
 /// verifyBlockPostChildren.
 LogicalResult OperationVerifier::verifyOperation(Operation &op) {
-  SmallVector<WorkItem> worklist{{&op}};
-  SmallPtrSet<WorkItem, 8> seen;
+  SmallVector<WorkItemEntry> worklist{{&op, false}};
   while (!worklist.empty()) {
-    WorkItem top = worklist.back();
+    WorkItemEntry &top = worklist.back();
 
     auto visit = [](auto &&visitor, WorkItem w) {
       if (w.is<Operation *>())
@@ -278,24 +279,28 @@ LogicalResult OperationVerifier::verifyOperation(Operation &op) {
       return visitor(w.get<Block *>());
     };
 
-    const bool isExit = !seen.insert(top).second;
+    const bool isExit = top.getInt();
+    top.setInt(true);
+    auto item = top.getPointer();
+
     // 2nd visit of this work item ("exit").
     if (isExit) {
-      worklist.pop_back();
-      if (failed(visit(
-              [this](auto *workItem) { return verifyOnExit(*workItem); }, top)))
+      if (failed(
+              visit([this](auto *workItem) { return verifyOnExit(*workItem); },
+                    item)))
         return failure();
+      worklist.pop_back();
       continue;
     }
 
     // 1st visit of this work item ("entrance").
     if (failed(visit(
             [this](auto *workItem) { return verifyOnEntrance(*workItem); },
-            top)))
+            item)))
       return failure();
 
-    if (top.is<Block *>()) {
-      Block &currentBlock = *top.get<Block *>();
+    if (item.is<Block *>()) {
+      Block &currentBlock = *item.get<Block *>();
       // Skip "isolated from above operations".
       for (Operation &o : llvm::reverse(currentBlock)) {
         if (o.getNumRegions() == 0 ||
@@ -305,7 +310,7 @@ LogicalResult OperationVerifier::verifyOperation(Operation &op) {
       continue;
     }
 
-    Operation &currentOp = *top.get<Operation *>();
+    Operation &currentOp = *item.get<Operation *>();
     if (verifyRecursively)
       for (Region &region : llvm::reverse(currentOp.getRegions()))
         for (Block &block : llvm::reverse(region))