[mlir] Add a TypeAttr class, allow type attributes
PiperOrigin-RevId: 207235956
diff --git a/lib/IR/MLIRContext.cpp b/lib/IR/MLIRContext.cpp
index f558ae1..3edb115 100644
--- a/lib/IR/MLIRContext.cpp
+++ b/lib/IR/MLIRContext.cpp
@@ -249,6 +249,7 @@
using ArrayAttrSet = DenseSet<ArrayAttr *, ArrayAttrKeyInfo>;
ArrayAttrSet arrayAttrs;
DenseMap<AffineMap *, AffineMapAttr *> affineMapAttrs;
+ DenseMap<Type *, TypeAttr *> typeAttrs;
using AttributeListSet =
DenseSet<AttributeListStorage *, AttributeListKeyInfo>;
AttributeListSet attributeLists;
@@ -622,6 +623,16 @@
return result;
}
+TypeAttr *TypeAttr::get(Type *type, MLIRContext *context) {
+ auto *&result = context->getImpl().typeAttrs[type];
+ if (result)
+ return result;
+
+ result = context->getImpl().allocator.Allocate<TypeAttr>();
+ new (result) TypeAttr(type);
+ return result;
+}
+
/// Perform a three-way comparison between the names of the specified
/// NamedAttributes.
static int compareNamedAttributes(const NamedAttribute *lhs,