[Mlir-commits] [mlir] [mlir] Make symbol-user type verification proportional to participating IR (PR #212354)
Jared Hoberock
llvmlistbot at llvm.org
Mon Jul 27 14:51:23 PDT 2026
https://github.com/jaredhoberock created https://github.com/llvm/llvm-project/pull/212354
Motivation: #198435 made the SymbolTable verifier walk every op's types and full attribute dictionary per pass; @ehein6 measured 14.1s -> 26.1s on a real pipeline with zero participating types (#212160).
On a synthetic reproduction (8000 functions, big dictionaries, 40 passes) this patch takes this pipeline from ~1.27s to ~0.64s. This PR composes with, and does not subsume, the registration-level skip proposed in #212160.
Whether or not an attribute or type contains a SymbolRefAttr is a fact known at uniquing time, so precompute and store this fact in the uniqued storage.
This PR adds a field to the storage of types & attributes and records, at uniquing time, whether the thing in question might contain a SymbolRefAttr somewhere in its tree. "Might", because implementations with mutable storage need to answer this question conservatively.
At SymbolTable verification time, we look at that field to avoid traversing types and attributes which do not contain any symbols.
Additionally:
1. One walker with one deduped visited set is hoisted once per SymbolTable scope so that the attributes & types which do get traversed get traversed at most once per symbol table scope, rather than once per operation they appear on.
2. Avoid materializing a DictionaryAttr value when checking an operation's attributes & properties.
One semantic change: the contract on SymbolUserTypeInterface & SymbolUserAttrInterface now requires that the symbols an implementation references be spelled as SymbolRefAttr sub-elements. So, a use is a structural fact about an attribute or type. As a consequence, an implementation that references symbols in some other encoding (e.g., a string) is treated as referencing no symbols and may no longer have verifySymbolUses invoked.
Other than tightening the contract on the two symbol user interfaces, the set of things verified against each symbol table scope is unchanged. This PR just removes all the pointless traversal.
Assisted-by: Claude Code (Anthropic)
>From a0ba771be0b2463c5f0b3e9e666a9b91f73c46cd Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Mon, 27 Jul 2026 13:51:26 -0500
Subject: [PATCH] [mlir] Make symbol-user type verification proportional to
participating IR
Motivation: #198435 made the SymbolTable verifier walk every op's types
and full attribute dictionary per pass; @ehein6 measured 14.1s -> 26.1s
on a real pipeline with zero participating types (#212160).
On a synthetic reproduction (8000 functions, big dictionaries, 40
passes) this patch takes this pipeline from ~1.27s to ~0.64s. This PR
composes with, and does not subsume, the registration-level skip
proposed in #212160.
Whether or not an attribute or type contains a SymbolRefAttr is a fact
known at uniquing time, so precompute and store this fact in the
uniqued storage.
This PR adds a field to the storage of types & attributes and records,
at uniquing time, whether the thing in question might contain a
SymbolRefAttr somewhere in its tree. "Might", because implementations
with mutable storage need to answer this question conservatively.
At SymbolTable verification time, we look at that field to avoid
traversing types and attributes which do not contain any symbols.
Additionally:
1. One walker with one deduped visited set is hoisted once per
SymbolTable scope so that the attributes & types which do get
traversed get traversed at most once per symbol table scope, rather
than once per operation they appear on.
2. Avoid materializing a DictionaryAttr value when checking an
operation's attributes & properties.
One semantic change: the contract on SymbolUserTypeInterface &
SymbolUserAttrInterface now requires that the symbols an implementation
references be spelled as SymbolRefAttr sub-elements. So, a use is a
structural fact about an attribute or type. As a consequence, an
implementation that references symbols in some other encoding (e.g., a
string) is treated as referencing no symbols and may no longer have
verifySymbolUses invoked.
Other than tightening the contract on the two symbol user interfaces,
the set of things verified against each symbol table scope is
unchanged. This PR just removes all the pointless traversal.
Assisted-by: Claude Code (Anthropic)
---
mlir/include/mlir/IR/AttributeSupport.h | 34 +++
mlir/include/mlir/IR/Attributes.h | 9 +
mlir/include/mlir/IR/SymbolInterfaces.td | 16 ++
mlir/include/mlir/IR/TypeSupport.h | 32 +++
mlir/include/mlir/IR/Types.h | 8 +
mlir/lib/IR/AttributeDetail.h | 4 +
mlir/lib/IR/Attributes.cpp | 28 +++
mlir/lib/IR/SymbolTable.cpp | 87 +++++---
mlir/lib/IR/Types.cpp | 28 +++
mlir/unittests/IR/AttributeTest.cpp | 24 +-
mlir/unittests/IR/CMakeLists.txt | 1 +
.../IR/SymbolReferenceContainmentTest.cpp | 211 ++++++++++++++++++
12 files changed, 440 insertions(+), 42 deletions(-)
create mode 100644 mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
diff --git a/mlir/include/mlir/IR/AttributeSupport.h b/mlir/include/mlir/IR/AttributeSupport.h
index 9f36ee4aae278..07aeb7cf47790 100644
--- a/mlir/include/mlir/IR/AttributeSupport.h
+++ b/mlir/include/mlir/IR/AttributeSupport.h
@@ -159,9 +159,21 @@ class AbstractAttribute {
// AttributeStorage
//===----------------------------------------------------------------------===//
+class AttributeStorage;
+
namespace detail {
class AttributeUniquer;
class DistinctAttributeUniquer;
+
+/// Compute, and record on `storage`, whether this attribute transitively
+/// contains a SymbolRefAttr, i.e. whether it is one or has a sub-element that
+/// is. Kinds carrying the mutable-storage trait, and anything containing them,
+/// report true conservatively. Called once per newly uniqued instance, after
+/// its abstract descriptor and sub-elements are set and before it is published,
+/// so that the bits of its immediate sub-elements (interned first) are final.
+/// Runs under the uniquer lock; walk implementations must not create new
+/// attributes or types.
+void populateAttrContainsSymbolReferences(AttributeStorage *storage);
} // namespace detail
/// Base storage class appearing in an attribute. Derived storage classes should
@@ -170,6 +182,7 @@ class alignas(8) AttributeStorage : public StorageUniquer::BaseStorage {
friend detail::AttributeUniquer;
friend detail::DistinctAttributeUniquer;
friend StorageUniquer;
+ friend void detail::populateAttrContainsSymbolReferences(AttributeStorage *);
public:
/// Return the abstract descriptor for this attribute.
@@ -178,6 +191,14 @@ class alignas(8) AttributeStorage : public StorageUniquer::BaseStorage {
return *abstractAttribute;
}
+ /// Return whether this attribute may transitively contain a SymbolRefAttr.
+ /// False is authoritative: no SymbolRefAttr is reachable through the
+ /// sub-element tree. True may be conservative: mutable-storage kinds, and
+ /// anything containing them, report true. A pure function of the attribute's
+ /// interned structure, synthesized at uniquing; not part of the unique key
+ /// and taking no part in hashing or equality.
+ bool mayContainSymbolRefs() const { return mayContainSymbolRefsBit; }
+
protected:
/// Set the abstract attribute for this storage instance. This is used by the
/// AttributeUniquer when initializing a newly constructed storage object.
@@ -190,8 +211,18 @@ class alignas(8) AttributeStorage : public StorageUniquer::BaseStorage {
void initialize(MLIRContext *context) {}
private:
+ /// Record the symbol-reference containment bit. Set once at uniquing.
+ void setMayContainSymbolRefs(bool mayContain) {
+ mayContainSymbolRefsBit = mayContain;
+ }
+
/// The abstract descriptor for this attribute.
const AbstractAttribute *abstractAttribute = nullptr;
+
+ /// Whether this attribute may transitively contain a SymbolRefAttr; a pure
+ /// function of the interned structure, taking no part in the unique key,
+ /// hashing, or equality.
+ bool mayContainSymbolRefsBit = false;
};
/// Default storage type for attributes that require no additional
@@ -243,6 +274,8 @@ class AttributeUniquer {
// Execute any additional attribute storage initialization with the
// context.
static_cast<typename T::ImplType *>(storage)->initialize(ctx);
+
+ populateAttrContainsSymbolReferences(storage);
},
typeID, std::forward<Args>(args)...);
}
@@ -299,6 +332,7 @@ class AttributeUniquer {
.registerSingletonStorageType<typename T::ImplType>(
typeID, [ctx, typeID](AttributeStorage *storage) {
initializeAttributeStorage(storage, ctx, typeID);
+ populateAttrContainsSymbolReferences(storage);
});
}
diff --git a/mlir/include/mlir/IR/Attributes.h b/mlir/include/mlir/IR/Attributes.h
index a2f9c5fcda873..9035a6c39597f 100644
--- a/mlir/include/mlir/IR/Attributes.h
+++ b/mlir/include/mlir/IR/Attributes.h
@@ -143,6 +143,15 @@ class Attribute {
/// Return the internal Attribute implementation.
ImplType *getImpl() const { return impl; }
+ /// Return whether this attribute may transitively contain a SymbolRefAttr.
+ /// False is authoritative: no SymbolRefAttr is reachable through the
+ /// sub-element tree. True may be conservative: mutable-storage kinds, and
+ /// anything containing them, report true. A pure function of the interned
+ /// structure; not part of the unique key, hashing, or equality. Symbol-table
+ /// verification uses it to skip attributes that cannot hold a symbol
+ /// reference.
+ bool mayContainSymbolRefs() const { return impl->mayContainSymbolRefs(); }
+
protected:
ImplType *impl{nullptr};
};
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/TypeSupport.h b/mlir/include/mlir/IR/TypeSupport.h
index 36ef696ba4f6f..81b4563b3ce2a 100644
--- a/mlir/include/mlir/IR/TypeSupport.h
+++ b/mlir/include/mlir/IR/TypeSupport.h
@@ -158,14 +158,26 @@ class AbstractType {
// TypeStorage
//===----------------------------------------------------------------------===//
+class TypeStorage;
+
namespace detail {
struct TypeUniquer;
+
+/// Compute, and record on `storage`, whether this type transitively contains a
+/// SymbolRefAttr. Kinds carrying the mutable-storage trait, and anything
+/// containing them, report true conservatively. Called once per newly uniqued
+/// instance, after its abstract descriptor and sub-elements are set and before
+/// it is published, so that the bits of its immediate sub-elements (interned
+/// first) are final. Runs under the uniquer lock; walk implementations must not
+/// create new attributes or types.
+void populateTypeContainsSymbolReferences(TypeStorage *storage);
} // namespace detail
/// Base storage class appearing in a Type.
class TypeStorage : public StorageUniquer::BaseStorage {
friend detail::TypeUniquer;
friend StorageUniquer;
+ friend void detail::populateTypeContainsSymbolReferences(TypeStorage *);
public:
/// Return the abstract type descriptor for this type.
@@ -174,6 +186,14 @@ class TypeStorage : public StorageUniquer::BaseStorage {
return *abstractType;
}
+ /// Return whether this type may transitively contain a SymbolRefAttr. False
+ /// is authoritative: no SymbolRefAttr is reachable through the sub-element
+ /// tree. True may be conservative: mutable-storage kinds, and anything
+ /// containing them, report true. A pure function of the type's interned
+ /// structure, synthesized at uniquing; not part of the unique key and taking
+ /// no part in hashing or equality.
+ bool mayContainSymbolRefs() const { return mayContainSymbolRefsBit; }
+
protected:
/// This constructor is used by derived classes as part of the TypeUniquer.
TypeStorage() {}
@@ -185,8 +205,18 @@ class TypeStorage : public StorageUniquer::BaseStorage {
abstractType = const_cast<AbstractType *>(&abstractTy);
}
+ /// Record the symbol-reference containment bit. Set once at uniquing.
+ void setMayContainSymbolRefs(bool mayContain) {
+ mayContainSymbolRefsBit = mayContain;
+ }
+
/// The abstract description for this type.
AbstractType *abstractType{nullptr};
+
+ /// Whether this type may transitively contain a SymbolRefAttr; a pure
+ /// function of the interned structure, taking no part in the unique key,
+ /// hashing, or equality.
+ bool mayContainSymbolRefsBit = false;
};
/// Default storage type for types that require no additional initialization or
@@ -233,6 +263,7 @@ struct TypeUniquer {
return ctx->getTypeUniquer().get<typename T::ImplType>(
[&, typeID](TypeStorage *storage) {
storage->initialize(AbstractType::lookup(typeID, ctx));
+ populateTypeContainsSymbolReferences(storage);
},
typeID, std::forward<Args>(args)...);
}
@@ -290,6 +321,7 @@ struct TypeUniquer {
ctx->getTypeUniquer().registerSingletonStorageType<TypeStorage>(
typeID, [&ctx, typeID](TypeStorage *storage) {
storage->initialize(AbstractType::lookup(typeID, ctx));
+ populateTypeContainsSymbolReferences(storage);
});
}
};
diff --git a/mlir/include/mlir/IR/Types.h b/mlir/include/mlir/IR/Types.h
index 70599b448f76e..879a2e7e1b281 100644
--- a/mlir/include/mlir/IR/Types.h
+++ b/mlir/include/mlir/IR/Types.h
@@ -192,6 +192,14 @@ class Type {
/// Return the Type implementation.
ImplType *getImpl() const { return impl; }
+ /// Return whether this type may transitively contain a SymbolRefAttr. False
+ /// is authoritative: no SymbolRefAttr is reachable through the sub-element
+ /// tree. True may be conservative: mutable-storage kinds, and anything
+ /// containing them, report true. A pure function of the interned structure;
+ /// not part of the unique key, hashing, or equality. Symbol-table
+ /// verification uses it to skip types that cannot hold a symbol reference.
+ bool mayContainSymbolRefs() const { return impl->mayContainSymbolRefs(); }
+
/// Walk all of the immediately nested sub-attributes and sub-types. This
/// method does not recurse into sub elements.
void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 4b7e7dd5677e8..63ef097003b7b 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -300,6 +300,10 @@ class DistinctAttributeUniquer {
context, std::forward<Args>(args)...);
storage->initializeAbstractAttribute(
AbstractAttribute::lookup(DistinctAttr::getTypeID(), context));
+ // Distinct attributes bypass the uniquer, so synthesize their
+ // "contains a SymbolRefAttr" bit here from the already-final bit of the
+ // referenced attribute.
+ populateAttrContainsSymbolReferences(storage);
return storage;
}
diff --git a/mlir/lib/IR/Attributes.cpp b/mlir/lib/IR/Attributes.cpp
index ff1cd8432fb07..c0123b1049639 100644
--- a/mlir/lib/IR/Attributes.cpp
+++ b/mlir/lib/IR/Attributes.cpp
@@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/Dialect.h"
using namespace mlir;
@@ -72,3 +73,30 @@ bool NamedAttribute::operator<(const NamedAttribute &rhs) const {
bool NamedAttribute::operator<(StringRef rhs) const {
return getName().getValue().compare(rhs) < 0;
}
+
+//===----------------------------------------------------------------------===//
+// Symbol-reference containment bit synthesis
+//===----------------------------------------------------------------------===//
+
+/// Synthesize the interning-time "contains a SymbolRefAttr" bit for a newly
+/// uniqued attribute. The bit is set when this attribute is a SymbolRefAttr
+/// (including FlatSymbolRefAttr), when it carries a mutable component, or when
+/// any immediate sub-element already carries the bit.
+void mlir::detail::populateAttrContainsSymbolReferences(
+ AttributeStorage *storage) {
+ Attribute attr(storage);
+ const AbstractAttribute &abstractAttr = storage->getAbstractAttribute();
+ bool mayContain = isa<SymbolRefAttr>(attr) ||
+ abstractAttr.hasTrait<StorageUserTrait::IsMutable>();
+ if (!mayContain) {
+ abstractAttr.walkImmediateSubElements(
+ attr,
+ [&](Attribute subAttr) {
+ mayContain |= subAttr && subAttr.mayContainSymbolRefs();
+ },
+ [&](Type subType) {
+ mayContain |= subType && subType.mayContainSymbolRefs();
+ });
+ }
+ storage->setMayContainSymbolRefs(mayContain);
+}
diff --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index 100eba7146d81..e170ba0a024aa 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -476,49 +476,57 @@ 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.
+/// `typeWalker` carries the SymbolUserTypeInterface check as a walk callback,
+/// anchored at `op` for symbol lookups, and its shared visited set makes each
+/// uniqued type, which may recur across many positions and operations, verified
+/// against the enclosing symbol table at most once. A type or attribute whose
+/// interning-time bit is clear is skipped: it provably contains no
+/// SymbolRefAttr, and a SymbolUserTypeInterface type is required to spell its
+/// references as SymbolRefAttr sub-elements, so walking it would verify
+/// nothing. 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();
+ AttrTypeWalker &typeWalker) {
+ auto verifyType = [&](Type type) {
+ if (!type.mayContainSymbolRefs())
return WalkResult::advance();
- });
+ return typeWalker.walk<WalkOrder::PreOrder>(type);
+ };
+ auto verifyAttr = [&](Attribute attr) {
+ if (!attr || !attr.mayContainSymbolRefs())
+ return WalkResult::advance();
+ return typeWalker.walk<WalkOrder::PreOrder>(attr);
};
for (Type type : op->getOperandTypes())
- if (verify(type).wasInterrupted())
+ if (verifyType(type).wasInterrupted())
return failure();
for (Type type : op->getResultTypes())
- if (verify(type).wasInterrupted())
+ if (verifyType(type).wasInterrupted())
return failure();
for (Region ®ion : op->getRegions())
for (Block &block : region)
for (BlockArgument argument : block.getArguments())
- if (verify(argument.getType()).wasInterrupted())
+ if (verifyType(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. Read the raw stored
+ // attribute dictionary rather than getAttrDictionary(): the latter allocates
+ // and uniques a fresh dictionary for every operation that keeps its inherent
+ // attributes in properties. The raw dictionary already covers inherent
+ // attributes for operations that do not use properties; the properties-held
+ // inherent attributes are walked separately below.
+ if (verifyAttr(op->getRawDictionaryAttrs()).wasInterrupted())
+ return failure();
+ if (op->getPropertiesStorageSize()) {
+ NamedAttrList inherentAttrs;
+ op->getName().populateInherentAttrs(op, inherentAttrs);
+ for (const NamedAttribute &namedAttr : inherentAttrs)
+ if (verifyAttr(namedAttr.getValue()).wasInterrupted())
+ return failure();
+ }
+ return success();
}
LogicalResult detail::verifySymbolTable(Operation *op) {
@@ -557,7 +565,21 @@ 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 set records each uniqued
+ // type once, so a type recurring across operand/result/block-argument and
+ // attribute positions is verified only at its first occurrence, whose
+ // operation supplies the lookup anchor.
+ 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 +592,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();
};
diff --git a/mlir/lib/IR/Types.cpp b/mlir/lib/IR/Types.cpp
index ec10a5ce9e2e7..f2abede851e13 100644
--- a/mlir/lib/IR/Types.cpp
+++ b/mlir/lib/IR/Types.cpp
@@ -6,6 +6,7 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dialect.h"
@@ -127,3 +128,30 @@ unsigned Type::getIntOrFloatBitWidth() const {
return intType.getWidth();
return llvm::cast<FloatType>(*this).getWidth();
}
+
+//===----------------------------------------------------------------------===//
+// Symbol-reference containment bit synthesis
+//===----------------------------------------------------------------------===//
+
+/// Synthesize the interning-time "contains a SymbolRefAttr" bit for a newly
+/// uniqued type. A type is never itself a SymbolRefAttr, so the bit is set when
+/// it carries a mutable component (whose sub-elements may change after this bit
+/// is fixed, so it must report conservatively) or when any immediate
+/// sub-element already carries the bit. Sub-elements are interned before their
+/// parents, so their bits are final.
+void mlir::detail::populateTypeContainsSymbolReferences(TypeStorage *storage) {
+ const AbstractType &abstractType = storage->getAbstractType();
+ bool mayContain = abstractType.hasTrait<StorageUserTrait::IsMutable>();
+ if (!mayContain) {
+ Type type(storage);
+ abstractType.walkImmediateSubElements(
+ type,
+ [&](Attribute attr) {
+ mayContain |= attr && attr.mayContainSymbolRefs();
+ },
+ [&](Type subType) {
+ mayContain |= subType && subType.mayContainSymbolRefs();
+ });
+ }
+ storage->setMayContainSymbolRefs(mayContain);
+}
diff --git a/mlir/unittests/IR/AttributeTest.cpp b/mlir/unittests/IR/AttributeTest.cpp
index 17d8e567fbebc..47892fe18de1b 100644
--- a/mlir/unittests/IR/AttributeTest.cpp
+++ b/mlir/unittests/IR/AttributeTest.cpp
@@ -475,15 +475,17 @@ TEST(CopyCountAttr, CopyCountGet) {
int counter1 = test::CopyCount::counter;
test::CopyCount::counter = 0;
test::TestCopyCountAttr::get(&context, std::move(copyCount));
+ // Synthesizing the "contains a symbol reference" bit for a newly uniqued
+ // attribute walks its immediate sub-elements once, which reconstructs the
+ // storage key and copies the payload.
#ifndef NDEBUG
- // One verification enabled only in assert-mode requires two copies: one for
- // calling 'verifyInvariants' and one for calling 'verify' inside
- // 'verifyInvariants'.
- EXPECT_EQ(counter1, 2);
- EXPECT_EQ(test::CopyCount::counter, 2);
+ // Two further copies come from the assert-mode verification: one for calling
+ // 'verifyInvariants' and one for calling 'verify' inside 'verifyInvariants'.
+ EXPECT_EQ(counter1, 3);
+ EXPECT_EQ(test::CopyCount::counter, 3);
#else
- EXPECT_EQ(counter1, 0);
- EXPECT_EQ(test::CopyCount::counter, 0);
+ EXPECT_EQ(counter1, 1);
+ EXPECT_EQ(test::CopyCount::counter, 1);
#endif
}
@@ -500,9 +502,11 @@ TEST(CopyCountAttr, CopyCountGetChecked) {
test::CopyCount::counter = 0;
test::TestCopyCountAttr::getChecked(loc, &context, std::move(copyCount));
// The verifiers require two copies: one for calling 'verifyInvariants' and
- // one for calling 'verify' inside 'verifyInvariants'.
- EXPECT_EQ(counter1, 2);
- EXPECT_EQ(test::CopyCount::counter, 2);
+ // one for calling 'verify' inside 'verifyInvariants'. Uniquing a new instance
+ // adds a third: synthesizing its "contains a symbol reference" bit walks its
+ // immediate sub-elements once, reconstructing the storage key.
+ EXPECT_EQ(counter1, 3);
+ EXPECT_EQ(test::CopyCount::counter, 3);
}
// Test stripped printing using test dialect attribute.
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..dda6347fecbe3
--- /dev/null
+++ b/mlir/unittests/IR/SymbolReferenceContainmentTest.cpp
@@ -0,0 +1,211 @@
+//===- SymbolReferenceContainmentTest.cpp - Containment bit 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 the interning-time bit that records whether a type or attribute
+// transitively contains a SymbolRefAttr (conservatively true for mutable
+// storage). Symbol-table verification and the symbol-use walks rely on it to
+// skip types and attributes that provably hold no symbol reference. Because a
+// SymbolUserTypeInterface / SymbolUserAttrInterface implementation must spell
+// its references as SymbolRefAttr sub-elements, the bit being clear 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"
+
+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 carries no bit.
+TEST_F(SymbolReferenceContainmentTest, LeafTypeIsClear) {
+ EXPECT_FALSE(IntegerType::get(&context, 32).mayContainSymbolRefs());
+}
+
+// A plain attribute holding no symbol reference carries no bit.
+TEST_F(SymbolReferenceContainmentTest, LeafAttrIsClear) {
+ EXPECT_FALSE(StringAttr::get(&context, "hi").mayContainSymbolRefs());
+ EXPECT_FALSE(
+ TypeAttr::get(IntegerType::get(&context, 32)).mayContainSymbolRefs());
+}
+
+// A SymbolRefAttr itself carries the bit.
+TEST_F(SymbolReferenceContainmentTest, FlatSymbolRefAttrHasBit) {
+ EXPECT_TRUE(symbolRef().mayContainSymbolRefs());
+}
+
+// A non-flat SymbolRefAttr, which nests further references, carries the bit.
+TEST_F(SymbolReferenceContainmentTest, NestedSymbolRefAttrHasBit) {
+ SymbolRefAttr ref =
+ SymbolRefAttr::get(StringAttr::get(&context, "root"),
+ {FlatSymbolRefAttr::get(&context, "n")});
+ EXPECT_TRUE(ref.mayContainSymbolRefs());
+}
+
+// A conforming symbol-user type carries the bit through its SymbolRefAttr
+// parameter (not through the interface, which plays no part in the bit).
+TEST_F(SymbolReferenceContainmentTest, ConformingSymbolUserTypeHasBit) {
+ EXPECT_TRUE(symbolUserType().mayContainSymbolRefs());
+}
+
+// A conforming symbol-user attribute carries the bit through its SymbolRefAttr
+// parameter.
+TEST_F(SymbolReferenceContainmentTest, ConformingSymbolUserAttrHasBit) {
+ EXPECT_TRUE(symbolUserAttr().mayContainSymbolRefs());
+}
+
+// A type nesting a symbol-ref-bearing type propagates the bit.
+TEST_F(SymbolReferenceContainmentTest, TypeNestingSymbolRefBearingType) {
+ EXPECT_TRUE(
+ TupleType::get(&context, {symbolUserType()}).mayContainSymbolRefs());
+}
+
+// A tuple of ordinary types stays clear.
+TEST_F(SymbolReferenceContainmentTest, TypeNestingOrdinaryTypesIsClear) {
+ Type i32 = IntegerType::get(&context, 32);
+ EXPECT_FALSE(TupleType::get(&context, {i32, i32}).mayContainSymbolRefs());
+}
+
+// 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(encoding.mayContainSymbolRefs());
+ RankedTensorType tensor =
+ RankedTensorType::get({2}, IntegerType::get(&context, 32), encoding);
+ EXPECT_TRUE(tensor.mayContainSymbolRefs());
+}
+
+// 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(tensor.mayContainSymbolRefs());
+}
+
+// A dictionary attribute containing a plain SymbolRefAttr carries the bit.
+TEST_F(SymbolReferenceContainmentTest, DictionaryAttrContainingSymbolRef) {
+ NamedAttribute named(StringAttr::get(&context, "callee"), symbolRef());
+ EXPECT_TRUE(DictionaryAttr::get(&context, {named}).mayContainSymbolRefs());
+}
+
+// A dictionary attribute containing a symbol-ref-bearing type inside a TypeAttr
+// carries the bit; the bit on the dictionary summarizes its whole nested tree.
+TEST_F(SymbolReferenceContainmentTest, DictionaryAttrContainingSymbolRefType) {
+ NamedAttribute named(StringAttr::get(&context, "key"),
+ TypeAttr::get(symbolUserType()));
+ EXPECT_TRUE(DictionaryAttr::get(&context, {named}).mayContainSymbolRefs());
+}
+
+// A dictionary attribute with no symbol reference stays clear.
+TEST_F(SymbolReferenceContainmentTest, DictionaryAttrIsClear) {
+ NamedAttribute named(StringAttr::get(&context, "key"),
+ TypeAttr::get(IntegerType::get(&context, 32)));
+ EXPECT_FALSE(DictionaryAttr::get(&context, {named}).mayContainSymbolRefs());
+}
+
+// A type carrying a mutable component reports the bit conservatively, since its
+// sub-elements may change after the bit is fixed at uniquing.
+TEST_F(SymbolReferenceContainmentTest, MutableTypeReportsConservatively) {
+ test::TestRecursiveType recursive =
+ test::TestRecursiveType::get(&context, "rec");
+ EXPECT_TRUE(recursive.mayContainSymbolRefs());
+}
+
+// Interface membership plays no part in the bit, so late attachment needs no
+// fallback: a type that structurally holds a SymbolRefAttr (a tensor with a
+// symbol-ref encoding) has its bit set 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)
+// has a clear bit 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)));
+}
+
+} // namespace
More information about the Mlir-commits
mailing list