diff --git a/mlir/docs/LangRef.md b/mlir/docs/LangRef.md index b3ee0c958c32..ff3af3e4e496 100644 --- a/mlir/docs/LangRef.md +++ b/mlir/docs/LangRef.md @@ -1334,9 +1334,10 @@ dense-elements-attribute ::= `dense` `<` attribute-value `>` `:` ``` A dense elements attribute is an elements attribute where the storage for the -constant vector or tensor value has been packed to the element bitwidth. The -element type of the vector or tensor constant must be of integer, index, or -floating point type. +constant vector or tensor value has been densely packed. The attribute supports +storing integer or floating point elements, with integer/index/floating element +types. It also support storing string elements with a custom dialect string +element type. ##### Opaque Elements Attribute diff --git a/mlir/include/mlir/IR/Attributes.h b/mlir/include/mlir/IR/Attributes.h index 43628b8ce506..6d24cd087648 100644 --- a/mlir/include/mlir/IR/Attributes.h +++ b/mlir/include/mlir/IR/Attributes.h @@ -703,7 +703,8 @@ public: /// Constructs a dense elements attribute from an array of element values. /// Each element attribute value is expected to be an element of 'type'. - /// 'type' must be a vector or tensor with static shape. + /// 'type' must be a vector or tensor with static shape. If the element of + /// `type` is non-integer/index/float it is assumed to be a string type. static DenseElementsAttr get(ShapedType type, ArrayRef values); /// Constructs a dense integer elements attribute from an array of integer diff --git a/mlir/lib/IR/Attributes.cpp b/mlir/lib/IR/Attributes.cpp index f51b273a9c89..f7540f6feffd 100644 --- a/mlir/lib/IR/Attributes.cpp +++ b/mlir/lib/IR/Attributes.cpp @@ -615,6 +615,8 @@ Attribute DenseElementsAttr::AttributeElementIterator::operator*() const { FloatElementIterator floatIt(floatEltTy.getFloatSemantics(), intIt); return FloatAttr::get(eltTy, *floatIt); } + if (owner.isa()) + return StringAttr::get(owner.getRawStringData()[index], eltTy); llvm_unreachable("unexpected element type"); } @@ -655,11 +657,23 @@ DenseElementsAttr::FloatElementIterator::FloatElementIterator( DenseElementsAttr DenseElementsAttr::get(ShapedType type, ArrayRef values) { - assert(type.getElementType().isIntOrIndexOrFloat() && - "expected int or index or float element type"); assert(hasSameElementsOrSplat(type, values)); + // If the element type is not based on int/float/index, assume it is a string + // type. auto eltType = type.getElementType(); + if (!type.getElementType().isIntOrIndexOrFloat()) { + SmallVector stringValues; + stringValues.reserve(values.size()); + for (Attribute attr : values) { + assert(attr.isa() && + "expected string value for non integer/index/float element"); + stringValues.push_back(attr.cast().getValue()); + } + return get(type, stringValues); + } + + // Otherwise, get the raw storage width to use for the allocation. size_t bitWidth = getDenseElementBitWidth(eltType); size_t storageBitWidth = getDenseElementStorageWidth(bitWidth); diff --git a/mlir/unittests/IR/AttributeTest.cpp b/mlir/unittests/IR/AttributeTest.cpp index ad4b422eae91..7d22b5e5a07f 100644 --- a/mlir/unittests/IR/AttributeTest.cpp +++ b/mlir/unittests/IR/AttributeTest.cpp @@ -154,4 +154,12 @@ TEST(DenseSplatTest, StringSplat) { testSplat(stringType, value); } +TEST(DenseSplatTest, StringAttrSplat) { + MLIRContext context; + Type stringType = + OpaqueType::get(Identifier::get("test", &context), "string", &context); + Attribute stringAttr = StringAttr::get("test-string", stringType); + testSplat(stringType, stringAttr); +} + } // end namespace