[mlir] Verifier: steal bit to track seen instead of set. (#102626)

Tracking a set containing every block and operation visited can become very expensive and is unnecessary. Co-authored-by: Will Dietz <[email protected]>
llvm · Aug 9, 2024 · 7a98071 · 7a98071
1 parent 52126dc
commit 7a98071
Showing 1 changed file with 16 additions and 11 deletions.
diff --git 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,35 +269,38 @@ 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 *>())
         return visitor(w.get<Operation *>());
       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))