[mlir] Add a TypeAttr class, allow type attributes
PiperOrigin-RevId: 207235956
diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index 6177e19..b2ebaa6 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -280,6 +280,9 @@
case Attribute::Kind::AffineMap:
printAffineMapReference(cast<AffineMapAttr>(attr)->getValue());
break;
+ case Attribute::Kind::Type:
+ printType(cast<TypeAttr>(attr)->getValue());
+ break;
}
}
diff --git a/lib/IR/Builders.cpp b/lib/IR/Builders.cpp
index 199f35d..715e460 100644
--- a/lib/IR/Builders.cpp
+++ b/lib/IR/Builders.cpp
@@ -100,6 +100,10 @@
return AffineMapAttr::get(value, context);
}
+TypeAttr *Builder::getTypeAttr(Type *type) {
+ return TypeAttr::get(type, context);
+}
+
//===----------------------------------------------------------------------===//
// Affine Expressions and Affine Map.
//===----------------------------------------------------------------------===//
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,