[Mlir-commits] [mlir] [mlir] Make symbol-user type verification proportional to participating IR (PR #212354)

Jared Hoberock llvmlistbot at llvm.org
Wed Aug 19 21:54:39 PDT 2026


https://github.com/jaredhoberock updated https://github.com/llvm/llvm-project/pull/212354

>From e459bedcf1f3b5d4ae0b85a5c0109de1c6a44cf7 Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Tue, 18 Aug 2026 18:30:24 -0500
Subject: [PATCH 1/7] [mlir] Skip verified subtrees and reuse one walker per
 scope in symbol-use verification

The SymbolTable verifier walked every operation's types and attribute
dictionary with a fresh AttrTypeWalker per type position, deduplicating
verified types through a scope-wide SetVector<Type>. On a symbol table with
many operations sharing uniqued types, that re-descended already-verified
subtrees repeatedly.

Hoist a single AttrTypeWalker per symbol-table verification and rebind the
lookup anchor to the operation currently being verified. Route every root --
operand, result, and block-argument types, and the attribute dictionary --
through that one walker, so its visited memo both deduplicates and prunes: a
uniqued type or attribute is walked at its first occurrence and every later
occurrence returns from the memo without re-descending. This subsumes the
SetVector<Type>, which is removed.

Skipping the re-descent is sound because the walk is pre-order and
verification fails fast: the first occurrence of a subtree descends and
verifies all of it before any re-encounter, so a re-encounter that returns
from the memo has nothing left to verify. The walker and its memo are locals
of one verifySymbolTable invocation, so no state crosses scopes or verify
passes, and a uniqued type verifies identically regardless of which operation
anchors the lookup because walkSymbolTable does not descend into nested symbol
tables.

Assisted-by: Claude Code (Anthropic)
---
 mlir/lib/IR/SymbolTable.cpp | 67 ++++++++++++++++++++-----------------
 1 file changed, 36 insertions(+), 31 deletions(-)

diff --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index 100eba7146d81..9bc97ab66ab16 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -476,28 +476,16 @@ raw_ostream &mlir::operator<<(raw_ostream &os,
 // SymbolTable Trait Types
 //===----------------------------------------------------------------------===//
 
-/// Verify the symbol uses held by the types owned by `op`: its operand,
-/// result, and block-argument types, and any types nested within its
-/// attributes. `op` is the anchor used for symbol lookups. `verifiedTypes`
-/// records the types already verified within the current symbol table so that
-/// each type, which may be uniqued and shared across many positions or
-/// operations, is verified at most once. Verification fails fast on the first
-/// invalid symbol use.
+/// Verify the symbol uses held by the types owned by `op`: its operand, result,
+/// and block-argument types, and any types nested within its attributes.
+/// `walker` carries the SymbolUserTypeInterface check as a type-walk callback,
+/// anchored at the operation currently being verified, and its visited memo
+/// makes each uniqued type, which may recur across many positions and
+/// operations, verified against the enclosing symbol table at most once.
+/// Verification fails fast on the first invalid symbol use.
 static LogicalResult verifyOpTypeSymbolUses(Operation *op,
-                                            SymbolTableCollection &symbolTable,
-                                            SetVector<Type> &verifiedTypes) {
-  // Walk `type` and any nested type parameters reachable from it, verifying
-  // each not-yet-seen type and interrupting on the first failure.
-  auto verify = [&](Type type) {
-    return type.walk<WalkOrder::PreOrder>([&](Type nestedType) {
-      if (!verifiedTypes.insert(nestedType))
-        return WalkResult::advance();
-      if (auto user = dyn_cast<SymbolUserTypeInterface>(nestedType))
-        if (failed(user.verifySymbolUses(op, symbolTable)))
-          return WalkResult::interrupt();
-      return WalkResult::advance();
-    });
-  };
+                                            AttrTypeWalker &walker) {
+  auto verify = [&](Type type) { return walker.walk<WalkOrder::PreOrder>(type); };
 
   for (Type type : op->getOperandTypes())
     if (verify(type).wasInterrupted())
@@ -511,14 +499,13 @@ static LogicalResult verifyOpTypeSymbolUses(Operation *op,
         if (verify(argument.getType()).wasInterrupted())
           return failure();
 
-  // Verify types nested within the operation's attributes.
-  WalkResult attrResult =
-      op->getAttrDictionary().walk<WalkOrder::PreOrder>([&](Type type) {
-        if (verify(type).wasInterrupted())
-          return WalkResult::interrupt();
-        return WalkResult::advance();
-      });
-  return failure(attrResult.wasInterrupted());
+  // Verify types nested within the operation's attributes. Route the attribute
+  // dictionary through the same walker, and thus the same visited memo, as the
+  // type positions above, so a type recurring across positions and attributes
+  // is verified only at its first occurrence.
+  if (walker.walk<WalkOrder::PreOrder>(op->getAttrDictionary()).wasInterrupted())
+    return failure();
+  return success();
 }
 
 LogicalResult detail::verifySymbolTable(Operation *op) {
@@ -557,7 +544,24 @@ LogicalResult detail::verifySymbolTable(Operation *op) {
   // regardless of which operation anchors the lookup, so each is verified at
   // most once across the whole scope.
   SetVector<Attribute> verifiedAttrs;
-  SetVector<Type> verifiedTypes;
+
+  // A single walker, shared across the whole scope, checks the symbol uses of
+  // every SymbolUserTypeInterface type. Its visited memo records each uniqued
+  // type and attribute once, so a subtree recurring across operand, result,
+  // block-argument, and attribute positions and across operations is walked
+  // only at its first occurrence, whose operation supplies the lookup anchor;
+  // a re-encounter returns from the memo without re-descending. Because a first
+  // occurrence descends the whole subtree pre-order and verification fails
+  // fast, skipping the re-descent verifies nothing new.
+  Operation *typeSymbolUseAnchor = nullptr;
+  AttrTypeWalker typeWalker;
+  typeWalker.addWalk([&](Type type) -> WalkResult {
+    if (auto user = dyn_cast<SymbolUserTypeInterface>(type))
+      if (failed(user.verifySymbolUses(typeSymbolUseAnchor, symbolTable)))
+        return WalkResult::interrupt();
+    return WalkResult::advance();
+  });
+
   auto verifySymbolUserFn = [&](Operation *op) -> std::optional<WalkResult> {
     if (SymbolUserOpInterface user = dyn_cast<SymbolUserOpInterface>(op))
       if (failed(user.verifySymbolUses(symbolTable)))
@@ -570,7 +574,8 @@ LogicalResult detail::verifySymbolTable(Operation *op) {
           return WalkResult::interrupt();
       }
     }
-    if (failed(verifyOpTypeSymbolUses(op, symbolTable, verifiedTypes)))
+    typeSymbolUseAnchor = op;
+    if (failed(verifyOpTypeSymbolUses(op, typeWalker)))
       return WalkResult::interrupt();
     return WalkResult::advance();
   };

>From c1f6fb315f7ab1810ec7540bf0b13cf18059bca8 Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Tue, 18 Aug 2026 18:30:45 -0500
Subject: [PATCH 2/7] [mlir] Enumerate attribute roots without materializing
 dictionaries in symbol-use verification

The attribute half of symbol-use verification called getAttrDictionary(),
which for every operation that keeps its inherent attributes in properties
sorts and uniques a fresh DictionaryAttr. Repeated after every pass, this
interns a dictionary per operation and grows memory with the number of
operations even when nothing participates in symbol verification.

Enumerate the attribute roots directly instead. getRawDictionaryAttrs()
returns the already-stored dictionary -- all attributes for operations that
keep none in properties, the discardable attributes otherwise -- with no
allocation. For operations that carry properties, populateInherentAttrs
appends the inherent attributes into a stack-local NamedAttrList that is
walked and discarded. Coverage is unchanged because getAttrDictionary() itself
obtains the inherent attributes through the same populateInherentAttrs hook.

Assisted-by: Claude Code (Anthropic)
---
 mlir/lib/IR/SymbolTable.cpp | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index 9bc97ab66ab16..c373233dd3728 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -499,12 +499,26 @@ static LogicalResult verifyOpTypeSymbolUses(Operation *op,
         if (verify(argument.getType()).wasInterrupted())
           return failure();
 
-  // Verify types nested within the operation's attributes. Route the attribute
-  // dictionary through the same walker, and thus the same visited memo, as the
-  // type positions above, so a type recurring across positions and attributes
-  // is verified only at its first occurrence.
-  if (walker.walk<WalkOrder::PreOrder>(op->getAttrDictionary()).wasInterrupted())
+  // Verify types nested within the operation's attributes, routed through the
+  // same walker (and thus the same visited memo) as the type positions above.
+  // Read the raw stored attribute dictionary rather than getAttrDictionary():
+  // the latter allocates and uniques a fresh DictionaryAttr for every operation
+  // that keeps its inherent attributes in properties. The raw dictionary
+  // already covers inherent attributes for operations that do not use
+  // properties, and the discardable attributes otherwise; the properties-held
+  // inherent attributes are appended into a stack-local NamedAttrList via
+  // populateInherentAttrs and walked separately -- the same coverage
+  // getAttrDictionary() provides, since it calls the same populateInherentAttrs.
+  if (walker.walk<WalkOrder::PreOrder>(op->getRawDictionaryAttrs())
+          .wasInterrupted())
     return failure();
+  if (op->getPropertiesStorageSize()) {
+    NamedAttrList inherentAttrs;
+    op->getName().populateInherentAttrs(op, inherentAttrs);
+    for (const NamedAttribute &namedAttr : inherentAttrs)
+      if (walker.walk<WalkOrder::PreOrder>(namedAttr.getValue()).wasInterrupted())
+        return failure();
+  }
   return success();
 }
 

>From e4acfd73c370ad0640dd0a8f2df6af58cf7356d8 Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Tue, 18 Aug 2026 18:36:16 -0500
Subject: [PATCH 3/7] [mlir] Cache symbol-reference containment per uniqued
 object and prune symbol-use verification

The question symbol-table verification asks of each type and attribute -- can
this transitively contain a SymbolRefAttr? -- is a pure function of uniqued,
immutable structure. Cache the answer per uniqued object on the MLIRContext,
fill it lazily the first time verification meets each object, and prune the
symbol-use walk on a cached no: a type or attribute that provably holds no
SymbolRefAttr is skipped rather than descended into.

The cache lives on MLIRContextImpl (one member beside distinctAttributeAllocator
and one accessor, in a new lib-internal header), mirroring the tree's existing
context-owned IR-internal state. Two maps keyed on the uniqued opaque pointer
mirror the two uniquers. Uniqued storage is immortal and immutable for the
context's lifetime, so a cached answer can never go stale and a pointer can
never be recycled to alias another object; entries are write-once and never
invalidated.

Kinds with mutable storage are never cached no: they conservatively report
may-contain, and the fill never reads their contents (the verifier still visits
them exactly as before). So a concurrent setBody can neither invalidate a
cached answer nor race a fill.

Concurrency: operations are verified in parallel, so the cache expects
concurrent lookups and fills. It uses the same discipline as the context's
lazily-grown OperationName map -- read-locked probe, writer-locked
double-checked insert, lock elision when multithreading is disabled -- with one
required difference: a fill computes its answer entirely before taking the
writer lock, so it never holds a lock across the recursion and takes no uniquer
lock at all. Racing duplicate fills compute the same immutable fact and the
losing insert is a no-op.

One semantic change, documented on both interfaces: a SymbolUserTypeInterface /
SymbolUserAttrInterface implementation must spell the symbols it references as
SymbolRefAttr sub-elements, so a symbol use is a structural fact. An
implementation that encodes a reference some other way (e.g., a string) is
treated as referencing no symbols and may no longer have verifySymbolUses
invoked. Otherwise the set of things verified against each scope is unchanged.

Assisted-by: Claude Code (Anthropic)
---
 mlir/include/mlir/IR/SymbolInterfaces.td      |  16 +
 mlir/include/mlir/IR/SymbolTable.h            |  13 +
 mlir/lib/IR/MLIRContext.cpp                   |  11 +
 mlir/lib/IR/SymbolRefContainmentCache.h       |  99 ++++++
 mlir/lib/IR/SymbolTable.cpp                   |  90 ++++-
 mlir/unittests/IR/CMakeLists.txt              |   1 +
 .../IR/SymbolReferenceContainmentTest.cpp     | 336 ++++++++++++++++++
 7 files changed, 561 insertions(+), 5 deletions(-)
 create mode 100644 mlir/lib/IR/SymbolRefContainmentCache.h
 create mode 100644 mlir/unittests/IR/SymbolReferenceContainmentTest.cpp

diff --git a/mlir/include/mlir/IR/SymbolInterfaces.td b/mlir/include/mlir/IR/SymbolInterfaces.td
index 292c355cbe157..49d7fe5ac53cc 100644
--- a/mlir/include/mlir/IR/SymbolInterfaces.td
+++ b/mlir/include/mlir/IR/SymbolInterfaces.td
@@ -229,6 +229,14 @@ def SymbolUserAttrInterface : AttrInterface<"SymbolUserAttrInterface"> {
     symbol related utilities that are either costly or otherwise disallowed
     within an operation (e.g., recreating symbol users per op verified rather
     than per symbol table, or querying symbols usage of siblings).
+
+    Implementations must represent the symbols they reference as `SymbolRefAttr`s
+    nested anywhere within their sub-element tree, i.e. reachable by recursive
+    application of `walkImmediateSubElements`, rather than in some other encoding
+    such as a string. This is what lets the symbol machinery cheaply dismiss
+    instances that cannot reference a symbol: an instance with no `SymbolRefAttr`
+    anywhere in its sub-element tree is treated as referencing no symbols, so
+    symbol-table verification may never invoke its `verifySymbolUses`.
   }];
   let cppNamespace = "::mlir";
 
@@ -248,6 +256,14 @@ def SymbolUserTypeInterface : TypeInterface<"SymbolUserTypeInterface"> {
     costly or otherwise disallowed within type construction and uniquing.
     `op` is the operation whose verification triggered the check and should be
     used as the anchor for symbol lookups.
+
+    Implementations must represent the symbols they reference as `SymbolRefAttr`s
+    nested anywhere within their sub-element tree, i.e. reachable by recursive
+    application of `walkImmediateSubElements`, rather than in some other encoding
+    such as a string. This is what lets the symbol machinery cheaply dismiss
+    instances that cannot reference a symbol: an instance with no `SymbolRefAttr`
+    anywhere in its sub-element tree is treated as referencing no symbols, so
+    symbol-table verification may never invoke its `verifySymbolUses`.
   }];
   let cppNamespace = "::mlir";
 
diff --git a/mlir/include/mlir/IR/SymbolTable.h b/mlir/include/mlir/IR/SymbolTable.h
index e4790037d37b2..ea768e0241612 100644
--- a/mlir/include/mlir/IR/SymbolTable.h
+++ b/mlir/include/mlir/IR/SymbolTable.h
@@ -242,6 +242,19 @@ class SymbolTable {
   static bool symbolKnownUseEmpty(StringAttr symbol, Region *from);
   static bool symbolKnownUseEmpty(Operation *symbol, Region *from);
 
+  /// Return whether the given type or attribute may transitively contain a
+  /// `SymbolRefAttr`, i.e. whether one is reachable through its sub-element
+  /// tree by recursive application of `walkImmediateSubElements`. A false
+  /// answer is authoritative: the object provably references no symbol, so a
+  /// SymbolUserTypeInterface / SymbolUserAttrInterface implementation that
+  /// honors its contract has a vacuous `verifySymbolUses` and need never be
+  /// visited. A true answer is conservative: mutable-storage kinds, and
+  /// anything containing them, always report true. The answer is a pure
+  /// function of the uniqued, immutable structure, computed once per object and
+  /// cached on the context for its lifetime.
+  static bool mayContainSymbolRefs(Type type);
+  static bool mayContainSymbolRefs(Attribute attr);
+
   /// Attempt to replace all uses of the given symbol 'oldSymbol' with the
   /// provided symbol 'newSymbol' that are nested within the given operation
   /// 'from'. This does not traverse into any nested symbol tables. If there are
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index da891a7e6e014..6dc200c813dde 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -11,6 +11,7 @@
 #include "AffineMapDetail.h"
 #include "AttributeDetail.h"
 #include "IntegerSetDetail.h"
+#include "SymbolRefContainmentCache.h"
 #include "TypeDetail.h"
 #include "mlir/IR/Action.h"
 #include "mlir/IR/AffineExpr.h"
@@ -269,6 +270,11 @@ class MLIRContextImpl {
   /// destruction.
   DistinctAttributeAllocator distinctAttributeAllocator;
 
+  /// Cache recording, per uniqued type and attribute, whether it may
+  /// transitively contain a SymbolRefAttr. Filled lazily by symbol-table
+  /// verification and never invalidated; see SymbolRefContainmentCache.h.
+  SymbolRefContainmentCache symbolRefContainmentCache;
+
 public:
   MLIRContextImpl(bool threadingIsEnabled)
       : threadingIsEnabled(threadingIsEnabled) {
@@ -1162,6 +1168,11 @@ detail::DistinctAttributeUniquer::allocateStorage(MLIRContext *context,
   return context->getImpl().distinctAttributeAllocator.allocate(referencedAttr);
 }
 
+detail::SymbolRefContainmentCache &
+detail::getSymbolRefContainmentCache(MLIRContext *ctx) {
+  return ctx->getImpl().symbolRefContainmentCache;
+}
+
 /// Return empty dictionary.
 DictionaryAttr DictionaryAttr::getEmpty(MLIRContext *context) {
   return context->getImpl().emptyDictionaryAttr;
diff --git a/mlir/lib/IR/SymbolRefContainmentCache.h b/mlir/lib/IR/SymbolRefContainmentCache.h
new file mode 100644
index 0000000000000..6bc4f83055ad3
--- /dev/null
+++ b/mlir/lib/IR/SymbolRefContainmentCache.h
@@ -0,0 +1,99 @@
+//===- SymbolRefContainmentCache.h - Symbol-ref containment cache ---------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// A context-owned store recording, per uniqued type and attribute, whether it
+// may transitively contain a SymbolRefAttr. Symbol-table verification consults
+// it to prune the symbol-use walk.
+//
+// The store is filled lazily and never invalidated. Uniqued storage is immortal
+// and immutable for the context's lifetime, so a cached answer can never go
+// stale, and a pointer can never be recycled to alias another object within one
+// context. Every entry is write-once: false only for a provably
+// reference-free immutable subtree, true for everything else, including
+// mutable-storage kinds, whose contents the fill never reads.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_LIB_IR_SYMBOLREFCONTAINMENTCACHE_H
+#define MLIR_LIB_IR_SYMBOLREFCONTAINMENTCACHE_H
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Types.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/RWMutex.h"
+#include <optional>
+
+namespace mlir {
+class MLIRContext;
+namespace detail {
+
+/// Per-context store of the "may transitively contain a SymbolRefAttr" fact for
+/// uniqued types and attributes. Two maps mirror the context's two uniquers, so
+/// type and attribute opaque pointers never need to be argued disjoint. The
+/// `lock` flag threaded through each operation is the context's runtime
+/// multithreading flag: when it is false the store is touched single-threaded
+/// and no lock is taken, mirroring MLIRContext's ScopedWriterLock.
+class SymbolRefContainmentCache {
+public:
+  SymbolRefContainmentCache() = default;
+  SymbolRefContainmentCache(const SymbolRefContainmentCache &) = delete;
+  SymbolRefContainmentCache &
+  operator=(const SymbolRefContainmentCache &) = delete;
+
+  /// Return the cached answer for `type`/`attr`, or nullopt if not yet filled.
+  std::optional<bool> lookup(Type type, bool lock) const {
+    return lookupImpl(typeEntries, type.getAsOpaquePointer(), lock);
+  }
+  std::optional<bool> lookup(Attribute attr, bool lock) const {
+    return lookupImpl(attrEntries, attr.getAsOpaquePointer(), lock);
+  }
+
+  /// Record `value` for `type`/`attr` if no entry exists yet and return the
+  /// resident answer. A racing duplicate fill computes the same immutable fact,
+  /// so the losing insert is a no-op.
+  bool insert(Type type, bool value, bool lock) {
+    return insertImpl(typeEntries, type.getAsOpaquePointer(), value, lock);
+  }
+  bool insert(Attribute attr, bool value, bool lock) {
+    return insertImpl(attrEntries, attr.getAsOpaquePointer(), value, lock);
+  }
+
+private:
+  using Map = DenseMap<const void *, bool>;
+
+  std::optional<bool> lookupImpl(const Map &map, const void *key,
+                                 bool lock) const {
+    std::optional<llvm::sys::SmartScopedReader<true>> guard;
+    if (lock)
+      guard.emplace(mutex);
+    auto it = map.find(key);
+    if (it == map.end())
+      return std::nullopt;
+    return it->second;
+  }
+
+  bool insertImpl(Map &map, const void *key, bool value, bool lock) {
+    std::optional<llvm::sys::SmartScopedWriter<true>> guard;
+    if (lock)
+      guard.emplace(mutex);
+    return map.try_emplace(key, value).first->second;
+  }
+
+  mutable llvm::sys::SmartRWMutex<true> mutex;
+  Map typeEntries;
+  Map attrEntries;
+};
+
+/// Return the given context's symbol-reference containment cache. Defined in
+/// MLIRContext.cpp, where MLIRContextImpl is visible.
+SymbolRefContainmentCache &getSymbolRefContainmentCache(MLIRContext *ctx);
+
+} // namespace detail
+} // namespace mlir
+
+#endif // MLIR_LIB_IR_SYMBOLREFCONTAINMENTCACHE_H
diff --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index c373233dd3728..7f6445c84991f 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/SymbolTable.h"
+#include "SymbolRefContainmentCache.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/ADT/SetVector.h"
@@ -472,6 +473,62 @@ raw_ostream &mlir::operator<<(raw_ostream &os,
   llvm_unreachable("Unexpected visibility");
 }
 
+//===----------------------------------------------------------------------===//
+// SymbolRefAttr containment query
+//===----------------------------------------------------------------------===//
+
+static bool computeMayContainSymbolRefs(Type type);
+static bool computeMayContainSymbolRefs(Attribute attr);
+
+/// Fill (once) and return the "may transitively contain a SymbolRefAttr" fact
+/// for a uniqued type or attribute `obj`, seeded by `selfIsRef` for an object
+/// that is itself a symbol reference. The answer is computed entirely outside
+/// the cache lock -- which is taken only for the write-once insert -- so no
+/// lock is held across the recursion and a racing duplicate fill computes the
+/// same fact idempotently. A mutable-storage kind may gain sub-elements after
+/// this point, so its contents are never read; it and its containers report
+/// may-contain conservatively.
+template <typename T>
+static bool fillMayContainSymbolRefs(T obj, bool selfIsRef) {
+  MLIRContext *ctx = obj.getContext();
+  bool lock = ctx->isMultithreadingEnabled();
+  detail::SymbolRefContainmentCache &cache =
+      detail::getSymbolRefContainmentCache(ctx);
+  if (std::optional<bool> cached = cache.lookup(obj, lock))
+    return *cached;
+
+  bool mayContain =
+      selfIsRef || obj.template hasTrait<detail::StorageUserTrait::IsMutable>();
+  // The recursion needs no in-progress guard: immutable objects form a DAG
+  // (sub-elements are interned before their parents), so every cycle passes
+  // through a mutable kind, where the fill stops above before descending.
+  if (!mayContain)
+    obj.walkImmediateSubElements(
+        [&](Attribute sub) {
+          mayContain |= sub && computeMayContainSymbolRefs(sub);
+        },
+        [&](Type sub) {
+          mayContain |= sub && computeMayContainSymbolRefs(sub);
+        });
+  return cache.insert(obj, mayContain, lock);
+}
+
+// A type is never itself a SymbolRefAttr; an attribute is one exactly when
+// isa<SymbolRefAttr> holds (covering FlatSymbolRefAttr).
+static bool computeMayContainSymbolRefs(Type type) {
+  return fillMayContainSymbolRefs(type, /*selfIsRef=*/false);
+}
+static bool computeMayContainSymbolRefs(Attribute attr) {
+  return fillMayContainSymbolRefs(attr, /*selfIsRef=*/isa<SymbolRefAttr>(attr));
+}
+
+bool SymbolTable::mayContainSymbolRefs(Type type) {
+  return computeMayContainSymbolRefs(type);
+}
+bool SymbolTable::mayContainSymbolRefs(Attribute attr) {
+  return computeMayContainSymbolRefs(attr);
+}
+
 //===----------------------------------------------------------------------===//
 // SymbolTable Trait Types
 //===----------------------------------------------------------------------===//
@@ -485,7 +542,18 @@ raw_ostream &mlir::operator<<(raw_ostream &os,
 /// Verification fails fast on the first invalid symbol use.
 static LogicalResult verifyOpTypeSymbolUses(Operation *op,
                                             AttrTypeWalker &walker) {
-  auto verify = [&](Type type) { return walker.walk<WalkOrder::PreOrder>(type); };
+  // A root that provably contains no SymbolRefAttr is not worth entering; the
+  // walker's callbacks prune interior subtrees the same way.
+  auto verify = [&](Type type) {
+    if (!SymbolTable::mayContainSymbolRefs(type))
+      return WalkResult::advance();
+    return walker.walk<WalkOrder::PreOrder>(type);
+  };
+  auto verifyAttr = [&](Attribute attr) {
+    if (!attr || !SymbolTable::mayContainSymbolRefs(attr))
+      return WalkResult::advance();
+    return walker.walk<WalkOrder::PreOrder>(attr);
+  };
 
   for (Type type : op->getOperandTypes())
     if (verify(type).wasInterrupted())
@@ -508,15 +576,15 @@ static LogicalResult verifyOpTypeSymbolUses(Operation *op,
   // properties, and the discardable attributes otherwise; the properties-held
   // inherent attributes are appended into a stack-local NamedAttrList via
   // populateInherentAttrs and walked separately -- the same coverage
-  // getAttrDictionary() provides, since it calls the same populateInherentAttrs.
-  if (walker.walk<WalkOrder::PreOrder>(op->getRawDictionaryAttrs())
-          .wasInterrupted())
+  // getAttrDictionary() provides, since it calls the same
+  // populateInherentAttrs.
+  if (verifyAttr(op->getRawDictionaryAttrs()).wasInterrupted())
     return failure();
   if (op->getPropertiesStorageSize()) {
     NamedAttrList inherentAttrs;
     op->getName().populateInherentAttrs(op, inherentAttrs);
     for (const NamedAttribute &namedAttr : inherentAttrs)
-      if (walker.walk<WalkOrder::PreOrder>(namedAttr.getValue()).wasInterrupted())
+      if (verifyAttr(namedAttr.getValue()).wasInterrupted())
         return failure();
   }
   return success();
@@ -570,11 +638,23 @@ LogicalResult detail::verifySymbolTable(Operation *op) {
   Operation *typeSymbolUseAnchor = nullptr;
   AttrTypeWalker typeWalker;
   typeWalker.addWalk([&](Type type) -> WalkResult {
+    // Prune subtrees that provably hold no SymbolRefAttr: a conforming
+    // SymbolUserTypeInterface spells its references as SymbolRefAttr
+    // sub-elements, so such a type has a vacuous verifySymbolUses.
+    if (!SymbolTable::mayContainSymbolRefs(type))
+      return WalkResult::skip();
     if (auto user = dyn_cast<SymbolUserTypeInterface>(type))
       if (failed(user.verifySymbolUses(typeSymbolUseAnchor, symbolTable)))
         return WalkResult::interrupt();
     return WalkResult::advance();
   });
+  typeWalker.addWalk([&](Attribute attr) -> WalkResult {
+    // Prune reference-free attribute subtrees so the walk descends only where a
+    // SymbolRefAttr, and thus a symbol-using type, may live.
+    if (!SymbolTable::mayContainSymbolRefs(attr))
+      return WalkResult::skip();
+    return WalkResult::advance();
+  });
 
   auto verifySymbolUserFn = [&](Operation *op) -> std::optional<WalkResult> {
     if (SymbolUserOpInterface user = dyn_cast<SymbolUserOpInterface>(op))
diff --git a/mlir/unittests/IR/CMakeLists.txt b/mlir/unittests/IR/CMakeLists.txt
index f5b522bb5cf22..346d27fe47269 100644
--- a/mlir/unittests/IR/CMakeLists.txt
+++ b/mlir/unittests/IR/CMakeLists.txt
@@ -17,6 +17,7 @@ add_mlir_unittest(MLIRIRTests
   PatternMatchTest.cpp
   RemarkTest.cpp  
   ShapedTypeTest.cpp
+  SymbolReferenceContainmentTest.cpp
   SymbolTableTest.cpp
   TypeTest.cpp
   TypeAttrNamesTest.cpp
diff --git a/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp b/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
new file mode 100644
index 0000000000000..b3a6bd76ee3c2
--- /dev/null
+++ b/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
@@ -0,0 +1,336 @@
+//===- SymbolReferenceContainmentTest.cpp - Containment query unit tests --===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Tests for SymbolTable::mayContainSymbolRefs, which records per uniqued type
+// or attribute whether it transitively contains a SymbolRefAttr (conservatively
+// true for mutable storage). Symbol-table verification relies on it to prune
+// types and attributes that provably hold no symbol reference. Because a
+// SymbolUserTypeInterface / SymbolUserAttrInterface implementation must spell
+// its references as SymbolRefAttr sub-elements, a false answer is a sound
+// reason to skip an instance even after the interface is attached late.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/OwningOpRef.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/IR/Verifier.h"
+#include "mlir/Parser/Parser.h"
+#include "gtest/gtest.h"
+
+#include "../../test/lib/Dialect/Test/TestAttributes.h"
+#include "../../test/lib/Dialect/Test/TestDialect.h"
+#include "../../test/lib/Dialect/Test/TestTypes.h"
+
+#include <atomic>
+#include <thread>
+#include <vector>
+
+using namespace mlir;
+
+namespace {
+
+// Symbol-user models whose verification always fails, attached externally to
+// exercise late interface attachment. One targets a type that structurally
+// holds a SymbolRefAttr (a tensor with a symbol-ref encoding); the other
+// targets f32, which holds none.
+struct FailingTensorSymbolUserModel
+    : public SymbolUserTypeInterface::ExternalModel<
+          FailingTensorSymbolUserModel, RankedTensorType> {
+  LogicalResult verifySymbolUses(Type type, Operation *op,
+                                 SymbolTableCollection &symbolTable) const {
+    return op->emitError("tensor rejected by its attached symbol-user model");
+  }
+};
+struct FailingF32SymbolUserModel
+    : public SymbolUserTypeInterface::ExternalModel<FailingF32SymbolUserModel,
+                                                    Float32Type> {
+  LogicalResult verifySymbolUses(Type type, Operation *op,
+                                 SymbolTableCollection &symbolTable) const {
+    return op->emitError("f32 rejected by its attached symbol-user model");
+  }
+};
+
+class SymbolReferenceContainmentTest : public ::testing::Test {
+protected:
+  SymbolReferenceContainmentTest() {
+    context.loadDialect<test::TestDialect>();
+    context.allowUnregisteredDialects();
+  }
+
+  FlatSymbolRefAttr symbolRef() {
+    return FlatSymbolRefAttr::get(&context, "sym");
+  }
+
+  // A conforming type implementing SymbolUserTypeInterface, spelling its
+  // reference as a FlatSymbolRefAttr parameter: !test.symbol_ref<@sym>.
+  test::TestSymbolUserType symbolUserType() {
+    return test::TestSymbolUserType::get(&context, symbolRef());
+  }
+
+  // A conforming attribute implementing SymbolUserAttrInterface, spelling its
+  // reference as a FlatSymbolRefAttr parameter: #test.symbol_ref_attr<@sym>.
+  test::TestSymbolRefAttr symbolUserAttr() {
+    return test::TestSymbolRefAttr::get(&context, symbolRef());
+  }
+
+  MLIRContext context;
+};
+
+// A leaf type holding no symbol reference answers false.
+TEST_F(SymbolReferenceContainmentTest, LeafTypeIsClear) {
+  EXPECT_FALSE(
+      SymbolTable::mayContainSymbolRefs(IntegerType::get(&context, 32)));
+}
+
+// A plain attribute holding no symbol reference answers false.
+TEST_F(SymbolReferenceContainmentTest, LeafAttrIsClear) {
+  EXPECT_FALSE(
+      SymbolTable::mayContainSymbolRefs(StringAttr::get(&context, "hi")));
+  EXPECT_FALSE(SymbolTable::mayContainSymbolRefs(
+      TypeAttr::get(IntegerType::get(&context, 32))));
+}
+
+// A SymbolRefAttr itself answers true.
+TEST_F(SymbolReferenceContainmentTest, FlatSymbolRefAttrIsTrue) {
+  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(symbolRef()));
+}
+
+// A non-flat SymbolRefAttr, which nests further references, answers true.
+TEST_F(SymbolReferenceContainmentTest, NestedSymbolRefAttrIsTrue) {
+  SymbolRefAttr ref =
+      SymbolRefAttr::get(StringAttr::get(&context, "root"),
+                         {FlatSymbolRefAttr::get(&context, "n")});
+  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(ref));
+}
+
+// A conforming symbol-user type answers true through its SymbolRefAttr
+// parameter (not through the interface, which plays no part in the answer).
+TEST_F(SymbolReferenceContainmentTest, ConformingSymbolUserTypeIsTrue) {
+  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(symbolUserType()));
+}
+
+// A conforming symbol-user attribute answers true through its SymbolRefAttr
+// parameter.
+TEST_F(SymbolReferenceContainmentTest, ConformingSymbolUserAttrIsTrue) {
+  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(symbolUserAttr()));
+}
+
+// A type nesting a symbol-ref-bearing type propagates true.
+TEST_F(SymbolReferenceContainmentTest, TypeNestingSymbolRefBearingType) {
+  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(
+      TupleType::get(&context, {symbolUserType()})));
+}
+
+// A tuple of ordinary types stays false.
+TEST_F(SymbolReferenceContainmentTest, TypeNestingOrdinaryTypesIsClear) {
+  Type i32 = IntegerType::get(&context, 32);
+  EXPECT_FALSE(
+      SymbolTable::mayContainSymbolRefs(TupleType::get(&context, {i32, i32})));
+}
+
+// A type reaches a SymbolRefAttr two levels deep, through an attribute
+// sub-element (a tensor encoding holding a TypeAttr of a symbol-ref type).
+TEST_F(SymbolReferenceContainmentTest, TypeReachesSymbolRefThroughAttribute) {
+  Attribute encoding = TypeAttr::get(symbolUserType());
+  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(encoding));
+  RankedTensorType tensor =
+      RankedTensorType::get({2}, IntegerType::get(&context, 32), encoding);
+  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(tensor));
+}
+
+// A type reaches a plain SymbolRefAttr through an attribute parameter (a tensor
+// encoding).
+TEST_F(SymbolReferenceContainmentTest, TypeWithSymbolRefAttrParameter) {
+  RankedTensorType tensor =
+      RankedTensorType::get({2}, IntegerType::get(&context, 32), symbolRef());
+  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(tensor));
+}
+
+// A dictionary attribute containing a plain SymbolRefAttr answers true.
+TEST_F(SymbolReferenceContainmentTest, DictionaryAttrContainingSymbolRef) {
+  NamedAttribute named(StringAttr::get(&context, "callee"), symbolRef());
+  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(
+      DictionaryAttr::get(&context, {named})));
+}
+
+// A dictionary attribute containing a symbol-ref-bearing type inside a TypeAttr
+// answers true; the answer on the dictionary summarizes its whole nested tree.
+TEST_F(SymbolReferenceContainmentTest, DictionaryAttrContainingSymbolRefType) {
+  NamedAttribute named(StringAttr::get(&context, "key"),
+                       TypeAttr::get(symbolUserType()));
+  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(
+      DictionaryAttr::get(&context, {named})));
+}
+
+// A dictionary attribute with no symbol reference stays false.
+TEST_F(SymbolReferenceContainmentTest, DictionaryAttrIsClear) {
+  NamedAttribute named(StringAttr::get(&context, "key"),
+                       TypeAttr::get(IntegerType::get(&context, 32)));
+  EXPECT_FALSE(SymbolTable::mayContainSymbolRefs(
+      DictionaryAttr::get(&context, {named})));
+}
+
+// A type carrying a mutable component answers true conservatively, since its
+// sub-elements may change after the answer is computed at first query; the
+// fill never reads its contents.
+TEST_F(SymbolReferenceContainmentTest, MutableTypeReportsConservatively) {
+  test::TestRecursiveType recursive =
+      test::TestRecursiveType::get(&context, "rec");
+  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(recursive));
+}
+
+// A container of a mutable-storage kind inherits true, even before the mutable
+// body is populated, so no later mutation can turn a cached false stale.
+TEST_F(SymbolReferenceContainmentTest, ContainerOfMutableTypeIsTrue) {
+  test::TestRecursiveType recursive =
+      test::TestRecursiveType::get(&context, "rec2");
+  EXPECT_TRUE(
+      SymbolTable::mayContainSymbolRefs(TupleType::get(&context, {recursive})));
+}
+
+// Interface membership plays no part in the answer, so late attachment needs no
+// fallback: a type that structurally holds a SymbolRefAttr (a tensor with a
+// symbol-ref encoding) answers true from interning, so verification visits it
+// and the newly-attached verifySymbolUses fires.
+TEST_F(SymbolReferenceContainmentTest, LateInterfaceAttachmentStillVerifies) {
+  OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(
+      "module { \"foo.op\"() : () -> tensor<4xf32, @sym> }", &context);
+  ASSERT_TRUE(module);
+
+  RankedTensorType::attachInterface<FailingTensorSymbolUserModel>(context);
+  ScopedDiagnosticHandler handler(&context,
+                                  [](Diagnostic &) { return success(); });
+  EXPECT_TRUE(failed(verify(*module)));
+}
+
+// The contract boundary: a type that references a symbol without spelling it as
+// a SymbolRefAttr (here f32, standing in for a non-conforming symbol-user type)
+// answers false and is therefore skipped -- its verifySymbolUses never fires,
+// so verification succeeds. This is the documented cost of the interface
+// contract.
+TEST_F(SymbolReferenceContainmentTest, NonConformingSymbolUserTypeIsSkipped) {
+  OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(
+      "module { \"foo.op\"() : () -> f32 }", &context);
+  ASSERT_TRUE(module);
+
+  Float32Type::attachInterface<FailingF32SymbolUserModel>(context);
+  ScopedDiagnosticHandler handler(&context,
+                                  [](Diagnostic &) { return success(); });
+  EXPECT_TRUE(succeeded(verify(*module)));
+}
+
+// Concurrency crux: many threads query the same cold, shared objects at once,
+// exactly as parallel symbol-table verification fills the context cache from
+// several isolated-op workers simultaneously. Every thread must agree with the
+// single-threaded ground truth, and the run must be clean under
+// ThreadSanitizer. Built cold (constructed but never queried before the
+// threads start) so the fills genuinely race.
+TEST_F(SymbolReferenceContainmentTest, ConcurrentFillIsRaceFree) {
+  ASSERT_TRUE(context.isMultithreadingEnabled());
+
+  // A set of shared objects spanning the interesting fill paths and sharing
+  // sub-elements, so concurrent fills overlap on the same interior objects.
+  Type i32 = IntegerType::get(&context, 32);
+  std::vector<Attribute> attrs = {
+      StringAttr::get(&context, "leaf"),
+      symbolRef(),
+      symbolUserAttr(),
+      TypeAttr::get(symbolUserType()),
+      DictionaryAttr::get(
+          &context,
+          {NamedAttribute(StringAttr::get(&context, "callee"), symbolRef())}),
+      DictionaryAttr::get(&context,
+                          {NamedAttribute(StringAttr::get(&context, "plain"),
+                                          TypeAttr::get(i32))}),
+  };
+  std::vector<Type> types = {
+      i32,
+      symbolUserType(),
+      TupleType::get(&context, {i32, symbolUserType()}),
+      TupleType::get(&context, {i32, i32}),
+      RankedTensorType::get({2}, i32, symbolRef()),
+  };
+
+  // Ground truth computed single-threaded (this also warms nothing new for the
+  // threads, since these very objects are what they race to fill).
+  std::vector<bool> attrTruth, typeTruth;
+  for (Attribute a : attrs)
+    attrTruth.push_back(SymbolTable::mayContainSymbolRefs(a));
+  for (Type t : types)
+    typeTruth.push_back(SymbolTable::mayContainSymbolRefs(t));
+
+  // Fresh context so the threads hit a cold cache and genuinely race the fills.
+  MLIRContext raced;
+  raced.loadDialect<test::TestDialect>();
+  raced.allowUnregisteredDialects();
+  ASSERT_TRUE(raced.isMultithreadingEnabled());
+  auto rebuildAttrs = [&](MLIRContext &ctx) {
+    Type ri32 = IntegerType::get(&ctx, 32);
+    FlatSymbolRefAttr rsym = FlatSymbolRefAttr::get(&ctx, "sym");
+    return std::vector<Attribute>{
+        StringAttr::get(&ctx, "leaf"),
+        rsym,
+        test::TestSymbolRefAttr::get(&ctx, rsym),
+        TypeAttr::get(test::TestSymbolUserType::get(&ctx, rsym)),
+        DictionaryAttr::get(
+            &ctx, {NamedAttribute(StringAttr::get(&ctx, "callee"), rsym)}),
+        DictionaryAttr::get(&ctx,
+                            {NamedAttribute(StringAttr::get(&ctx, "plain"),
+                                            TypeAttr::get(ri32))}),
+    };
+  };
+  auto rebuildTypes = [&](MLIRContext &ctx) {
+    Type ri32 = IntegerType::get(&ctx, 32);
+    FlatSymbolRefAttr rsym = FlatSymbolRefAttr::get(&ctx, "sym");
+    test::TestSymbolUserType user = test::TestSymbolUserType::get(&ctx, rsym);
+    return std::vector<Type>{
+        ri32,
+        user,
+        TupleType::get(&ctx, {ri32, user}),
+        TupleType::get(&ctx, {ri32, ri32}),
+        RankedTensorType::get({2}, ri32, rsym),
+    };
+  };
+  std::vector<Attribute> racedAttrs = rebuildAttrs(raced);
+  std::vector<Type> racedTypes = rebuildTypes(raced);
+
+  const int numThreads = 16;
+  std::vector<std::vector<bool>> attrResults(numThreads);
+  std::vector<std::vector<bool>> typeResults(numThreads);
+  std::atomic<int> ready{0};
+  std::atomic<bool> go{false};
+  std::vector<std::thread> threads;
+  for (int i = 0; i < numThreads; ++i)
+    threads.emplace_back([&, i] {
+      ready.fetch_add(1);
+      while (!go.load())
+        ;
+      for (Attribute a : racedAttrs)
+        attrResults[i].push_back(SymbolTable::mayContainSymbolRefs(a));
+      for (Type t : racedTypes)
+        typeResults[i].push_back(SymbolTable::mayContainSymbolRefs(t));
+    });
+  while (ready.load() < numThreads)
+    ;
+  go.store(true);
+  for (std::thread &t : threads)
+    t.join();
+
+  for (int i = 0; i < numThreads; ++i) {
+    ASSERT_EQ(attrResults[i].size(), attrTruth.size());
+    ASSERT_EQ(typeResults[i].size(), typeTruth.size());
+    EXPECT_EQ(attrResults[i], attrTruth) << "thread " << i;
+    EXPECT_EQ(typeResults[i], typeTruth) << "thread " << i;
+  }
+}
+
+} // namespace

>From 3f7c4ba7095a8aa8c622d25a2ed3a855758b7673 Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Wed, 19 Aug 2026 18:06:07 -0500
Subject: [PATCH 4/7] [mlir] Pack the symbol-ref containment cache and hoist
 the verify context

Replace the containment cache's two DenseMap<const void *, bool> with two
packed open-addressed tables, and stop re-resolving the per-scope verification
context on every containment probe.

Each table is a flat array of single-word slots. A live slot holds the uniqued
opaque pointer with its answer in bit 0 -- both storage families are 8-aligned,
so bit 0 is always free, and a zero word is an empty slot because a uniqued
pointer is never null. A lookup mixes the pointer with one multiply and probes
forward to the key or the first empty slot; an insert places the key at that
empty slot or returns the resident answer, so a racing duplicate fill is a
no-op. The tables keep the cache's existing SmartRWMutex: a lookup holds the
read lock for its whole probe and an insert -- including the growth it may
trigger at half load -- holds the write lock, so readers never observe a
half-grown table and the plain arrays need no per-slot atomics. This drops the
hash mix and the DenseMap bucket objects from the probe; the answer still lives
behind the same lock, in a cheaper container.

verifySymbolTable now resolves the context, its multithreading flag, and the
containment cache once per scope and threads them through the per-operation
type/attribute checks and the fill recursion, rather than re-deriving all three
per probe and per recursion level. The properties-attribute scratch list is
hoisted out of the per-operation loop and cleared before each population so its
buffer is reused without accumulating attributes across operations. The fill
stops descending a value's sub-elements once one has answered may-contain. The
root gates are kept: skipping a provably reference-free root before entering the
walker spares its out-of-line descent and keeps its per-scope visited memo
confined to subtrees that can carry a symbol use.

The set of things verified, the verdicts, and the diagnostics are unchanged.

Assisted-by: Claude Code (Anthropic)
---
 mlir/lib/IR/SymbolRefContainmentCache.h       | 136 ++++++++++++++++--
 mlir/lib/IR/SymbolTable.cpp                   | 103 ++++++++-----
 .../IR/SymbolReferenceContainmentTest.cpp     |  70 +++++++++
 3 files changed, 261 insertions(+), 48 deletions(-)

diff --git a/mlir/lib/IR/SymbolRefContainmentCache.h b/mlir/lib/IR/SymbolRefContainmentCache.h
index 6bc4f83055ad3..a9ca2b346b304 100644
--- a/mlir/lib/IR/SymbolRefContainmentCache.h
+++ b/mlir/lib/IR/SymbolRefContainmentCache.h
@@ -24,20 +24,41 @@
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Types.h"
-#include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/MathExtras.h"
 #include "llvm/Support/RWMutex.h"
+#include <cstdint>
 #include <optional>
+#include <vector>
 
 namespace mlir {
 class MLIRContext;
 namespace detail {
 
 /// Per-context store of the "may transitively contain a SymbolRefAttr" fact for
-/// uniqued types and attributes. Two maps mirror the context's two uniquers, so
-/// type and attribute opaque pointers never need to be argued disjoint. The
-/// `lock` flag threaded through each operation is the context's runtime
+/// uniqued types and attributes. Two tables mirror the context's two uniquers,
+/// so type and attribute opaque pointers never need to be argued disjoint.
+///
+/// Each table is a packed open-addressed array of single-word slots. A live
+/// slot holds the uniqued opaque pointer with its answer in bit 0; both storage
+/// families are 8-aligned, so bit 0 is always free to carry the fact (the
+/// static_asserts below stand watch over that). A zero word is an empty slot,
+/// which a live entry can never collide with because a uniqued pointer is never
+/// null.
+///
+/// Locking discipline: the tables live entirely behind one SmartRWMutex. A
+/// lookup holds the read lock for its whole probe; an insert -- including any
+/// table growth it triggers -- holds the write lock. Readers therefore never
+/// observe a half-grown table and the plain arrays need no per-slot atomics.
+/// The `lock` flag threaded through each operation is the context's runtime
 /// multithreading flag: when it is false the store is touched single-threaded
 /// and no lock is taken, mirroring MLIRContext's ScopedWriterLock.
+///
+/// Probe contract: a lookup probes forward from the mixed home slot until it
+/// meets the key (a hit) or an empty slot (a miss); it is never bounded short
+/// of an empty slot. Growth keeps the load factor at or below one half, so an
+/// empty slot always exists and every probe terminates. Insert places a key at
+/// the first empty slot on its probe, or returns the resident answer if the key
+/// is already present, so a racing duplicate fill is a no-op.
 class SymbolRefContainmentCache {
 public:
   SymbolRefContainmentCache() = default;
@@ -64,29 +85,116 @@ class SymbolRefContainmentCache {
   }
 
 private:
-  using Map = DenseMap<const void *, bool>;
+  // Both uniqued-storage families are 8-aligned, so bit 0 of a stored pointer
+  // is free to hold the cached answer.
+  static_assert(
+      alignof(TypeStorage) >= 8,
+      "type storage must be 8-aligned so bit 0 of its pointer is free "
+      "to tag the cached answer");
+  static_assert(
+      alignof(AttributeStorage) >= 8,
+      "attribute storage must be 8-aligned so bit 0 of its pointer is "
+      "free to tag the cached answer");
 
-  std::optional<bool> lookupImpl(const Map &map, const void *key,
+  /// One packed open-addressed table of tagged-pointer facts. It is not
+  /// self-synchronizing: the enclosing cache serializes all access with its
+  /// SmartRWMutex, so every method here runs under the read lock (lookup) or
+  /// the write lock (insert and the growth it may trigger).
+  struct Table {
+    // Power-of-two length, or empty before the first insert. 0 == empty slot;
+    // a live slot is (uniqued pointer | answer bit).
+    std::vector<uintptr_t> slots;
+    size_t count = 0;
+
+    std::optional<bool> lookup(uintptr_t key) const {
+      if (slots.empty())
+        return std::nullopt;
+      unsigned shift = 64u - llvm::Log2_64(slots.size());
+      size_t mask = slots.size() - 1;
+      for (size_t i = home(key, shift);; i = (i + 1) & mask) {
+        uintptr_t word = slots[i];
+        if (word == 0)
+          return std::nullopt;
+        if ((word & ~static_cast<uintptr_t>(1)) == key)
+          return static_cast<bool>(word & 1);
+      }
+    }
+
+    bool insert(uintptr_t key, bool value) {
+      // Grow before this insert could push the load past one half, so a lookup
+      // is always guaranteed an empty slot to terminate on.
+      if (slots.empty() || (count + 1) * 2 > slots.size())
+        grow();
+      for (;;) {
+        unsigned shift = 64u - llvm::Log2_64(slots.size());
+        size_t mask = slots.size() - 1;
+        size_t probe = 0;
+        for (size_t i = home(key, shift);; i = (i + 1) & mask) {
+          uintptr_t word = slots[i];
+          if (word == 0) {
+            slots[i] = key | static_cast<uintptr_t>(value);
+            ++count;
+            return value;
+          }
+          if ((word & ~static_cast<uintptr_t>(1)) == key)
+            return static_cast<bool>(word & 1);
+          // A cluster longer than the bound is a sign the mix has degenerated
+          // for this address run; grow to break it up rather than lengthen it.
+          if (++probe > kMaxProbe)
+            break;
+        }
+        grow();
+      }
+    }
+
+  private:
+    // Fibonacci hashing: one multiply by the 64-bit golden-ratio constant, then
+    // take the top log2(capacity) bits. Bump-allocated storage pointers advance
+    // in regular strides, which a bare mask would fold into long clusters; the
+    // multiply spreads those strides across the whole table.
+    static size_t home(uintptr_t key, unsigned shift) {
+      uint64_t mixed = static_cast<uint64_t>(key >> 3) * 0x9E3779B97F4A7C15ULL;
+      return static_cast<size_t>(mixed >> shift);
+    }
+
+    void grow() {
+      size_t newCapacity = slots.empty() ? kInitialCapacity : slots.size() * 2;
+      std::vector<uintptr_t> old = std::move(slots);
+      slots.assign(newCapacity, 0);
+      unsigned shift = 64u - llvm::Log2_64(newCapacity);
+      size_t mask = newCapacity - 1;
+      for (uintptr_t word : old) {
+        if (word == 0)
+          continue;
+        size_t i = home(word & ~static_cast<uintptr_t>(1), shift);
+        while (slots[i] != 0)
+          i = (i + 1) & mask;
+        slots[i] = word;
+      }
+    }
+
+    static constexpr size_t kInitialCapacity = 8;
+    static constexpr size_t kMaxProbe = 16;
+  };
+
+  std::optional<bool> lookupImpl(const Table &table, const void *key,
                                  bool lock) const {
     std::optional<llvm::sys::SmartScopedReader<true>> guard;
     if (lock)
       guard.emplace(mutex);
-    auto it = map.find(key);
-    if (it == map.end())
-      return std::nullopt;
-    return it->second;
+    return table.lookup(reinterpret_cast<uintptr_t>(key));
   }
 
-  bool insertImpl(Map &map, const void *key, bool value, bool lock) {
+  bool insertImpl(Table &table, const void *key, bool value, bool lock) {
     std::optional<llvm::sys::SmartScopedWriter<true>> guard;
     if (lock)
       guard.emplace(mutex);
-    return map.try_emplace(key, value).first->second;
+    return table.insert(reinterpret_cast<uintptr_t>(key), value);
   }
 
   mutable llvm::sys::SmartRWMutex<true> mutex;
-  Map typeEntries;
-  Map attrEntries;
+  Table typeEntries;
+  Table attrEntries;
 };
 
 /// Return the given context's symbol-reference containment cache. Defined in
diff --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index 7f6445c84991f..ca4f6d0a64883 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -477,23 +477,26 @@ raw_ostream &mlir::operator<<(raw_ostream &os,
 // SymbolRefAttr containment query
 //===----------------------------------------------------------------------===//
 
-static bool computeMayContainSymbolRefs(Type type);
-static bool computeMayContainSymbolRefs(Attribute attr);
+static bool
+computeMayContainSymbolRefs(Type type, bool lock,
+                            detail::SymbolRefContainmentCache &cache);
+static bool
+computeMayContainSymbolRefs(Attribute attr, bool lock,
+                            detail::SymbolRefContainmentCache &cache);
 
 /// Fill (once) and return the "may transitively contain a SymbolRefAttr" fact
 /// for a uniqued type or attribute `obj`, seeded by `selfIsRef` for an object
-/// that is itself a symbol reference. The answer is computed entirely outside
-/// the cache lock -- which is taken only for the write-once insert -- so no
-/// lock is held across the recursion and a racing duplicate fill computes the
-/// same fact idempotently. A mutable-storage kind may gain sub-elements after
-/// this point, so its contents are never read; it and its containers report
-/// may-contain conservatively.
+/// that is itself a symbol reference. The context's runtime multithreading flag
+/// `lock` and its containment `cache` are resolved once per verification scope
+/// and threaded through the recursion, so no probe re-derives them. The answer
+/// is computed entirely outside the cache lock -- which is taken only for the
+/// write-once insert -- so no lock is held across the recursion and a racing
+/// duplicate fill computes the same fact idempotently. A mutable-storage kind
+/// may gain sub-elements after this point, so its contents are never read; it
+/// and its containers report may-contain conservatively.
 template <typename T>
-static bool fillMayContainSymbolRefs(T obj, bool selfIsRef) {
-  MLIRContext *ctx = obj.getContext();
-  bool lock = ctx->isMultithreadingEnabled();
-  detail::SymbolRefContainmentCache &cache =
-      detail::getSymbolRefContainmentCache(ctx);
+static bool fillMayContainSymbolRefs(T obj, bool selfIsRef, bool lock,
+                                     detail::SymbolRefContainmentCache &cache) {
   if (std::optional<bool> cached = cache.lookup(obj, lock))
     return *cached;
 
@@ -501,32 +504,46 @@ static bool fillMayContainSymbolRefs(T obj, bool selfIsRef) {
       selfIsRef || obj.template hasTrait<detail::StorageUserTrait::IsMutable>();
   // The recursion needs no in-progress guard: immutable objects form a DAG
   // (sub-elements are interned before their parents), so every cycle passes
-  // through a mutable kind, where the fill stops above before descending.
+  // through a mutable kind, where the fill stops above before descending. Once
+  // one sub-element has answered may-contain, the rest need not be visited:
+  // walkImmediateSubElements cannot be interrupted, but the callback is a no-op
+  // once `mayContain` is set.
   if (!mayContain)
     obj.walkImmediateSubElements(
         [&](Attribute sub) {
-          mayContain |= sub && computeMayContainSymbolRefs(sub);
+          if (!mayContain && sub)
+            mayContain = computeMayContainSymbolRefs(sub, lock, cache);
         },
         [&](Type sub) {
-          mayContain |= sub && computeMayContainSymbolRefs(sub);
+          if (!mayContain && sub)
+            mayContain = computeMayContainSymbolRefs(sub, lock, cache);
         });
   return cache.insert(obj, mayContain, lock);
 }
 
 // A type is never itself a SymbolRefAttr; an attribute is one exactly when
 // isa<SymbolRefAttr> holds (covering FlatSymbolRefAttr).
-static bool computeMayContainSymbolRefs(Type type) {
-  return fillMayContainSymbolRefs(type, /*selfIsRef=*/false);
+static bool
+computeMayContainSymbolRefs(Type type, bool lock,
+                            detail::SymbolRefContainmentCache &cache) {
+  return fillMayContainSymbolRefs(type, /*selfIsRef=*/false, lock, cache);
 }
-static bool computeMayContainSymbolRefs(Attribute attr) {
-  return fillMayContainSymbolRefs(attr, /*selfIsRef=*/isa<SymbolRefAttr>(attr));
+static bool
+computeMayContainSymbolRefs(Attribute attr, bool lock,
+                            detail::SymbolRefContainmentCache &cache) {
+  return fillMayContainSymbolRefs(attr, /*selfIsRef=*/isa<SymbolRefAttr>(attr),
+                                  lock, cache);
 }
 
 bool SymbolTable::mayContainSymbolRefs(Type type) {
-  return computeMayContainSymbolRefs(type);
+  MLIRContext *ctx = type.getContext();
+  return computeMayContainSymbolRefs(type, ctx->isMultithreadingEnabled(),
+                                     detail::getSymbolRefContainmentCache(ctx));
 }
 bool SymbolTable::mayContainSymbolRefs(Attribute attr) {
-  return computeMayContainSymbolRefs(attr);
+  MLIRContext *ctx = attr.getContext();
+  return computeMayContainSymbolRefs(attr, ctx->isMultithreadingEnabled(),
+                                     detail::getSymbolRefContainmentCache(ctx));
 }
 
 //===----------------------------------------------------------------------===//
@@ -540,17 +557,22 @@ bool SymbolTable::mayContainSymbolRefs(Attribute attr) {
 /// makes each uniqued type, which may recur across many positions and
 /// operations, verified against the enclosing symbol table at most once.
 /// Verification fails fast on the first invalid symbol use.
-static LogicalResult verifyOpTypeSymbolUses(Operation *op,
-                                            AttrTypeWalker &walker) {
-  // A root that provably contains no SymbolRefAttr is not worth entering; the
-  // walker's callbacks prune interior subtrees the same way.
+static LogicalResult
+verifyOpTypeSymbolUses(Operation *op, AttrTypeWalker &walker, bool lock,
+                       detail::SymbolRefContainmentCache &cache,
+                       NamedAttrList &inherentAttrs) {
+  // Skip a root that provably holds no SymbolRefAttr before entering the walker
+  // at all: the cheap containment probe spares the walker's out-of-line descent
+  // and keeps its per-scope visited memo confined to subtrees that can actually
+  // carry a symbol use, rather than growing it by every distinct clean root.
+  // The walker's callbacks prune interior subtrees on the same fact.
   auto verify = [&](Type type) {
-    if (!SymbolTable::mayContainSymbolRefs(type))
+    if (!computeMayContainSymbolRefs(type, lock, cache))
       return WalkResult::advance();
     return walker.walk<WalkOrder::PreOrder>(type);
   };
   auto verifyAttr = [&](Attribute attr) {
-    if (!attr || !SymbolTable::mayContainSymbolRefs(attr))
+    if (!attr || !computeMayContainSymbolRefs(attr, lock, cache))
       return WalkResult::advance();
     return walker.walk<WalkOrder::PreOrder>(attr);
   };
@@ -574,14 +596,16 @@ static LogicalResult verifyOpTypeSymbolUses(Operation *op,
   // that keeps its inherent attributes in properties. The raw dictionary
   // already covers inherent attributes for operations that do not use
   // properties, and the discardable attributes otherwise; the properties-held
-  // inherent attributes are appended into a stack-local NamedAttrList via
+  // inherent attributes are appended into a scope-reused NamedAttrList via
   // populateInherentAttrs and walked separately -- the same coverage
   // getAttrDictionary() provides, since it calls the same
-  // populateInherentAttrs.
+  // populateInherentAttrs. The list is cleared before each population so its
+  // heap buffer is reused across operations without accumulating their
+  // attributes.
   if (verifyAttr(op->getRawDictionaryAttrs()).wasInterrupted())
     return failure();
   if (op->getPropertiesStorageSize()) {
-    NamedAttrList inherentAttrs;
+    inherentAttrs.clear();
     op->getName().populateInherentAttrs(op, inherentAttrs);
     for (const NamedAttribute &namedAttr : inherentAttrs)
       if (verifyAttr(namedAttr.getValue()).wasInterrupted())
@@ -627,6 +651,16 @@ LogicalResult detail::verifySymbolTable(Operation *op) {
   // most once across the whole scope.
   SetVector<Attribute> verifiedAttrs;
 
+  // Resolve the containment cache, the context's multithreading flag, and (via
+  // the hoisted NamedAttrList) the properties-attribute scratch buffer once for
+  // the whole scope, then thread them through every containment probe and the
+  // fill recursion, rather than re-deriving them per type occurrence.
+  MLIRContext *ctx = op->getContext();
+  bool cacheLock = ctx->isMultithreadingEnabled();
+  detail::SymbolRefContainmentCache &cache =
+      detail::getSymbolRefContainmentCache(ctx);
+  NamedAttrList inherentAttrs;
+
   // A single walker, shared across the whole scope, checks the symbol uses of
   // every SymbolUserTypeInterface type. Its visited memo records each uniqued
   // type and attribute once, so a subtree recurring across operand, result,
@@ -641,7 +675,7 @@ LogicalResult detail::verifySymbolTable(Operation *op) {
     // Prune subtrees that provably hold no SymbolRefAttr: a conforming
     // SymbolUserTypeInterface spells its references as SymbolRefAttr
     // sub-elements, so such a type has a vacuous verifySymbolUses.
-    if (!SymbolTable::mayContainSymbolRefs(type))
+    if (!computeMayContainSymbolRefs(type, cacheLock, cache))
       return WalkResult::skip();
     if (auto user = dyn_cast<SymbolUserTypeInterface>(type))
       if (failed(user.verifySymbolUses(typeSymbolUseAnchor, symbolTable)))
@@ -651,7 +685,7 @@ LogicalResult detail::verifySymbolTable(Operation *op) {
   typeWalker.addWalk([&](Attribute attr) -> WalkResult {
     // Prune reference-free attribute subtrees so the walk descends only where a
     // SymbolRefAttr, and thus a symbol-using type, may live.
-    if (!SymbolTable::mayContainSymbolRefs(attr))
+    if (!computeMayContainSymbolRefs(attr, cacheLock, cache))
       return WalkResult::skip();
     return WalkResult::advance();
   });
@@ -669,7 +703,8 @@ LogicalResult detail::verifySymbolTable(Operation *op) {
       }
     }
     typeSymbolUseAnchor = op;
-    if (failed(verifyOpTypeSymbolUses(op, typeWalker)))
+    if (failed(verifyOpTypeSymbolUses(op, typeWalker, cacheLock, cache,
+                                      inherentAttrs)))
       return WalkResult::interrupt();
     return WalkResult::advance();
   };
diff --git a/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp b/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
index b3a6bd76ee3c2..15a4bd79cabf4 100644
--- a/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
+++ b/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
@@ -26,6 +26,7 @@
 #include "mlir/Parser/Parser.h"
 #include "gtest/gtest.h"
 
+#include "../../lib/IR/SymbolRefContainmentCache.h"
 #include "../../test/lib/Dialect/Test/TestAttributes.h"
 #include "../../test/lib/Dialect/Test/TestDialect.h"
 #include "../../test/lib/Dialect/Test/TestTypes.h"
@@ -333,4 +334,73 @@ TEST_F(SymbolReferenceContainmentTest, ConcurrentFillIsRaceFree) {
   }
 }
 
+// Growth preserves every earlier fact. The cache's internal table starts
+// minimal and doubles as it fills, so inserting far past the growth trigger
+// forces several rehashes; open addressing must re-home every live slot on each
+// grow. Synthetic 8-aligned pointers in a regular stride stand in for the
+// bump-allocated storage addresses the cache keys on; the keys are opaque and
+// never dereferenced. Because a probe runs to the empty slot, preservation
+// holds for any home function, so this exercises the rehash, not the mix -- the
+// mix instead keeps the table's capacity bounded when storage addresses arrive
+// in a stride, which this test does not measure.
+TEST_F(SymbolReferenceContainmentTest, TableGrowthPreservesEntries) {
+  detail::SymbolRefContainmentCache cache;
+  const int n = 5000; // many doublings past the minimal table
+  std::vector<Type> keys;
+  keys.reserve(n);
+  for (int i = 0; i < n; ++i) {
+    auto *p = reinterpret_cast<const void *>(
+        static_cast<uintptr_t>(0x100000 + i * 8));
+    Type key = Type::getFromOpaquePointer(p);
+    bool value = (i % 3 == 0); // a deterministic mix of true/false facts
+    // The first insert of a key returns the value it records.
+    EXPECT_EQ(cache.insert(key, value, /*lock=*/false), value);
+    keys.push_back(key);
+  }
+  // After all the growth, every earlier fact is still found with its value.
+  for (int i = 0; i < n; ++i) {
+    std::optional<bool> got = cache.lookup(keys[i], /*lock=*/false);
+    ASSERT_TRUE(got.has_value()) << "lost entry " << i;
+    EXPECT_EQ(*got, (i % 3 == 0)) << "entry " << i;
+  }
+  // A key never inserted is a miss, not a stray cluster hit.
+  Type absent = Type::getFromOpaquePointer(
+      reinterpret_cast<const void *>(static_cast<uintptr_t>(0x100000 + n * 8)));
+  EXPECT_FALSE(cache.lookup(absent, /*lock=*/false).has_value());
+}
+
+// A DistinctAttr is keyed by the address of its own storage, which comes from a
+// separate always-allocating allocator rather than the attribute uniquer; the
+// cache must tag and recover that pointer just like a uniqued one. Its
+// 8-alignment, which frees bit 0 for the answer tag, is asserted at compile
+// time beside the table; this exercises the runtime path.
+TEST_F(SymbolReferenceContainmentTest, DistinctAttrIsHandled) {
+  DistinctAttr d1 = DistinctAttr::create(UnitAttr::get(&context));
+  DistinctAttr d2 = DistinctAttr::create(UnitAttr::get(&context));
+  ASSERT_NE(d1, d2); // always-allocating: distinct instances, distinct pointers
+
+  // The separate allocator must still hand back 8-aligned storage, so bit 0 of
+  // the opaque pointer is free for the answer tag.
+  EXPECT_EQ(
+      reinterpret_cast<uintptr_t>(Attribute(d1).getAsOpaquePointer()) & 1u, 0u);
+
+  // Through the public query a DistinctAttr exposes no SymbolRefAttr
+  // sub-element, so it answers false, stably across the cold fill and the warm
+  // cache hit.
+  bool cold = SymbolTable::mayContainSymbolRefs(d1);
+  EXPECT_FALSE(cold);
+  EXPECT_EQ(SymbolTable::mayContainSymbolRefs(d1), cold);
+  EXPECT_FALSE(SymbolTable::mayContainSymbolRefs(d2));
+
+  // Drive the tag round-trip with bit 0 set on a distinct-allocator pointer: a
+  // forced true must survive the insert and be returned by the warm lookup,
+  // while an uninserted distinct pointer stays a miss.
+  detail::SymbolRefContainmentCache cache;
+  EXPECT_TRUE(cache.insert(d1, /*value=*/true, /*lock=*/false));
+  std::optional<bool> warm = cache.lookup(d1, /*lock=*/false);
+  ASSERT_TRUE(warm.has_value());
+  EXPECT_TRUE(*warm);
+  EXPECT_FALSE(cache.lookup(d2, /*lock=*/false).has_value());
+}
+
 } // namespace

>From 7daf50e40b4de0b255a0e0e993a7faa6a5f02ae2 Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Wed, 19 Aug 2026 20:42:14 -0500
Subject: [PATCH 5/7] [mlir] Record only reference-free objects in the
 symbol-ref containment cache

Replace each uniquer's packed open-addressed table of tagged-pointer facts with
one llvm::DenseSet<const void *> of the objects proven free of a transitive
SymbolRefAttr, behind the same SmartRWMutex.

Only the "provably reference-free" fact is worth storing. A may-contain answer
is recomputed cheaply on each encounter because the fill early-exits at the
first SymbolRefAttr, so it need never be recorded: membership in the set means
clear (answer false), and non-membership collapses "not yet filled" and
"may-contain" into one state the caller recomputes. lookup returns false for a
recorded-clear object and nullopt otherwise, never true; insert records
membership only for a clear object and otherwise returns may-contain without
touching the set. A mutable-storage kind is therefore conservatively may-contain
forever with no special handling.

This drops the bespoke table, its Fibonacci pointer mix, and its growth and
re-homing machinery; the low-bit answer tag is gone, so the storage-alignment
static_asserts that guarded it go with it. The set keeps the cache's existing
locking discipline: a lookup holds the read lock for its whole probe and an
insert -- including any set growth -- holds the write lock, so readers never
observe a half-grown set.

The answers the cache serves, the verification verdicts, and the diagnostics are
unchanged.

Assisted-by: Claude Code (Anthropic)
---
 mlir/lib/IR/SymbolRefContainmentCache.h       | 188 ++++++------------
 .../IR/SymbolReferenceContainmentTest.cpp     |  72 +++----
 2 files changed, 96 insertions(+), 164 deletions(-)

diff --git a/mlir/lib/IR/SymbolRefContainmentCache.h b/mlir/lib/IR/SymbolRefContainmentCache.h
index a9ca2b346b304..0804f2a102f97 100644
--- a/mlir/lib/IR/SymbolRefContainmentCache.h
+++ b/mlir/lib/IR/SymbolRefContainmentCache.h
@@ -6,16 +6,16 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// A context-owned store recording, per uniqued type and attribute, whether it
-// may transitively contain a SymbolRefAttr. Symbol-table verification consults
-// it to prune the symbol-use walk.
+// A context-owned store recording which uniqued types and attributes are
+// provably free of a transitive SymbolRefAttr. Symbol-table verification
+// consults it to prune the symbol-use walk.
 //
 // The store is filled lazily and never invalidated. Uniqued storage is immortal
-// and immutable for the context's lifetime, so a cached answer can never go
+// and immutable for the context's lifetime, so a recorded answer can never go
 // stale, and a pointer can never be recycled to alias another object within one
-// context. Every entry is write-once: false only for a provably
-// reference-free immutable subtree, true for everything else, including
-// mutable-storage kinds, whose contents the fill never reads.
+// context. Only the "provably reference-free" fact is recorded; a may-contain
+// answer -- including every mutable-storage kind, whose contents the fill never
+// reads -- is left unrecorded and recomputed on each encounter.
 //
 //===----------------------------------------------------------------------===//
 
@@ -24,41 +24,34 @@
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Types.h"
-#include "llvm/Support/MathExtras.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/RWMutex.h"
-#include <cstdint>
 #include <optional>
-#include <vector>
 
 namespace mlir {
 class MLIRContext;
 namespace detail {
 
-/// Per-context store of the "may transitively contain a SymbolRefAttr" fact for
-/// uniqued types and attributes. Two tables mirror the context's two uniquers,
-/// so type and attribute opaque pointers never need to be argued disjoint.
+/// Per-context store of the "provably free of a transitive SymbolRefAttr" fact
+/// for uniqued types and attributes. Two sets mirror the context's two
+/// uniquers, so type and attribute opaque pointers never need to be argued
+/// disjoint.
 ///
-/// Each table is a packed open-addressed array of single-word slots. A live
-/// slot holds the uniqued opaque pointer with its answer in bit 0; both storage
-/// families are 8-aligned, so bit 0 is always free to carry the fact (the
-/// static_asserts below stand watch over that). A zero word is an empty slot,
-/// which a live entry can never collide with because a uniqued pointer is never
-/// null.
+/// Each set holds only the opaque pointers of objects proven reference-free.
+/// Membership means "clear" (answer false); non-membership means "not yet
+/// proven clear", which the fill treats as unfilled and may-contain both at
+/// once -- it recomputes containment, and a genuine may-contain object is
+/// recomputed cheaply because the walk early-exits at its first SymbolRefAttr.
+/// A may-contain answer is therefore never stored, so a mutable-storage kind is
+/// conservatively may-contain forever with no special handling.
 ///
-/// Locking discipline: the tables live entirely behind one SmartRWMutex. A
-/// lookup holds the read lock for its whole probe; an insert -- including any
-/// table growth it triggers -- holds the write lock. Readers therefore never
-/// observe a half-grown table and the plain arrays need no per-slot atomics.
-/// The `lock` flag threaded through each operation is the context's runtime
-/// multithreading flag: when it is false the store is touched single-threaded
-/// and no lock is taken, mirroring MLIRContext's ScopedWriterLock.
-///
-/// Probe contract: a lookup probes forward from the mixed home slot until it
-/// meets the key (a hit) or an empty slot (a miss); it is never bounded short
-/// of an empty slot. Growth keeps the load factor at or below one half, so an
-/// empty slot always exists and every probe terminates. Insert places a key at
-/// the first empty slot on its probe, or returns the resident answer if the key
-/// is already present, so a racing duplicate fill is a no-op.
+/// Locking discipline: the sets live entirely behind one SmartRWMutex. A lookup
+/// holds the read lock for its whole probe; an insert -- including any set
+/// growth it triggers -- holds the write lock. Readers therefore never observe
+/// a half-grown set and the plain sets need no per-slot atomics. The `lock`
+/// flag threaded through each operation is the context's runtime multithreading
+/// flag: when it is false the store is touched single-threaded and no lock is
+/// taken, mirroring MLIRContext's ScopedWriterLock.
 class SymbolRefContainmentCache {
 public:
   SymbolRefContainmentCache() = default;
@@ -66,7 +59,9 @@ class SymbolRefContainmentCache {
   SymbolRefContainmentCache &
   operator=(const SymbolRefContainmentCache &) = delete;
 
-  /// Return the cached answer for `type`/`attr`, or nullopt if not yet filled.
+  /// Return false for `type`/`attr` if it is recorded clear, or nullopt if it
+  /// is not yet proven clear (unfilled, treated as may-contain by the caller).
+  /// A true answer is never recorded, so lookup never returns true.
   std::optional<bool> lookup(Type type, bool lock) const {
     return lookupImpl(typeEntries, type.getAsOpaquePointer(), lock);
   }
@@ -74,9 +69,11 @@ class SymbolRefContainmentCache {
     return lookupImpl(attrEntries, attr.getAsOpaquePointer(), lock);
   }
 
-  /// Record `value` for `type`/`attr` if no entry exists yet and return the
-  /// resident answer. A racing duplicate fill computes the same immutable fact,
-  /// so the losing insert is a no-op.
+  /// Record `type`/`attr` as clear when `value` is false and return the answer
+  /// unchanged. A may-contain (`value` true) fact is not stored -- it is
+  /// recomputed cheaply on each encounter -- so insert then returns true
+  /// without touching the set. A racing duplicate clear fill records the same
+  /// membership and is a no-op.
   bool insert(Type type, bool value, bool lock) {
     return insertImpl(typeEntries, type.getAsOpaquePointer(), value, lock);
   }
@@ -85,116 +82,47 @@ class SymbolRefContainmentCache {
   }
 
 private:
-  // Both uniqued-storage families are 8-aligned, so bit 0 of a stored pointer
-  // is free to hold the cached answer.
-  static_assert(
-      alignof(TypeStorage) >= 8,
-      "type storage must be 8-aligned so bit 0 of its pointer is free "
-      "to tag the cached answer");
-  static_assert(
-      alignof(AttributeStorage) >= 8,
-      "attribute storage must be 8-aligned so bit 0 of its pointer is "
-      "free to tag the cached answer");
-
-  /// One packed open-addressed table of tagged-pointer facts. It is not
-  /// self-synchronizing: the enclosing cache serializes all access with its
-  /// SmartRWMutex, so every method here runs under the read lock (lookup) or
-  /// the write lock (insert and the growth it may trigger).
-  struct Table {
-    // Power-of-two length, or empty before the first insert. 0 == empty slot;
-    // a live slot is (uniqued pointer | answer bit).
-    std::vector<uintptr_t> slots;
-    size_t count = 0;
-
-    std::optional<bool> lookup(uintptr_t key) const {
-      if (slots.empty())
-        return std::nullopt;
-      unsigned shift = 64u - llvm::Log2_64(slots.size());
-      size_t mask = slots.size() - 1;
-      for (size_t i = home(key, shift);; i = (i + 1) & mask) {
-        uintptr_t word = slots[i];
-        if (word == 0)
-          return std::nullopt;
-        if ((word & ~static_cast<uintptr_t>(1)) == key)
-          return static_cast<bool>(word & 1);
-      }
-    }
-
-    bool insert(uintptr_t key, bool value) {
-      // Grow before this insert could push the load past one half, so a lookup
-      // is always guaranteed an empty slot to terminate on.
-      if (slots.empty() || (count + 1) * 2 > slots.size())
-        grow();
-      for (;;) {
-        unsigned shift = 64u - llvm::Log2_64(slots.size());
-        size_t mask = slots.size() - 1;
-        size_t probe = 0;
-        for (size_t i = home(key, shift);; i = (i + 1) & mask) {
-          uintptr_t word = slots[i];
-          if (word == 0) {
-            slots[i] = key | static_cast<uintptr_t>(value);
-            ++count;
-            return value;
-          }
-          if ((word & ~static_cast<uintptr_t>(1)) == key)
-            return static_cast<bool>(word & 1);
-          // A cluster longer than the bound is a sign the mix has degenerated
-          // for this address run; grow to break it up rather than lengthen it.
-          if (++probe > kMaxProbe)
-            break;
-        }
-        grow();
-      }
-    }
-
-  private:
-    // Fibonacci hashing: one multiply by the 64-bit golden-ratio constant, then
-    // take the top log2(capacity) bits. Bump-allocated storage pointers advance
-    // in regular strides, which a bare mask would fold into long clusters; the
-    // multiply spreads those strides across the whole table.
-    static size_t home(uintptr_t key, unsigned shift) {
-      uint64_t mixed = static_cast<uint64_t>(key >> 3) * 0x9E3779B97F4A7C15ULL;
-      return static_cast<size_t>(mixed >> shift);
+  /// One uniquer's clear-object set. It is not self-synchronizing: the
+  /// enclosing cache serializes all access with its SmartRWMutex, so every
+  /// method here runs under the read lock (lookup) or the write lock (insert
+  /// and the growth it may trigger).
+  struct ClearSet {
+    llvm::DenseSet<const void *> clear;
+
+    std::optional<bool> lookup(const void *key) const {
+      if (clear.count(key))
+        return false;
+      return std::nullopt;
     }
 
-    void grow() {
-      size_t newCapacity = slots.empty() ? kInitialCapacity : slots.size() * 2;
-      std::vector<uintptr_t> old = std::move(slots);
-      slots.assign(newCapacity, 0);
-      unsigned shift = 64u - llvm::Log2_64(newCapacity);
-      size_t mask = newCapacity - 1;
-      for (uintptr_t word : old) {
-        if (word == 0)
-          continue;
-        size_t i = home(word & ~static_cast<uintptr_t>(1), shift);
-        while (slots[i] != 0)
-          i = (i + 1) & mask;
-        slots[i] = word;
-      }
+    bool insert(const void *key, bool value) {
+      // Only the clear fact is durable; a may-contain answer is left out to be
+      // recomputed cheaply on each encounter. A racing duplicate clear fill
+      // records the same membership and is a no-op.
+      if (!value)
+        clear.insert(key);
+      return value;
     }
-
-    static constexpr size_t kInitialCapacity = 8;
-    static constexpr size_t kMaxProbe = 16;
   };
 
-  std::optional<bool> lookupImpl(const Table &table, const void *key,
+  std::optional<bool> lookupImpl(const ClearSet &set, const void *key,
                                  bool lock) const {
     std::optional<llvm::sys::SmartScopedReader<true>> guard;
     if (lock)
       guard.emplace(mutex);
-    return table.lookup(reinterpret_cast<uintptr_t>(key));
+    return set.lookup(key);
   }
 
-  bool insertImpl(Table &table, const void *key, bool value, bool lock) {
+  bool insertImpl(ClearSet &set, const void *key, bool value, bool lock) {
     std::optional<llvm::sys::SmartScopedWriter<true>> guard;
     if (lock)
       guard.emplace(mutex);
-    return table.insert(reinterpret_cast<uintptr_t>(key), value);
+    return set.insert(key, value);
   }
 
   mutable llvm::sys::SmartRWMutex<true> mutex;
-  Table typeEntries;
-  Table attrEntries;
+  ClearSet typeEntries;
+  ClearSet attrEntries;
 };
 
 /// Return the given context's symbol-reference containment cache. Defined in
diff --git a/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp b/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
index 15a4bd79cabf4..c0dc2aa25434e 100644
--- a/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
+++ b/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
@@ -334,35 +334,42 @@ TEST_F(SymbolReferenceContainmentTest, ConcurrentFillIsRaceFree) {
   }
 }
 
-// Growth preserves every earlier fact. The cache's internal table starts
-// minimal and doubles as it fills, so inserting far past the growth trigger
-// forces several rehashes; open addressing must re-home every live slot on each
-// grow. Synthetic 8-aligned pointers in a regular stride stand in for the
-// bump-allocated storage addresses the cache keys on; the keys are opaque and
-// never dereferenced. Because a probe runs to the empty slot, preservation
-// holds for any home function, so this exercises the rehash, not the mix -- the
-// mix instead keeps the table's capacity bounded when storage addresses arrive
-// in a stride, which this test does not measure.
-TEST_F(SymbolReferenceContainmentTest, TableGrowthPreservesEntries) {
+// Growth preserves every recorded clear fact. The clear-object set starts empty
+// and grows as it fills, so inserting far past its initial capacity forces
+// several rehashes; every live key must survive each grow. Synthetic pointers
+// in a regular stride stand in for the bump-allocated storage addresses the
+// cache keys on; the keys are opaque and never dereferenced. Only clear facts
+// are recorded, so this checks that each survives the rehash and that
+// may-contain facts, which the store never keeps, stay misses.
+TEST_F(SymbolReferenceContainmentTest, SetGrowthPreservesEntries) {
   detail::SymbolRefContainmentCache cache;
-  const int n = 5000; // many doublings past the minimal table
-  std::vector<Type> keys;
-  keys.reserve(n);
+  const int n = 5000; // many doublings past the minimal set
+  std::vector<Type> clearKeys;
+  std::vector<Type> mayContainKeys;
   for (int i = 0; i < n; ++i) {
     auto *p = reinterpret_cast<const void *>(
         static_cast<uintptr_t>(0x100000 + i * 8));
     Type key = Type::getFromOpaquePointer(p);
-    bool value = (i % 3 == 0); // a deterministic mix of true/false facts
-    // The first insert of a key returns the value it records.
-    EXPECT_EQ(cache.insert(key, value, /*lock=*/false), value);
-    keys.push_back(key);
+    if (i % 3 == 0) {
+      // A may-contain fact is never stored: insert returns true but records
+      // nothing, so a later lookup stays a miss and the caller recomputes.
+      EXPECT_TRUE(cache.insert(key, /*value=*/true, /*lock=*/false));
+      mayContainKeys.push_back(key);
+    } else {
+      // A clear fact is recorded and returned unchanged.
+      EXPECT_FALSE(cache.insert(key, /*value=*/false, /*lock=*/false));
+      clearKeys.push_back(key);
+    }
   }
-  // After all the growth, every earlier fact is still found with its value.
-  for (int i = 0; i < n; ++i) {
-    std::optional<bool> got = cache.lookup(keys[i], /*lock=*/false);
-    ASSERT_TRUE(got.has_value()) << "lost entry " << i;
-    EXPECT_EQ(*got, (i % 3 == 0)) << "entry " << i;
+  // After all the growth, every clear fact is still found as false.
+  for (Type key : clearKeys) {
+    std::optional<bool> got = cache.lookup(key, /*lock=*/false);
+    ASSERT_TRUE(got.has_value()) << "lost clear entry";
+    EXPECT_FALSE(*got);
   }
+  // A may-contain fact was never stored, so it stays a miss.
+  for (Type key : mayContainKeys)
+    EXPECT_FALSE(cache.lookup(key, /*lock=*/false).has_value());
   // A key never inserted is a miss, not a stray cluster hit.
   Type absent = Type::getFromOpaquePointer(
       reinterpret_cast<const void *>(static_cast<uintptr_t>(0x100000 + n * 8)));
@@ -371,19 +378,12 @@ TEST_F(SymbolReferenceContainmentTest, TableGrowthPreservesEntries) {
 
 // A DistinctAttr is keyed by the address of its own storage, which comes from a
 // separate always-allocating allocator rather than the attribute uniquer; the
-// cache must tag and recover that pointer just like a uniqued one. Its
-// 8-alignment, which frees bit 0 for the answer tag, is asserted at compile
-// time beside the table; this exercises the runtime path.
+// cache must record and recover that pointer just like a uniqued one.
 TEST_F(SymbolReferenceContainmentTest, DistinctAttrIsHandled) {
   DistinctAttr d1 = DistinctAttr::create(UnitAttr::get(&context));
   DistinctAttr d2 = DistinctAttr::create(UnitAttr::get(&context));
   ASSERT_NE(d1, d2); // always-allocating: distinct instances, distinct pointers
 
-  // The separate allocator must still hand back 8-aligned storage, so bit 0 of
-  // the opaque pointer is free for the answer tag.
-  EXPECT_EQ(
-      reinterpret_cast<uintptr_t>(Attribute(d1).getAsOpaquePointer()) & 1u, 0u);
-
   // Through the public query a DistinctAttr exposes no SymbolRefAttr
   // sub-element, so it answers false, stably across the cold fill and the warm
   // cache hit.
@@ -392,14 +392,18 @@ TEST_F(SymbolReferenceContainmentTest, DistinctAttrIsHandled) {
   EXPECT_EQ(SymbolTable::mayContainSymbolRefs(d1), cold);
   EXPECT_FALSE(SymbolTable::mayContainSymbolRefs(d2));
 
-  // Drive the tag round-trip with bit 0 set on a distinct-allocator pointer: a
-  // forced true must survive the insert and be returned by the warm lookup,
-  // while an uninserted distinct pointer stays a miss.
+  // A may-contain answer is never recorded: a forced-true insert on a
+  // distinct-allocator pointer returns true but stores nothing, so the warm
+  // lookup stays a miss and the caller recomputes. A clear answer, by contrast,
+  // is recorded and returned false warm, while an uninserted distinct pointer
+  // stays a miss -- the pointer round-trips just like a uniqued one.
   detail::SymbolRefContainmentCache cache;
   EXPECT_TRUE(cache.insert(d1, /*value=*/true, /*lock=*/false));
+  EXPECT_FALSE(cache.lookup(d1, /*lock=*/false).has_value());
+  EXPECT_FALSE(cache.insert(d1, /*value=*/false, /*lock=*/false));
   std::optional<bool> warm = cache.lookup(d1, /*lock=*/false);
   ASSERT_TRUE(warm.has_value());
-  EXPECT_TRUE(*warm);
+  EXPECT_FALSE(*warm);
   EXPECT_FALSE(cache.lookup(d2, /*lock=*/false).has_value());
 }
 

>From e224565f0669c79fb0a2565988c6e653ee58b130 Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Wed, 19 Aug 2026 20:42:14 -0500
Subject: [PATCH 6/7] [mlir] Extract the type/attribute symbol-use verification
 into a scope-lived verifier

verifySymbolTable had grown a long tail of type/attribute containment machinery
inline -- a hoisted cache, threading flag, walker, and scratch list, the two
walker callbacks, and the per-operation type/attribute enumeration -- that
crowded out its structural narrative.

Move that machinery into a file-local OpTypeSymbolUseVerifier. Its fields are the
per-scope state each verification would otherwise re-derive: the containment
cache, the context's multithreading flag, the AttrTypeWalker and its visited
memo, and the NamedAttrList reused for properties-held attributes. It exposes one
entry point, verify(op), and names the attribute-root enumeration --
forEachAttributeRoot -- whose doc comment carries the contract that it covers
exactly what getAttrDictionary() would walk without materializing a fresh
DictionaryAttr per operation.

verifySymbolTable returns to its narrative: the structural checks, the
symbol-name uniqueness check, one construction of the verifier, one verify() per
operation, and the symbol-name/use walk.

This is a pure reorganization; the verification verdicts and diagnostics are
unchanged.

Assisted-by: Claude Code (Anthropic)
---
 mlir/lib/IR/SymbolTable.cpp | 215 ++++++++++++++++++------------------
 1 file changed, 109 insertions(+), 106 deletions(-)

diff --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index ca4f6d0a64883..ab97ede4465eb 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -550,69 +550,112 @@ bool SymbolTable::mayContainSymbolRefs(Attribute attr) {
 // SymbolTable Trait Types
 //===----------------------------------------------------------------------===//
 
-/// Verify the symbol uses held by the types owned by `op`: its operand, result,
-/// and block-argument types, and any types nested within its attributes.
-/// `walker` carries the SymbolUserTypeInterface check as a type-walk callback,
-/// anchored at the operation currently being verified, and its visited memo
-/// makes each uniqued type, which may recur across many positions and
-/// operations, verified against the enclosing symbol table at most once.
-/// Verification fails fast on the first invalid symbol use.
-static LogicalResult
-verifyOpTypeSymbolUses(Operation *op, AttrTypeWalker &walker, bool lock,
-                       detail::SymbolRefContainmentCache &cache,
-                       NamedAttrList &inherentAttrs) {
-  // Skip a root that provably holds no SymbolRefAttr before entering the walker
-  // at all: the cheap containment probe spares the walker's out-of-line descent
-  // and keeps its per-scope visited memo confined to subtrees that can actually
-  // carry a symbol use, rather than growing it by every distinct clean root.
-  // The walker's callbacks prune interior subtrees on the same fact.
-  auto verify = [&](Type type) {
-    if (!computeMayContainSymbolRefs(type, lock, cache))
+namespace {
+/// Verifies the symbol uses carried by the types an operation owns -- its
+/// operand, result, and block-argument types, and the types nested within its
+/// attributes -- against the enclosing symbol table, failing fast on the first
+/// invalid use. One instance serves a whole verification scope and is invoked
+/// once per operation, holding the per-scope state each verification would
+/// otherwise re-derive: the context's containment cache and multithreading
+/// flag, a NamedAttrList reused as the scratch buffer for properties-held
+/// attributes, and a single AttrTypeWalker whose visited memo walks each
+/// uniqued type or attribute -- which may recur across many positions and
+/// operations -- at most once per scope.
+class OpTypeSymbolUseVerifier {
+public:
+  OpTypeSymbolUseVerifier(MLIRContext *ctx, SymbolTableCollection &symbolTable)
+      : lock(ctx->isMultithreadingEnabled()),
+        cache(detail::getSymbolRefContainmentCache(ctx)) {
+    typeWalker.addWalk([this, &symbolTable](Type type) -> WalkResult {
+      // Prune subtrees that provably hold no SymbolRefAttr: a conforming
+      // SymbolUserTypeInterface spells its references as SymbolRefAttr
+      // sub-elements, so such a type has a vacuous verifySymbolUses.
+      if (!computeMayContainSymbolRefs(type, lock, cache))
+        return WalkResult::skip();
+      if (auto user = dyn_cast<SymbolUserTypeInterface>(type))
+        if (failed(user.verifySymbolUses(anchor, symbolTable)))
+          return WalkResult::interrupt();
       return WalkResult::advance();
-    return walker.walk<WalkOrder::PreOrder>(type);
-  };
-  auto verifyAttr = [&](Attribute attr) {
-    if (!attr || !computeMayContainSymbolRefs(attr, lock, cache))
+    });
+    typeWalker.addWalk([this](Attribute attr) -> WalkResult {
+      // Prune reference-free attribute subtrees so the walk descends only where
+      // a SymbolRefAttr, and thus a symbol-using type, may live.
+      if (!computeMayContainSymbolRefs(attr, lock, cache))
+        return WalkResult::skip();
       return WalkResult::advance();
-    return walker.walk<WalkOrder::PreOrder>(attr);
-  };
+    });
+  }
 
-  for (Type type : op->getOperandTypes())
-    if (verify(type).wasInterrupted())
-      return failure();
-  for (Type type : op->getResultTypes())
-    if (verify(type).wasInterrupted())
-      return failure();
-  for (Region &region : op->getRegions())
-    for (Block &block : region)
-      for (BlockArgument argument : block.getArguments())
-        if (verify(argument.getType()).wasInterrupted())
-          return failure();
-
-  // Verify types nested within the operation's attributes, routed through the
-  // same walker (and thus the same visited memo) as the type positions above.
-  // Read the raw stored attribute dictionary rather than getAttrDictionary():
-  // the latter allocates and uniques a fresh DictionaryAttr for every operation
-  // that keeps its inherent attributes in properties. The raw dictionary
-  // already covers inherent attributes for operations that do not use
-  // properties, and the discardable attributes otherwise; the properties-held
-  // inherent attributes are appended into a scope-reused NamedAttrList via
-  // populateInherentAttrs and walked separately -- the same coverage
-  // getAttrDictionary() provides, since it calls the same
-  // populateInherentAttrs. The list is cleared before each population so its
-  // heap buffer is reused across operations without accumulating their
-  // attributes.
-  if (verifyAttr(op->getRawDictionaryAttrs()).wasInterrupted())
-    return failure();
-  if (op->getPropertiesStorageSize()) {
-    inherentAttrs.clear();
-    op->getName().populateInherentAttrs(op, inherentAttrs);
-    for (const NamedAttribute &namedAttr : inherentAttrs)
-      if (verifyAttr(namedAttr.getValue()).wasInterrupted())
+  /// Verify the symbol uses carried by `op`'s types, anchoring each
+  /// SymbolUserTypeInterface check at `op`.
+  LogicalResult verify(Operation *op) {
+    anchor = op;
+    // Skip a root that provably holds no SymbolRefAttr before entering the
+    // walker at all: the cheap containment probe spares the walker's
+    // out-of-line descent and keeps its visited memo confined to subtrees that
+    // can actually carry a symbol use, rather than growing it by every distinct
+    // clean root. The walker's callbacks prune interior subtrees on the same
+    // fact.
+    auto verifyType = [&](Type type) {
+      if (!computeMayContainSymbolRefs(type, lock, cache))
+        return WalkResult::advance();
+      return typeWalker.walk<WalkOrder::PreOrder>(type);
+    };
+    auto verifyAttr = [&](Attribute attr) {
+      if (!attr || !computeMayContainSymbolRefs(attr, lock, cache))
+        return WalkResult::advance();
+      return typeWalker.walk<WalkOrder::PreOrder>(attr);
+    };
+
+    for (Type type : op->getOperandTypes())
+      if (verifyType(type).wasInterrupted())
         return failure();
+    for (Type type : op->getResultTypes())
+      if (verifyType(type).wasInterrupted())
+        return failure();
+    for (Region &region : op->getRegions())
+      for (Block &block : region)
+        for (BlockArgument argument : block.getArguments())
+          if (verifyType(argument.getType()).wasInterrupted())
+            return failure();
+
+    return success(!forEachAttributeRoot(op, verifyAttr).wasInterrupted());
   }
-  return success();
-}
+
+private:
+  /// Invoke `root` on each attribute subtree an operation can carry a nested
+  /// type within: its raw stored attribute dictionary as one root, then each
+  /// properties-held inherent attribute, stopping early once `root` interrupts.
+  /// This covers exactly what getAttrDictionary() would walk without allocating
+  /// and uniquing a fresh DictionaryAttr per operation -- the raw dictionary
+  /// already holds the inherent attributes of operations that keep none in
+  /// properties and the discardable attributes otherwise, and
+  /// populateInherentAttrs supplies the properties-held remainder into the
+  /// scope-reused list, cleared before each population so its buffer serves
+  /// every operation without accumulating attributes across them.
+  WalkResult forEachAttributeRoot(Operation *op,
+                                  function_ref<WalkResult(Attribute)> root) {
+    if (root(op->getRawDictionaryAttrs()).wasInterrupted())
+      return WalkResult::interrupt();
+    if (op->getPropertiesStorageSize()) {
+      inherentAttrs.clear();
+      op->getName().populateInherentAttrs(op, inherentAttrs);
+      for (const NamedAttribute &namedAttr : inherentAttrs)
+        if (root(namedAttr.getValue()).wasInterrupted())
+          return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  }
+
+  bool lock;
+  detail::SymbolRefContainmentCache &cache;
+  AttrTypeWalker typeWalker;
+  NamedAttrList inherentAttrs;
+  // The operation currently being verified; the type-walk anchors each
+  // SymbolUserTypeInterface check here.
+  Operation *anchor = nullptr;
+};
+} // namespace
 
 LogicalResult detail::verifySymbolTable(Operation *op) {
   if (op->getNumRegions() != 1)
@@ -642,53 +685,15 @@ LogicalResult detail::verifySymbolTable(Operation *op) {
     }
   }
 
-  // Verify any nested symbol user operations.
+  // Verify any nested symbol user operations. walkSymbolTable does not descend
+  // into nested symbol tables, so every operation visited here shares one
+  // nearest symbol table; a uniqued attribute or type therefore resolves its
+  // symbol uses identically no matter which operation anchors the lookup, so
+  // each need be verified only once across the whole scope. The type verifier's
+  // walk memo and this attribute set both hold that once-per-scope state.
   SymbolTableCollection symbolTable;
-  // walkSymbolTable does not descend into nested symbol tables, so every
-  // operation visited here shares the same nearest symbol table. A uniqued
-  // attribute or type therefore resolves its symbol uses identically
-  // regardless of which operation anchors the lookup, so each is verified at
-  // most once across the whole scope.
   SetVector<Attribute> verifiedAttrs;
-
-  // Resolve the containment cache, the context's multithreading flag, and (via
-  // the hoisted NamedAttrList) the properties-attribute scratch buffer once for
-  // the whole scope, then thread them through every containment probe and the
-  // fill recursion, rather than re-deriving them per type occurrence.
-  MLIRContext *ctx = op->getContext();
-  bool cacheLock = ctx->isMultithreadingEnabled();
-  detail::SymbolRefContainmentCache &cache =
-      detail::getSymbolRefContainmentCache(ctx);
-  NamedAttrList inherentAttrs;
-
-  // A single walker, shared across the whole scope, checks the symbol uses of
-  // every SymbolUserTypeInterface type. Its visited memo records each uniqued
-  // type and attribute once, so a subtree recurring across operand, result,
-  // block-argument, and attribute positions and across operations is walked
-  // only at its first occurrence, whose operation supplies the lookup anchor;
-  // a re-encounter returns from the memo without re-descending. Because a first
-  // occurrence descends the whole subtree pre-order and verification fails
-  // fast, skipping the re-descent verifies nothing new.
-  Operation *typeSymbolUseAnchor = nullptr;
-  AttrTypeWalker typeWalker;
-  typeWalker.addWalk([&](Type type) -> WalkResult {
-    // Prune subtrees that provably hold no SymbolRefAttr: a conforming
-    // SymbolUserTypeInterface spells its references as SymbolRefAttr
-    // sub-elements, so such a type has a vacuous verifySymbolUses.
-    if (!computeMayContainSymbolRefs(type, cacheLock, cache))
-      return WalkResult::skip();
-    if (auto user = dyn_cast<SymbolUserTypeInterface>(type))
-      if (failed(user.verifySymbolUses(typeSymbolUseAnchor, symbolTable)))
-        return WalkResult::interrupt();
-    return WalkResult::advance();
-  });
-  typeWalker.addWalk([&](Attribute attr) -> WalkResult {
-    // Prune reference-free attribute subtrees so the walk descends only where a
-    // SymbolRefAttr, and thus a symbol-using type, may live.
-    if (!computeMayContainSymbolRefs(attr, cacheLock, cache))
-      return WalkResult::skip();
-    return WalkResult::advance();
-  });
+  OpTypeSymbolUseVerifier typeVerifier(op->getContext(), symbolTable);
 
   auto verifySymbolUserFn = [&](Operation *op) -> std::optional<WalkResult> {
     if (SymbolUserOpInterface user = dyn_cast<SymbolUserOpInterface>(op))
@@ -702,9 +707,7 @@ LogicalResult detail::verifySymbolTable(Operation *op) {
           return WalkResult::interrupt();
       }
     }
-    typeSymbolUseAnchor = op;
-    if (failed(verifyOpTypeSymbolUses(op, typeWalker, cacheLock, cache,
-                                      inherentAttrs)))
+    if (failed(typeVerifier.verify(op)))
       return WalkResult::interrupt();
     return WalkResult::advance();
   };

>From 351aa64eb8426e155f67cffcc5e7d57c790282f4 Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Wed, 19 Aug 2026 23:19:45 -0500
Subject: [PATCH 7/7] [mlir] Give the symbol-ref containment cache sole
 ownership of its query

The "may this type/attribute transitively contain a SymbolRefAttr?" query was
spread across three layers in SymbolTable.cpp -- two public statics, a pair of
compute overloads, and a fill template -- over a store that only ever recorded
the negative fact, behind a tri-state lookup that could never return true.

Move the whole recursive computation into SymbolRefContainmentCache as a
self-recursive private template, leaving the class two public methods:
mayContainSymbolRefs(Type) and mayContainSymbolRefs(Attribute). The negative-only
semantics are now the class's own: one DenseSet of proven-clear opaque pointers,
no ClearSet sub-struct, no typed overload pairs, no std::optional tri-state. The
two uniquers share the single set because uniqued type and attribute storage is
immortal and simultaneously live, so two distinct objects can never share an
address -- the same argument the attribute side already rests on across its two
allocators.

The reshape carries two behavior-identical improvements. The isa<SymbolRefAttr>
seed is evaluated only on a cache miss rather than eagerly on every probe, and
the write lock is taken only when there is a clear fact to record rather than on
every may-contain answer that stores nothing -- the latter a real writer-lock
contention defect under multithreaded verification.

The runtime multithreading flag stops being a parameter threaded through every
call. The cache is constructed with a reference to MLIRContextImpl's
threadingIsEnabled and reads it live at each probe, AND-ing in the compile-time
llvm_is_multithreaded() locally, so no probe makes an out-of-line call and no
surviving boolean is named for the lock rather than for what it is.

OpTypeSymbolUseVerifier keeps its scope-lived shape but sheds the lock member and
folds its two root pre-probes into one generic visit(); forEachAttributeRoot
becomes a concrete method that visits each root directly. Its dead null test goes
-- an operation's raw dictionary and a NamedAttribute value are both non-null by
construction. The root pre-probe stays distinct from the walker callbacks, which
would otherwise pay a memo insertion and a type-erased dispatch per clean root.

SymbolTable::mayContainSymbolRefs had no caller outside the unit test, which now
reaches the query through the lib-internal cache header, so the two public
statics are removed; SymbolTable.h carries no trace of the rework.

The SymbolUserAttrInterface description drops the sub-element-encoding contract:
attribute symbol-use verification runs unconditionally on an operation's
top-level discardable attributes and is never pruned by the containment walk, so
only symbol-using type verification -- and thus only SymbolUserTypeInterface --
depends on that contract.

This is behavior-preserving; the verification verdicts and diagnostics are
unchanged.

Assisted-by: Claude Code (Anthropic)
---
 mlir/include/mlir/IR/SymbolInterfaces.td      |   8 -
 mlir/include/mlir/IR/SymbolTable.h            |  13 -
 mlir/lib/IR/MLIRContext.cpp                   |   9 +-
 mlir/lib/IR/SymbolRefContainmentCache.h       | 158 ++++---
 mlir/lib/IR/SymbolTable.cpp                   | 163 ++-----
 .../IR/SymbolReferenceContainmentTest.cpp     | 403 ++++++------------
 6 files changed, 262 insertions(+), 492 deletions(-)

diff --git a/mlir/include/mlir/IR/SymbolInterfaces.td b/mlir/include/mlir/IR/SymbolInterfaces.td
index 49d7fe5ac53cc..b7790be460180 100644
--- a/mlir/include/mlir/IR/SymbolInterfaces.td
+++ b/mlir/include/mlir/IR/SymbolInterfaces.td
@@ -229,14 +229,6 @@ def SymbolUserAttrInterface : AttrInterface<"SymbolUserAttrInterface"> {
     symbol related utilities that are either costly or otherwise disallowed
     within an operation (e.g., recreating symbol users per op verified rather
     than per symbol table, or querying symbols usage of siblings).
-
-    Implementations must represent the symbols they reference as `SymbolRefAttr`s
-    nested anywhere within their sub-element tree, i.e. reachable by recursive
-    application of `walkImmediateSubElements`, rather than in some other encoding
-    such as a string. This is what lets the symbol machinery cheaply dismiss
-    instances that cannot reference a symbol: an instance with no `SymbolRefAttr`
-    anywhere in its sub-element tree is treated as referencing no symbols, so
-    symbol-table verification may never invoke its `verifySymbolUses`.
   }];
   let cppNamespace = "::mlir";
 
diff --git a/mlir/include/mlir/IR/SymbolTable.h b/mlir/include/mlir/IR/SymbolTable.h
index ea768e0241612..e4790037d37b2 100644
--- a/mlir/include/mlir/IR/SymbolTable.h
+++ b/mlir/include/mlir/IR/SymbolTable.h
@@ -242,19 +242,6 @@ class SymbolTable {
   static bool symbolKnownUseEmpty(StringAttr symbol, Region *from);
   static bool symbolKnownUseEmpty(Operation *symbol, Region *from);
 
-  /// Return whether the given type or attribute may transitively contain a
-  /// `SymbolRefAttr`, i.e. whether one is reachable through its sub-element
-  /// tree by recursive application of `walkImmediateSubElements`. A false
-  /// answer is authoritative: the object provably references no symbol, so a
-  /// SymbolUserTypeInterface / SymbolUserAttrInterface implementation that
-  /// honors its contract has a vacuous `verifySymbolUses` and need never be
-  /// visited. A true answer is conservative: mutable-storage kinds, and
-  /// anything containing them, always report true. The answer is a pure
-  /// function of the uniqued, immutable structure, computed once per object and
-  /// cached on the context for its lifetime.
-  static bool mayContainSymbolRefs(Type type);
-  static bool mayContainSymbolRefs(Attribute attr);
-
   /// Attempt to replace all uses of the given symbol 'oldSymbol' with the
   /// provided symbol 'newSymbol' that are nested within the given operation
   /// 'from'. This does not traverse into any nested symbol tables. If there are
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 6dc200c813dde..e8cfb40a04848 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -270,10 +270,11 @@ class MLIRContextImpl {
   /// destruction.
   DistinctAttributeAllocator distinctAttributeAllocator;
 
-  /// Cache recording, per uniqued type and attribute, whether it may
-  /// transitively contain a SymbolRefAttr. Filled lazily by symbol-table
-  /// verification and never invalidated; see SymbolRefContainmentCache.h.
-  SymbolRefContainmentCache symbolRefContainmentCache;
+  /// Cache recording which uniqued types and attributes are provably free of a
+  /// transitive SymbolRefAttr. Filled on demand by symbol-table verification
+  /// and never invalidated; reads `threadingIsEnabled` live to guard its set.
+  /// See SymbolRefContainmentCache.h.
+  SymbolRefContainmentCache symbolRefContainmentCache{threadingIsEnabled};
 
 public:
   MLIRContextImpl(bool threadingIsEnabled)
diff --git a/mlir/lib/IR/SymbolRefContainmentCache.h b/mlir/lib/IR/SymbolRefContainmentCache.h
index 0804f2a102f97..737db39066e27 100644
--- a/mlir/lib/IR/SymbolRefContainmentCache.h
+++ b/mlir/lib/IR/SymbolRefContainmentCache.h
@@ -6,16 +6,16 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// A context-owned store recording which uniqued types and attributes are
-// provably free of a transitive SymbolRefAttr. Symbol-table verification
-// consults it to prune the symbol-use walk.
+// A context-owned cache answering, for a uniqued type or attribute, whether it
+// may transitively contain a SymbolRefAttr. Symbol-table verification consults
+// it to prune its symbol-use walk to the subtrees that can carry a symbol use.
 //
-// The store is filled lazily and never invalidated. Uniqued storage is immortal
-// and immutable for the context's lifetime, so a recorded answer can never go
-// stale, and a pointer can never be recycled to alias another object within one
-// context. Only the "provably reference-free" fact is recorded; a may-contain
-// answer -- including every mutable-storage kind, whose contents the fill never
-// reads -- is left unrecorded and recomputed on each encounter.
+// Uniqued type and attribute storage is immortal, and immutable in everything
+// the cache reads -- mutable-storage kinds are never read or recorded -- so a
+// recorded answer never goes stale, and two live objects can never share an
+// address. One DenseSet of opaque pointers therefore keys both kinds without
+// arguing them disjoint -- the attribute side already pools two allocators, the
+// uniquer and the DistinctAttr allocator, on that same argument.
 //
 //===----------------------------------------------------------------------===//
 
@@ -23,106 +23,102 @@
 #define MLIR_LIB_IR_SYMBOLREFCONTAINMENTCACHE_H
 
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Types.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/RWMutex.h"
+#include "llvm/Support/Threading.h"
 #include <optional>
 
 namespace mlir {
 class MLIRContext;
 namespace detail {
 
-/// Per-context store of the "provably free of a transitive SymbolRefAttr" fact
-/// for uniqued types and attributes. Two sets mirror the context's two
-/// uniquers, so type and attribute opaque pointers never need to be argued
-/// disjoint.
-///
-/// Each set holds only the opaque pointers of objects proven reference-free.
-/// Membership means "clear" (answer false); non-membership means "not yet
-/// proven clear", which the fill treats as unfilled and may-contain both at
-/// once -- it recomputes containment, and a genuine may-contain object is
-/// recomputed cheaply because the walk early-exits at its first SymbolRefAttr.
-/// A may-contain answer is therefore never stored, so a mutable-storage kind is
-/// conservatively may-contain forever with no special handling.
-///
-/// Locking discipline: the sets live entirely behind one SmartRWMutex. A lookup
-/// holds the read lock for its whole probe; an insert -- including any set
-/// growth it triggers -- holds the write lock. Readers therefore never observe
-/// a half-grown set and the plain sets need no per-slot atomics. The `lock`
-/// flag threaded through each operation is the context's runtime multithreading
-/// flag: when it is false the store is touched single-threaded and no lock is
-/// taken, mirroring MLIRContext's ScopedWriterLock.
+/// Per-context cache of the "provably free of a transitive SymbolRefAttr" fact
+/// for uniqued types and attributes, computed on demand, memoizing proven-clear
+/// answers. A false answer is authoritative -- the object references no symbol
+/// -- so a conforming SymbolUserTypeInterface has a vacuous verifySymbolUses
+/// and need never be visited. A true answer is conservative: a mutable-storage
+/// kind, whose sub-elements may change after the query, and anything containing
+/// one, always answers true.
 class SymbolRefContainmentCache {
 public:
-  SymbolRefContainmentCache() = default;
+  /// `isMultithreaded` is the context's runtime multithreading flag, read live
+  /// on every probe to decide whether the set needs its lock; it must outlive
+  /// this cache.
+  explicit SymbolRefContainmentCache(const bool &isMultithreaded)
+      : isMultithreaded(isMultithreaded) {}
   SymbolRefContainmentCache(const SymbolRefContainmentCache &) = delete;
   SymbolRefContainmentCache &
   operator=(const SymbolRefContainmentCache &) = delete;
 
-  /// Return false for `type`/`attr` if it is recorded clear, or nullopt if it
-  /// is not yet proven clear (unfilled, treated as may-contain by the caller).
-  /// A true answer is never recorded, so lookup never returns true.
-  std::optional<bool> lookup(Type type, bool lock) const {
-    return lookupImpl(typeEntries, type.getAsOpaquePointer(), lock);
-  }
-  std::optional<bool> lookup(Attribute attr, bool lock) const {
-    return lookupImpl(attrEntries, attr.getAsOpaquePointer(), lock);
-  }
-
-  /// Record `type`/`attr` as clear when `value` is false and return the answer
-  /// unchanged. A may-contain (`value` true) fact is not stored -- it is
-  /// recomputed cheaply on each encounter -- so insert then returns true
-  /// without touching the set. A racing duplicate clear fill records the same
-  /// membership and is a no-op.
-  bool insert(Type type, bool value, bool lock) {
-    return insertImpl(typeEntries, type.getAsOpaquePointer(), value, lock);
-  }
-  bool insert(Attribute attr, bool value, bool lock) {
-    return insertImpl(attrEntries, attr.getAsOpaquePointer(), value, lock);
-  }
+  bool mayContainSymbolRefs(Type type) { return compute(type); }
+  bool mayContainSymbolRefs(Attribute attr) { return compute(attr); }
 
 private:
-  /// One uniquer's clear-object set. It is not self-synchronizing: the
-  /// enclosing cache serializes all access with its SmartRWMutex, so every
-  /// method here runs under the read lock (lookup) or the write lock (insert
-  /// and the growth it may trigger).
-  struct ClearSet {
-    llvm::DenseSet<const void *> clear;
-
-    std::optional<bool> lookup(const void *key) const {
-      if (clear.count(key))
-        return false;
-      return std::nullopt;
-    }
+  // A type is never itself a SymbolRefAttr; an attribute is one exactly when
+  // isa<SymbolRefAttr> holds (covering FlatSymbolRefAttr).
+  static bool isSelfSymbolRef(Type) { return false; }
+  static bool isSelfSymbolRef(Attribute attr) {
+    return isa<SymbolRefAttr>(attr);
+  }
 
-    bool insert(const void *key, bool value) {
-      // Only the clear fact is durable; a may-contain answer is left out to be
-      // recomputed cheaply on each encounter. A racing duplicate clear fill
-      // records the same membership and is a no-op.
-      if (!value)
-        clear.insert(key);
-      return value;
+  /// Return whether `obj` may transitively contain a SymbolRefAttr, recording a
+  /// proven-clear result so it need never recompute. The recursion needs no
+  /// in-progress guard: immutable objects form a DAG (sub-elements are interned
+  /// before their parents), so every cycle passes through a mutable kind, where
+  /// the descent stops above before reading contents. Once one sub-element
+  /// answers may-contain the rest need not be visited: walkImmediateSubElements
+  /// cannot be interrupted, but the callback is a no-op once `mayContain`
+  /// holds.
+  template <typename T>
+  bool compute(T obj) {
+    const void *key = obj.getAsOpaquePointer();
+    if (isKnownClear(key))
+      return false;
+    bool mayContain = isSelfSymbolRef(obj) ||
+                      obj.template hasTrait<StorageUserTrait::IsMutable>();
+    if (!mayContain) {
+      auto walkSub = [&](auto sub) {
+        if (!mayContain && sub)
+          mayContain = compute(sub);
+      };
+      obj.walkImmediateSubElements(walkSub, walkSub);
     }
-  };
+    if (!mayContain)
+      markClear(key);
+    return mayContain;
+  }
 
-  std::optional<bool> lookupImpl(const ClearSet &set, const void *key,
-                                 bool lock) const {
+  bool isKnownClear(const void *key) const {
     std::optional<llvm::sys::SmartScopedReader<true>> guard;
-    if (lock)
+    if (shouldLock())
       guard.emplace(mutex);
-    return set.lookup(key);
+    return clear.contains(key);
   }
-
-  bool insertImpl(ClearSet &set, const void *key, bool value, bool lock) {
+  void markClear(const void *key) {
     std::optional<llvm::sys::SmartScopedWriter<true>> guard;
-    if (lock)
+    if (shouldLock())
       guard.emplace(mutex);
-    return set.insert(key, value);
+    clear.insert(key);
+  }
+
+  /// The set is serialized only while the context runs multithreaded;
+  /// llvm_is_multithreaded() is a compile-time constant folded in here.
+  bool shouldLock() const {
+    return isMultithreaded && llvm::llvm_is_multithreaded();
   }
 
+  const bool &isMultithreaded;
+  // The set lives behind this lock: a read lock guards the contains probe, a
+  // write lock the insert and any growth it triggers, and neither is ever held
+  // across the recursion.
   mutable llvm::sys::SmartRWMutex<true> mutex;
-  ClearSet typeEntries;
-  ClearSet attrEntries;
+  // Only the proven-clear opaque pointers are recorded; a may-contain object --
+  // including every mutable-storage kind, whose contents are never read -- is
+  // left out and recomputed on each encounter, so no fact that could later go
+  // stale is ever stored.
+  llvm::DenseSet<const void *> clear;
 };
 
 /// Return the given context's symbol-reference containment cache. Defined in
diff --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index ab97ede4465eb..2609e6168ef7c 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -473,79 +473,6 @@ raw_ostream &mlir::operator<<(raw_ostream &os,
   llvm_unreachable("Unexpected visibility");
 }
 
-//===----------------------------------------------------------------------===//
-// SymbolRefAttr containment query
-//===----------------------------------------------------------------------===//
-
-static bool
-computeMayContainSymbolRefs(Type type, bool lock,
-                            detail::SymbolRefContainmentCache &cache);
-static bool
-computeMayContainSymbolRefs(Attribute attr, bool lock,
-                            detail::SymbolRefContainmentCache &cache);
-
-/// Fill (once) and return the "may transitively contain a SymbolRefAttr" fact
-/// for a uniqued type or attribute `obj`, seeded by `selfIsRef` for an object
-/// that is itself a symbol reference. The context's runtime multithreading flag
-/// `lock` and its containment `cache` are resolved once per verification scope
-/// and threaded through the recursion, so no probe re-derives them. The answer
-/// is computed entirely outside the cache lock -- which is taken only for the
-/// write-once insert -- so no lock is held across the recursion and a racing
-/// duplicate fill computes the same fact idempotently. A mutable-storage kind
-/// may gain sub-elements after this point, so its contents are never read; it
-/// and its containers report may-contain conservatively.
-template <typename T>
-static bool fillMayContainSymbolRefs(T obj, bool selfIsRef, bool lock,
-                                     detail::SymbolRefContainmentCache &cache) {
-  if (std::optional<bool> cached = cache.lookup(obj, lock))
-    return *cached;
-
-  bool mayContain =
-      selfIsRef || obj.template hasTrait<detail::StorageUserTrait::IsMutable>();
-  // The recursion needs no in-progress guard: immutable objects form a DAG
-  // (sub-elements are interned before their parents), so every cycle passes
-  // through a mutable kind, where the fill stops above before descending. Once
-  // one sub-element has answered may-contain, the rest need not be visited:
-  // walkImmediateSubElements cannot be interrupted, but the callback is a no-op
-  // once `mayContain` is set.
-  if (!mayContain)
-    obj.walkImmediateSubElements(
-        [&](Attribute sub) {
-          if (!mayContain && sub)
-            mayContain = computeMayContainSymbolRefs(sub, lock, cache);
-        },
-        [&](Type sub) {
-          if (!mayContain && sub)
-            mayContain = computeMayContainSymbolRefs(sub, lock, cache);
-        });
-  return cache.insert(obj, mayContain, lock);
-}
-
-// A type is never itself a SymbolRefAttr; an attribute is one exactly when
-// isa<SymbolRefAttr> holds (covering FlatSymbolRefAttr).
-static bool
-computeMayContainSymbolRefs(Type type, bool lock,
-                            detail::SymbolRefContainmentCache &cache) {
-  return fillMayContainSymbolRefs(type, /*selfIsRef=*/false, lock, cache);
-}
-static bool
-computeMayContainSymbolRefs(Attribute attr, bool lock,
-                            detail::SymbolRefContainmentCache &cache) {
-  return fillMayContainSymbolRefs(attr, /*selfIsRef=*/isa<SymbolRefAttr>(attr),
-                                  lock, cache);
-}
-
-bool SymbolTable::mayContainSymbolRefs(Type type) {
-  MLIRContext *ctx = type.getContext();
-  return computeMayContainSymbolRefs(type, ctx->isMultithreadingEnabled(),
-                                     detail::getSymbolRefContainmentCache(ctx));
-}
-bool SymbolTable::mayContainSymbolRefs(Attribute attr) {
-  MLIRContext *ctx = attr.getContext();
-  return computeMayContainSymbolRefs(attr, ctx->isMultithreadingEnabled(),
-                                     detail::getSymbolRefContainmentCache(ctx));
-}
-
 //===----------------------------------------------------------------------===//
 // SymbolTable Trait Types
 //===----------------------------------------------------------------------===//
@@ -554,23 +481,24 @@ namespace {
 /// Verifies the symbol uses carried by the types an operation owns -- its
 /// operand, result, and block-argument types, and the types nested within its
 /// attributes -- against the enclosing symbol table, failing fast on the first
-/// invalid use. One instance serves a whole verification scope and is invoked
-/// once per operation, holding the per-scope state each verification would
-/// otherwise re-derive: the context's containment cache and multithreading
-/// flag, a NamedAttrList reused as the scratch buffer for properties-held
+/// invalid use. A root is a type or attribute the operation holds directly: its
+/// operand, result, and block-argument types, and its top-level attributes (the
+/// raw stored dictionary and each properties-held inherent attribute); its
+/// sub-elements are everything beneath. One instance serves a whole
+/// verification scope and is invoked once per operation, holding the per-scope
+/// state each verification would otherwise re-derive: the context's containment
+/// cache, a NamedAttrList reused as the scratch buffer for properties-held
 /// attributes, and a single AttrTypeWalker whose visited memo walks each
 /// uniqued type or attribute -- which may recur across many positions and
 /// operations -- at most once per scope.
 class OpTypeSymbolUseVerifier {
 public:
   OpTypeSymbolUseVerifier(MLIRContext *ctx, SymbolTableCollection &symbolTable)
-      : lock(ctx->isMultithreadingEnabled()),
-        cache(detail::getSymbolRefContainmentCache(ctx)) {
+      : cache(detail::getSymbolRefContainmentCache(ctx)) {
     typeWalker.addWalk([this, &symbolTable](Type type) -> WalkResult {
-      // Prune subtrees that provably hold no SymbolRefAttr: a conforming
-      // SymbolUserTypeInterface spells its references as SymbolRefAttr
-      // sub-elements, so such a type has a vacuous verifySymbolUses.
-      if (!computeMayContainSymbolRefs(type, lock, cache))
+      // Prune subtrees that provably hold no SymbolRefAttr; nothing under them
+      // can carry a symbol use.
+      if (!cache.mayContainSymbolRefs(type))
         return WalkResult::skip();
       if (auto user = dyn_cast<SymbolUserTypeInterface>(type))
         if (failed(user.verifySymbolUses(anchor, symbolTable)))
@@ -580,7 +508,7 @@ class OpTypeSymbolUseVerifier {
     typeWalker.addWalk([this](Attribute attr) -> WalkResult {
       // Prune reference-free attribute subtrees so the walk descends only where
       // a SymbolRefAttr, and thus a symbol-using type, may live.
-      if (!computeMayContainSymbolRefs(attr, lock, cache))
+      if (!cache.mayContainSymbolRefs(attr))
         return WalkResult::skip();
       return WalkResult::advance();
     });
@@ -590,64 +518,59 @@ class OpTypeSymbolUseVerifier {
   /// SymbolUserTypeInterface check at `op`.
   LogicalResult verify(Operation *op) {
     anchor = op;
-    // Skip a root that provably holds no SymbolRefAttr before entering the
-    // walker at all: the cheap containment probe spares the walker's
-    // out-of-line descent and keeps its visited memo confined to subtrees that
-    // can actually carry a symbol use, rather than growing it by every distinct
-    // clean root. The walker's callbacks prune interior subtrees on the same
-    // fact.
-    auto verifyType = [&](Type type) {
-      if (!computeMayContainSymbolRefs(type, lock, cache))
-        return WalkResult::advance();
-      return typeWalker.walk<WalkOrder::PreOrder>(type);
-    };
-    auto verifyAttr = [&](Attribute attr) {
-      if (!attr || !computeMayContainSymbolRefs(attr, lock, cache))
-        return WalkResult::advance();
-      return typeWalker.walk<WalkOrder::PreOrder>(attr);
-    };
-
     for (Type type : op->getOperandTypes())
-      if (verifyType(type).wasInterrupted())
+      if (visit(type).wasInterrupted())
         return failure();
     for (Type type : op->getResultTypes())
-      if (verifyType(type).wasInterrupted())
+      if (visit(type).wasInterrupted())
         return failure();
     for (Region &region : op->getRegions())
       for (Block &block : region)
         for (BlockArgument argument : block.getArguments())
-          if (verifyType(argument.getType()).wasInterrupted())
+          if (visit(argument.getType()).wasInterrupted())
             return failure();
 
-    return success(!forEachAttributeRoot(op, verifyAttr).wasInterrupted());
+    return success(!verifyAttributeRoots(op).wasInterrupted());
   }
 
 private:
-  /// Invoke `root` on each attribute subtree an operation can carry a nested
-  /// type within: its raw stored attribute dictionary as one root, then each
-  /// properties-held inherent attribute, stopping early once `root` interrupts.
-  /// This covers exactly what getAttrDictionary() would walk without allocating
-  /// and uniquing a fresh DictionaryAttr per operation -- the raw dictionary
-  /// already holds the inherent attributes of operations that keep none in
-  /// properties and the discardable attributes otherwise, and
-  /// populateInherentAttrs supplies the properties-held remainder into the
-  /// scope-reused list, cleared before each population so its buffer serves
-  /// every operation without accumulating attributes across them.
-  WalkResult forEachAttributeRoot(Operation *op,
-                                  function_ref<WalkResult(Attribute)> root) {
-    if (root(op->getRawDictionaryAttrs()).wasInterrupted())
+  /// Verify one root, probing its containment before entering the walker. A
+  /// clean top-level root is skipped outright -- kept out of the walk and out
+  /// of the walker's visited memo, unlike the clean interior elements the walk
+  /// memoizes as it descends. A may-contain root is walked and reprobed by the
+  /// walker's root callback, so this spares only the clean roots, which
+  /// dominate.
+  template <typename T>
+  WalkResult visit(T root) {
+    if (!cache.mayContainSymbolRefs(root))
+      return WalkResult::advance();
+    return typeWalker.walk<WalkOrder::PreOrder>(root);
+  }
+
+  /// Verify each attribute root an operation can carry a nested type within:
+  /// its raw stored attribute dictionary as one root, then each properties-held
+  /// inherent attribute, stopping early on the first invalid use. This reaches
+  /// everything getAttrDictionary() would without allocating and uniquing a
+  /// fresh DictionaryAttr per operation -- the raw dictionary already holds the
+  /// inherent attributes of operations that keep none in properties and the
+  /// discardable attributes otherwise, and populateInherentAttrs supplies the
+  /// properties-held remainder into the scope-reused list, cleared before each
+  /// population so its buffer serves every operation without accumulating
+  /// attributes across them. An operation's raw dictionary and a NamedAttribute
+  /// value are both non-null by construction, so no root is ever null.
+  WalkResult verifyAttributeRoots(Operation *op) {
+    if (visit(op->getRawDictionaryAttrs()).wasInterrupted())
       return WalkResult::interrupt();
     if (op->getPropertiesStorageSize()) {
       inherentAttrs.clear();
       op->getName().populateInherentAttrs(op, inherentAttrs);
       for (const NamedAttribute &namedAttr : inherentAttrs)
-        if (root(namedAttr.getValue()).wasInterrupted())
+        if (visit(namedAttr.getValue()).wasInterrupted())
           return WalkResult::interrupt();
     }
     return WalkResult::advance();
   }
 
-  bool lock;
   detail::SymbolRefContainmentCache &cache;
   AttrTypeWalker typeWalker;
   NamedAttrList inherentAttrs;
diff --git a/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp b/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
index c0dc2aa25434e..c7b3fe3af5242 100644
--- a/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
+++ b/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
@@ -6,13 +6,14 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// Tests for SymbolTable::mayContainSymbolRefs, which records per uniqued type
-// or attribute whether it transitively contains a SymbolRefAttr (conservatively
-// true for mutable storage). Symbol-table verification relies on it to prune
-// types and attributes that provably hold no symbol reference. Because a
-// SymbolUserTypeInterface / SymbolUserAttrInterface implementation must spell
-// its references as SymbolRefAttr sub-elements, a false answer is a sound
-// reason to skip an instance even after the interface is attached late.
+// Tests for the per-context symbol-reference containment cache, which records
+// which uniqued types and attributes are provably free of a transitive
+// SymbolRefAttr; everything else, including mutable storage, answers true
+// conservatively. Symbol-table verification relies on it to prune the
+// symbol-using types it would otherwise walk. Because a SymbolUserTypeInterface
+// implementation must spell its references as SymbolRefAttr sub-elements, a
+// false answer is a sound reason to skip such a type even after the interface
+// is attached late.
 //
 //===----------------------------------------------------------------------===//
 
@@ -32,6 +33,7 @@
 #include "../../test/lib/Dialect/Test/TestTypes.h"
 
 #include <atomic>
+#include <string>
 #include <thread>
 #include <vector>
 
@@ -39,24 +41,15 @@ using namespace mlir;
 
 namespace {
 
-// Symbol-user models whose verification always fails, attached externally to
-// exercise late interface attachment. One targets a type that structurally
-// holds a SymbolRefAttr (a tensor with a symbol-ref encoding); the other
-// targets f32, which holds none.
-struct FailingTensorSymbolUserModel
+// A symbol-user model whose verification always fails, attached externally to
+// exercise late interface attachment on an arbitrary concrete type.
+template <typename ConcreteType>
+struct FailingSymbolUserModel
     : public SymbolUserTypeInterface::ExternalModel<
-          FailingTensorSymbolUserModel, RankedTensorType> {
+          FailingSymbolUserModel<ConcreteType>, ConcreteType> {
   LogicalResult verifySymbolUses(Type type, Operation *op,
                                  SymbolTableCollection &symbolTable) const {
-    return op->emitError("tensor rejected by its attached symbol-user model");
-  }
-};
-struct FailingF32SymbolUserModel
-    : public SymbolUserTypeInterface::ExternalModel<FailingF32SymbolUserModel,
-                                                    Float32Type> {
-  LogicalResult verifySymbolUses(Type type, Operation *op,
-                                 SymbolTableCollection &symbolTable) const {
-    return op->emitError("f32 rejected by its attached symbol-user model");
+    return op->emitError("rejected by its attached symbol-user model");
   }
 };
 
@@ -83,119 +76,87 @@ class SymbolReferenceContainmentTest : public ::testing::Test {
     return test::TestSymbolRefAttr::get(&context, symbolRef());
   }
 
+  // Query the cache for `obj`. The row index names a failed row without
+  // printing `obj`, which a mutable-storage kind with an unset body cannot
+  // survive.
+  template <typename T>
+  void expect(T obj, bool expected) {
+    SCOPED_TRACE("containment row " + std::to_string(++row));
+    EXPECT_EQ(
+        detail::getSymbolRefContainmentCache(&context).mayContainSymbolRefs(
+            obj),
+        expected);
+  }
+
+  unsigned row = 0;
   MLIRContext context;
 };
 
-// A leaf type holding no symbol reference answers false.
-TEST_F(SymbolReferenceContainmentTest, LeafTypeIsClear) {
-  EXPECT_FALSE(
-      SymbolTable::mayContainSymbolRefs(IntegerType::get(&context, 32)));
-}
-
-// A plain attribute holding no symbol reference answers false.
-TEST_F(SymbolReferenceContainmentTest, LeafAttrIsClear) {
-  EXPECT_FALSE(
-      SymbolTable::mayContainSymbolRefs(StringAttr::get(&context, "hi")));
-  EXPECT_FALSE(SymbolTable::mayContainSymbolRefs(
-      TypeAttr::get(IntegerType::get(&context, 32))));
-}
-
-// A SymbolRefAttr itself answers true.
-TEST_F(SymbolReferenceContainmentTest, FlatSymbolRefAttrIsTrue) {
-  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(symbolRef()));
-}
-
-// A non-flat SymbolRefAttr, which nests further references, answers true.
-TEST_F(SymbolReferenceContainmentTest, NestedSymbolRefAttrIsTrue) {
-  SymbolRefAttr ref =
-      SymbolRefAttr::get(StringAttr::get(&context, "root"),
-                         {FlatSymbolRefAttr::get(&context, "n")});
-  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(ref));
-}
-
-// A conforming symbol-user type answers true through its SymbolRefAttr
-// parameter (not through the interface, which plays no part in the answer).
-TEST_F(SymbolReferenceContainmentTest, ConformingSymbolUserTypeIsTrue) {
-  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(symbolUserType()));
-}
-
-// A conforming symbol-user attribute answers true through its SymbolRefAttr
-// parameter.
-TEST_F(SymbolReferenceContainmentTest, ConformingSymbolUserAttrIsTrue) {
-  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(symbolUserAttr()));
-}
-
-// A type nesting a symbol-ref-bearing type propagates true.
-TEST_F(SymbolReferenceContainmentTest, TypeNestingSymbolRefBearingType) {
-  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(
-      TupleType::get(&context, {symbolUserType()})));
-}
-
-// A tuple of ordinary types stays false.
-TEST_F(SymbolReferenceContainmentTest, TypeNestingOrdinaryTypesIsClear) {
+// The containment answer across leaves, self-references, conforming user
+// types/attributes, and every nesting path a symbol reference can hide behind.
+TEST_F(SymbolReferenceContainmentTest, ContainmentTruthTable) {
   Type i32 = IntegerType::get(&context, 32);
-  EXPECT_FALSE(
-      SymbolTable::mayContainSymbolRefs(TupleType::get(&context, {i32, i32})));
-}
-
-// A type reaches a SymbolRefAttr two levels deep, through an attribute
-// sub-element (a tensor encoding holding a TypeAttr of a symbol-ref type).
-TEST_F(SymbolReferenceContainmentTest, TypeReachesSymbolRefThroughAttribute) {
-  Attribute encoding = TypeAttr::get(symbolUserType());
-  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(encoding));
-  RankedTensorType tensor =
-      RankedTensorType::get({2}, IntegerType::get(&context, 32), encoding);
-  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(tensor));
-}
-
-// A type reaches a plain SymbolRefAttr through an attribute parameter (a tensor
-// encoding).
-TEST_F(SymbolReferenceContainmentTest, TypeWithSymbolRefAttrParameter) {
-  RankedTensorType tensor =
-      RankedTensorType::get({2}, IntegerType::get(&context, 32), symbolRef());
-  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(tensor));
-}
-
-// A dictionary attribute containing a plain SymbolRefAttr answers true.
-TEST_F(SymbolReferenceContainmentTest, DictionaryAttrContainingSymbolRef) {
-  NamedAttribute named(StringAttr::get(&context, "callee"), symbolRef());
-  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(
-      DictionaryAttr::get(&context, {named})));
-}
-
-// A dictionary attribute containing a symbol-ref-bearing type inside a TypeAttr
-// answers true; the answer on the dictionary summarizes its whole nested tree.
-TEST_F(SymbolReferenceContainmentTest, DictionaryAttrContainingSymbolRefType) {
-  NamedAttribute named(StringAttr::get(&context, "key"),
-                       TypeAttr::get(symbolUserType()));
-  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(
-      DictionaryAttr::get(&context, {named})));
-}
 
-// A dictionary attribute with no symbol reference stays false.
-TEST_F(SymbolReferenceContainmentTest, DictionaryAttrIsClear) {
-  NamedAttribute named(StringAttr::get(&context, "key"),
-                       TypeAttr::get(IntegerType::get(&context, 32)));
-  EXPECT_FALSE(SymbolTable::mayContainSymbolRefs(
-      DictionaryAttr::get(&context, {named})));
-}
-
-// A type carrying a mutable component answers true conservatively, since its
-// sub-elements may change after the answer is computed at first query; the
-// fill never reads its contents.
-TEST_F(SymbolReferenceContainmentTest, MutableTypeReportsConservatively) {
-  test::TestRecursiveType recursive =
-      test::TestRecursiveType::get(&context, "rec");
-  EXPECT_TRUE(SymbolTable::mayContainSymbolRefs(recursive));
-}
-
-// A container of a mutable-storage kind inherits true, even before the mutable
-// body is populated, so no later mutation can turn a cached false stale.
-TEST_F(SymbolReferenceContainmentTest, ContainerOfMutableTypeIsTrue) {
-  test::TestRecursiveType recursive =
-      test::TestRecursiveType::get(&context, "rec2");
-  EXPECT_TRUE(
-      SymbolTable::mayContainSymbolRefs(TupleType::get(&context, {recursive})));
+  // Leaves hold no reference.
+  expect(i32, false);
+  expect(StringAttr::get(&context, "hi"), false);
+  expect(TypeAttr::get(i32), false);
+
+  // A SymbolRefAttr, flat or nesting further references, is itself a reference.
+  expect(symbolRef(), true);
+  expect(SymbolRefAttr::get(StringAttr::get(&context, "root"),
+                            {FlatSymbolRefAttr::get(&context, "n")}),
+         true);
+
+  // A conforming symbol-user type/attribute answers true through its
+  // SymbolRefAttr parameter, not through the interface (which plays no part).
+  expect(symbolUserType(), true);
+  expect(symbolUserAttr(), true);
+
+  // Nesting propagates the reference; ordinary nesting stays clear.
+  expect(TupleType::get(&context, {symbolUserType()}), true);
+  expect(TupleType::get(&context, {i32, i32}), false);
+
+  // A type reaches a reference through an attribute sub-element -- a plain
+  // SymbolRefAttr encoding, or a TypeAttr of a symbol-ref-bearing type.
+  expect(RankedTensorType::get({2}, i32, symbolRef()), true);
+  expect(TypeAttr::get(symbolUserType()), true);
+  expect(RankedTensorType::get({2}, i32, TypeAttr::get(symbolUserType())),
+         true);
+
+  // A dictionary summarizes its whole nested tree, whichever way a reference
+  // hides, and stays clear when none does.
+  expect(DictionaryAttr::get(
+             &context, {NamedAttribute(StringAttr::get(&context, "callee"),
+                                       symbolRef())}),
+         true);
+  expect(DictionaryAttr::get(&context,
+                             {NamedAttribute(StringAttr::get(&context, "key"),
+                                             TypeAttr::get(symbolUserType()))}),
+         true);
+  expect(DictionaryAttr::get(&context,
+                             {NamedAttribute(StringAttr::get(&context, "key"),
+                                             TypeAttr::get(i32))}),
+         false);
+
+  // A mutable-storage kind, and any container of one, answers true
+  // conservatively: its sub-elements may change after the query, so its
+  // contents are never read and no later mutation can turn a cached false
+  // stale.
+  expect(test::TestRecursiveType::get(&context, "rec"), true);
+  expect(
+      TupleType::get(&context, {test::TestRecursiveType::get(&context, "c")}),
+      true);
+
+  // A DistinctAttr is keyed by its own always-allocated storage address, not
+  // the attribute uniquer; two instances answer false independently, and a
+  // repeat query is stable across the cold fill and the warm hit.
+  DistinctAttr d1 = DistinctAttr::create(UnitAttr::get(&context));
+  DistinctAttr d2 = DistinctAttr::create(UnitAttr::get(&context));
+  ASSERT_NE(d1, d2);
+  expect(d1, false);
+  expect(d1, false);
+  expect(d2, false);
 }
 
 // Interface membership plays no part in the answer, so late attachment needs no
@@ -207,7 +168,8 @@ TEST_F(SymbolReferenceContainmentTest, LateInterfaceAttachmentStillVerifies) {
       "module { \"foo.op\"() : () -> tensor<4xf32, @sym> }", &context);
   ASSERT_TRUE(module);
 
-  RankedTensorType::attachInterface<FailingTensorSymbolUserModel>(context);
+  RankedTensorType::attachInterface<FailingSymbolUserModel<RankedTensorType>>(
+      context);
   ScopedDiagnosticHandler handler(&context,
                                   [](Diagnostic &) { return success(); });
   EXPECT_TRUE(failed(verify(*module)));
@@ -223,86 +185,68 @@ TEST_F(SymbolReferenceContainmentTest, NonConformingSymbolUserTypeIsSkipped) {
       "module { \"foo.op\"() : () -> f32 }", &context);
   ASSERT_TRUE(module);
 
-  Float32Type::attachInterface<FailingF32SymbolUserModel>(context);
+  Float32Type::attachInterface<FailingSymbolUserModel<Float32Type>>(context);
   ScopedDiagnosticHandler handler(&context,
                                   [](Diagnostic &) { return success(); });
   EXPECT_TRUE(succeeded(verify(*module)));
 }
 
+// The corpus both the ground-truth and the raced context are filled from,
+// spanning the interesting fill paths and sharing sub-elements so concurrent
+// fills overlap on the same interior objects. Built from one helper so the two
+// contexts stay in lock-step by construction.
+std::vector<Attribute> buildAttrs(MLIRContext &ctx) {
+  Type i32 = IntegerType::get(&ctx, 32);
+  FlatSymbolRefAttr sym = FlatSymbolRefAttr::get(&ctx, "sym");
+  return {
+      StringAttr::get(&ctx, "leaf"),
+      sym,
+      test::TestSymbolRefAttr::get(&ctx, sym),
+      TypeAttr::get(test::TestSymbolUserType::get(&ctx, sym)),
+      DictionaryAttr::get(
+          &ctx, {NamedAttribute(StringAttr::get(&ctx, "callee"), sym)}),
+      DictionaryAttr::get(&ctx, {NamedAttribute(StringAttr::get(&ctx, "plain"),
+                                                TypeAttr::get(i32))}),
+  };
+}
+std::vector<Type> buildTypes(MLIRContext &ctx) {
+  Type i32 = IntegerType::get(&ctx, 32);
+  FlatSymbolRefAttr sym = FlatSymbolRefAttr::get(&ctx, "sym");
+  test::TestSymbolUserType user = test::TestSymbolUserType::get(&ctx, sym);
+  return {
+      i32,
+      user,
+      TupleType::get(&ctx, {i32, user}),
+      TupleType::get(&ctx, {i32, i32}),
+      RankedTensorType::get({2}, i32, sym),
+  };
+}
+
 // Concurrency crux: many threads query the same cold, shared objects at once,
 // exactly as parallel symbol-table verification fills the context cache from
 // several isolated-op workers simultaneously. Every thread must agree with the
 // single-threaded ground truth, and the run must be clean under
-// ThreadSanitizer. Built cold (constructed but never queried before the
-// threads start) so the fills genuinely race.
+// ThreadSanitizer. The raced context is built cold (constructed but never
+// queried before the threads start) so the fills genuinely race.
 TEST_F(SymbolReferenceContainmentTest, ConcurrentFillIsRaceFree) {
   ASSERT_TRUE(context.isMultithreadingEnabled());
 
-  // A set of shared objects spanning the interesting fill paths and sharing
-  // sub-elements, so concurrent fills overlap on the same interior objects.
-  Type i32 = IntegerType::get(&context, 32);
-  std::vector<Attribute> attrs = {
-      StringAttr::get(&context, "leaf"),
-      symbolRef(),
-      symbolUserAttr(),
-      TypeAttr::get(symbolUserType()),
-      DictionaryAttr::get(
-          &context,
-          {NamedAttribute(StringAttr::get(&context, "callee"), symbolRef())}),
-      DictionaryAttr::get(&context,
-                          {NamedAttribute(StringAttr::get(&context, "plain"),
-                                          TypeAttr::get(i32))}),
-  };
-  std::vector<Type> types = {
-      i32,
-      symbolUserType(),
-      TupleType::get(&context, {i32, symbolUserType()}),
-      TupleType::get(&context, {i32, i32}),
-      RankedTensorType::get({2}, i32, symbolRef()),
-  };
-
-  // Ground truth computed single-threaded (this also warms nothing new for the
-  // threads, since these very objects are what they race to fill).
+  auto &truthCache = detail::getSymbolRefContainmentCache(&context);
+  std::vector<Attribute> attrs = buildAttrs(context);
+  std::vector<Type> types = buildTypes(context);
   std::vector<bool> attrTruth, typeTruth;
   for (Attribute a : attrs)
-    attrTruth.push_back(SymbolTable::mayContainSymbolRefs(a));
+    attrTruth.push_back(truthCache.mayContainSymbolRefs(a));
   for (Type t : types)
-    typeTruth.push_back(SymbolTable::mayContainSymbolRefs(t));
+    typeTruth.push_back(truthCache.mayContainSymbolRefs(t));
 
-  // Fresh context so the threads hit a cold cache and genuinely race the fills.
   MLIRContext raced;
   raced.loadDialect<test::TestDialect>();
   raced.allowUnregisteredDialects();
   ASSERT_TRUE(raced.isMultithreadingEnabled());
-  auto rebuildAttrs = [&](MLIRContext &ctx) {
-    Type ri32 = IntegerType::get(&ctx, 32);
-    FlatSymbolRefAttr rsym = FlatSymbolRefAttr::get(&ctx, "sym");
-    return std::vector<Attribute>{
-        StringAttr::get(&ctx, "leaf"),
-        rsym,
-        test::TestSymbolRefAttr::get(&ctx, rsym),
-        TypeAttr::get(test::TestSymbolUserType::get(&ctx, rsym)),
-        DictionaryAttr::get(
-            &ctx, {NamedAttribute(StringAttr::get(&ctx, "callee"), rsym)}),
-        DictionaryAttr::get(&ctx,
-                            {NamedAttribute(StringAttr::get(&ctx, "plain"),
-                                            TypeAttr::get(ri32))}),
-    };
-  };
-  auto rebuildTypes = [&](MLIRContext &ctx) {
-    Type ri32 = IntegerType::get(&ctx, 32);
-    FlatSymbolRefAttr rsym = FlatSymbolRefAttr::get(&ctx, "sym");
-    test::TestSymbolUserType user = test::TestSymbolUserType::get(&ctx, rsym);
-    return std::vector<Type>{
-        ri32,
-        user,
-        TupleType::get(&ctx, {ri32, user}),
-        TupleType::get(&ctx, {ri32, ri32}),
-        RankedTensorType::get({2}, ri32, rsym),
-    };
-  };
-  std::vector<Attribute> racedAttrs = rebuildAttrs(raced);
-  std::vector<Type> racedTypes = rebuildTypes(raced);
+  std::vector<Attribute> racedAttrs = buildAttrs(raced);
+  std::vector<Type> racedTypes = buildTypes(raced);
+  auto &racedCache = detail::getSymbolRefContainmentCache(&raced);
 
   const int numThreads = 16;
   std::vector<std::vector<bool>> attrResults(numThreads);
@@ -316,9 +260,9 @@ TEST_F(SymbolReferenceContainmentTest, ConcurrentFillIsRaceFree) {
       while (!go.load())
         ;
       for (Attribute a : racedAttrs)
-        attrResults[i].push_back(SymbolTable::mayContainSymbolRefs(a));
+        attrResults[i].push_back(racedCache.mayContainSymbolRefs(a));
       for (Type t : racedTypes)
-        typeResults[i].push_back(SymbolTable::mayContainSymbolRefs(t));
+        typeResults[i].push_back(racedCache.mayContainSymbolRefs(t));
     });
   while (ready.load() < numThreads)
     ;
@@ -334,77 +278,4 @@ TEST_F(SymbolReferenceContainmentTest, ConcurrentFillIsRaceFree) {
   }
 }
 
-// Growth preserves every recorded clear fact. The clear-object set starts empty
-// and grows as it fills, so inserting far past its initial capacity forces
-// several rehashes; every live key must survive each grow. Synthetic pointers
-// in a regular stride stand in for the bump-allocated storage addresses the
-// cache keys on; the keys are opaque and never dereferenced. Only clear facts
-// are recorded, so this checks that each survives the rehash and that
-// may-contain facts, which the store never keeps, stay misses.
-TEST_F(SymbolReferenceContainmentTest, SetGrowthPreservesEntries) {
-  detail::SymbolRefContainmentCache cache;
-  const int n = 5000; // many doublings past the minimal set
-  std::vector<Type> clearKeys;
-  std::vector<Type> mayContainKeys;
-  for (int i = 0; i < n; ++i) {
-    auto *p = reinterpret_cast<const void *>(
-        static_cast<uintptr_t>(0x100000 + i * 8));
-    Type key = Type::getFromOpaquePointer(p);
-    if (i % 3 == 0) {
-      // A may-contain fact is never stored: insert returns true but records
-      // nothing, so a later lookup stays a miss and the caller recomputes.
-      EXPECT_TRUE(cache.insert(key, /*value=*/true, /*lock=*/false));
-      mayContainKeys.push_back(key);
-    } else {
-      // A clear fact is recorded and returned unchanged.
-      EXPECT_FALSE(cache.insert(key, /*value=*/false, /*lock=*/false));
-      clearKeys.push_back(key);
-    }
-  }
-  // After all the growth, every clear fact is still found as false.
-  for (Type key : clearKeys) {
-    std::optional<bool> got = cache.lookup(key, /*lock=*/false);
-    ASSERT_TRUE(got.has_value()) << "lost clear entry";
-    EXPECT_FALSE(*got);
-  }
-  // A may-contain fact was never stored, so it stays a miss.
-  for (Type key : mayContainKeys)
-    EXPECT_FALSE(cache.lookup(key, /*lock=*/false).has_value());
-  // A key never inserted is a miss, not a stray cluster hit.
-  Type absent = Type::getFromOpaquePointer(
-      reinterpret_cast<const void *>(static_cast<uintptr_t>(0x100000 + n * 8)));
-  EXPECT_FALSE(cache.lookup(absent, /*lock=*/false).has_value());
-}
-
-// A DistinctAttr is keyed by the address of its own storage, which comes from a
-// separate always-allocating allocator rather than the attribute uniquer; the
-// cache must record and recover that pointer just like a uniqued one.
-TEST_F(SymbolReferenceContainmentTest, DistinctAttrIsHandled) {
-  DistinctAttr d1 = DistinctAttr::create(UnitAttr::get(&context));
-  DistinctAttr d2 = DistinctAttr::create(UnitAttr::get(&context));
-  ASSERT_NE(d1, d2); // always-allocating: distinct instances, distinct pointers
-
-  // Through the public query a DistinctAttr exposes no SymbolRefAttr
-  // sub-element, so it answers false, stably across the cold fill and the warm
-  // cache hit.
-  bool cold = SymbolTable::mayContainSymbolRefs(d1);
-  EXPECT_FALSE(cold);
-  EXPECT_EQ(SymbolTable::mayContainSymbolRefs(d1), cold);
-  EXPECT_FALSE(SymbolTable::mayContainSymbolRefs(d2));
-
-  // A may-contain answer is never recorded: a forced-true insert on a
-  // distinct-allocator pointer returns true but stores nothing, so the warm
-  // lookup stays a miss and the caller recomputes. A clear answer, by contrast,
-  // is recorded and returned false warm, while an uninserted distinct pointer
-  // stays a miss -- the pointer round-trips just like a uniqued one.
-  detail::SymbolRefContainmentCache cache;
-  EXPECT_TRUE(cache.insert(d1, /*value=*/true, /*lock=*/false));
-  EXPECT_FALSE(cache.lookup(d1, /*lock=*/false).has_value());
-  EXPECT_FALSE(cache.insert(d1, /*value=*/false, /*lock=*/false));
-  std::optional<bool> warm = cache.lookup(d1, /*lock=*/false);
-  ASSERT_TRUE(warm.has_value());
-  EXPECT_FALSE(*warm);
-  EXPECT_FALSE(cache.lookup(d2, /*lock=*/false).has_value());
-}
-
 } // namespace



More information about the Mlir-commits mailing list