[Mlir-commits] [mlir] [mlir][IR] Separate `DenseStringElementsAttr` from `DenseElementsAttr` (PR #181385)
Matthias Springer
llvmlistbot at llvm.org
Sun Feb 15 06:35:11 PST 2026
https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/181385
>From a108b9eeb062415edb20406ef7e7449dc736817e Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Sun, 1 Feb 2026 17:41:37 +0000
Subject: [PATCH 1/2] [mlir][WIP] `DenseElementsAttr` generalized
getter / iterator via interface
extraTraitClassDeclaration to provide default FloatType impls
address comments
simplify parser
---
mlir/include/mlir/IR/BuiltinAttributes.td | 49 ++++---
mlir/include/mlir/IR/BuiltinTypeInterfaces.h | 23 ++++
mlir/include/mlir/IR/BuiltinTypeInterfaces.td | 75 ++++++++++-
mlir/include/mlir/IR/BuiltinTypes.td | 14 +-
mlir/lib/AsmParser/AttributeParser.cpp | 125 +++++++++++++++++-
mlir/lib/IR/AsmPrinter.cpp | 44 +++++-
mlir/lib/IR/AttributeDetail.h | 10 +-
mlir/lib/IR/BuiltinAttributes.cpp | 117 ++++------------
mlir/lib/IR/BuiltinTypeInterfaces.cpp | 34 +++++
mlir/lib/IR/BuiltinTypes.cpp | 87 ++++++++++++
.../IR/dense-elements-type-interface.mlir | 83 ++++++++++++
mlir/test/lib/Dialect/Test/TestTypeDefs.td | 12 ++
mlir/test/lib/Dialect/Test/TestTypes.cpp | 28 ++++
mlir/test/lib/Dialect/Test/TestTypes.h | 1 +
14 files changed, 581 insertions(+), 121 deletions(-)
create mode 100644 mlir/test/IR/dense-elements-type-interface.mlir
diff --git a/mlir/include/mlir/IR/BuiltinAttributes.td b/mlir/include/mlir/IR/BuiltinAttributes.td
index 798d3c84f9618..dced379d1f979 100644
--- a/mlir/include/mlir/IR/BuiltinAttributes.td
+++ b/mlir/include/mlir/IR/BuiltinAttributes.td
@@ -239,29 +239,48 @@ def Builtin_DenseIntOrFPElementsAttr : Builtin_Attr<
"DenseElementsAttr"
> {
let summary = "An Attribute containing a dense multi-dimensional array of "
- "integer or floating-point values";
+ "values";
let description = [{
- Syntax:
-
- ```
- tensor-literal ::= integer-literal | float-literal | bool-literal | [] | [tensor-literal (, tensor-literal)* ]
- dense-intorfloat-elements-attribute ::= `dense` `<` tensor-literal `>` `:`
- ( tensor-type | vector-type )
- ```
-
- A dense int-or-float elements attribute is an elements attribute containing
- a densely packed vector or tensor of integer or floating-point values. The
- element type of this attribute is required to be either an `IntegerType` or
- a `FloatType`.
+ A dense elements attribute stores one or multiple elements of the same type.
+ The term "dense" refers to the fact that elements are not stored as
+ individual MLIR attributes, but in a raw buffer. The attribute provides a
+ covenience API to access elements in the form of MLIR attributes, but users
+ should avoid that API in performance-critical code and utilize APIs that
+ operate on raw bytes instead.
+
+ The number of elements is determined by the `type` shaped type. (Unranked
+ shaped types are not supported.) The element type of the shaped type must
+ implement the `DenseElementType` interface. This type interface defines the
+ bitwidth of an element and provides a serializer/deserializer to/from MLIR
+ attributes.
+
+ Storage format: Given an element bitwidth "w", element "i" starts at byte
+ offset "i * ceildiv(w, 8)". In other words, each element starts at a full
+ byte offset.
+
+ TODO: The name `DenseIntOrFPElements` is no longer accurate. The attribute
+ will be renamed in the future.
Examples:
```
- // A splat tensor of integer values.
+ // Literal-first syntax: A splat tensor of integer values.
dense<10> : tensor<2xi32>
- // A tensor of 2 float32 elements.
+
+ // Literal-first syntax: A tensor of 2 float32 elements.
dense<[10.0, 11.0]> : tensor<2xf32>
+
+ // Type-first syntax: A splat tensor of integer values.
+ dense<tensor<2xi32> : 10 : i32>
+
+ // Type-first syntax: A tensor of 2 float32 elements.
+ dense<tensor<2xf32> : [10.0, 11.0]>
```
+
+ Note: The literal-first syntax is supported only for complex, float, index,
+ int element types. The parser/print have special casing for these types.
+ Dense element attributes with other element types must use the type-first
+ syntax.
}];
let parameters = (ins AttributeSelfTypeParameter<"", "ShapedType">:$type,
"ArrayRef<char>":$rawData);
diff --git a/mlir/include/mlir/IR/BuiltinTypeInterfaces.h b/mlir/include/mlir/IR/BuiltinTypeInterfaces.h
index 5f14517d8dd71..9425d554b427c 100644
--- a/mlir/include/mlir/IR/BuiltinTypeInterfaces.h
+++ b/mlir/include/mlir/IR/BuiltinTypeInterfaces.h
@@ -19,6 +19,29 @@ struct fltSemantics;
namespace mlir {
class FloatType;
class MLIRContext;
+
+namespace detail {
+/// Float type implementation of
+/// DenseElementTypeInterface::getDenseElementBitSize.
+size_t getFloatTypeDenseElementBitSize(Type type);
+
+/// Float type implementation of DenseElementTypeInterface::convertToAttribute.
+Attribute convertFloatTypeToAttribute(Type type, llvm::ArrayRef<char> rawData);
+
+/// Float type implementation of
+/// DenseElementTypeInterface::convertFromAttribute.
+LogicalResult
+convertFloatTypeFromAttribute(Type type, Attribute attr,
+ llvm::SmallVectorImpl<char> &result);
+
+/// Read `bitWidth` bits from byte-aligned position in `rawData` and return as
+/// an APInt. Handles endianness correctly.
+llvm::APInt readBits(const char *rawData, size_t bitPos, size_t bitWidth);
+
+/// Write `value` to byte-aligned position `bitPos` in `rawData`. Handles
+/// endianness correctly.
+void writeBits(char *rawData, size_t bitPos, llvm::APInt value);
+} // namespace detail
} // namespace mlir
#include "mlir/IR/BuiltinTypeInterfaces.h.inc"
diff --git a/mlir/include/mlir/IR/BuiltinTypeInterfaces.td b/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
index 9ef08b7020b99..93c8c0694b467 100644
--- a/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
+++ b/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
@@ -41,12 +41,70 @@ def VectorElementTypeInterface : TypeInterface<"VectorElementTypeInterface"> {
}];
}
+//===----------------------------------------------------------------------===//
+// DenseElementTypeInterface
+//===----------------------------------------------------------------------===//
+
+def DenseElementTypeInterface : TypeInterface<"DenseElementType"> {
+ let cppNamespace = "::mlir";
+ let description = [{
+ This interface allows custom types to be used as element types in
+ DenseElementsAttr. Types implementing this interface define:
+
+ 1. The bit size for element storage.
+ 2. Helper methods for converting from/to Attribute. This assumes that there
+ is a corresponding attribute for each type that implements this
+ interface.
+
+ The helper methods for converting from/to Attribute are utilized when
+ parsing/printing IR or iterating over the elements via Attribute.
+ }];
+
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the number of bits required to store one element in dense
+ storage.
+
+ Note: The DenseElementsAttr infrastructure will automatically align
+ every element to a full byte in storage. This limitation could be lifted
+ in the future to support dense packing of non-byte-sized elements.
+ }],
+ /*retTy=*/"size_t",
+ /*methodName=*/"getDenseElementBitSize",
+ /*args=*/(ins)
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Attribute deserialization / attribute factory: Convert raw storage bytes
+ into an MLIR attribute. The size of `rawData` is
+ "ceilDiv(getDenseElementBitSize(), 8)".
+ }],
+ /*retTy=*/"::mlir::Attribute",
+ /*methodName=*/"convertToAttribute",
+ /*args=*/(ins "::llvm::ArrayRef<char>":$rawData)
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Attribute serialization: Convert an MLIR attribute into raw bytes.
+ Implementations must append "getDenseElementBitSize() / 8" values to
+ `result`. Return "failure" if the attribute is incompatible with this
+ element type.
+ }],
+ /*retTy=*/"::llvm::LogicalResult",
+ /*methodName=*/"convertFromAttribute",
+ /*args=*/(ins "::mlir::Attribute":$attr,
+ "::llvm::SmallVectorImpl<char>&":$result)
+ >,
+ ];
+}
+
//===----------------------------------------------------------------------===//
// FloatTypeInterface
//===----------------------------------------------------------------------===//
def FloatTypeInterface : TypeInterface<"FloatType",
- [VectorElementTypeInterface]> {
+ [DenseElementTypeInterface, VectorElementTypeInterface]> {
let cppNamespace = "::mlir";
let description = [{
This type interface should be implemented by all floating-point types. It
@@ -83,6 +141,21 @@ def FloatTypeInterface : TypeInterface<"FloatType",
/// The width includes the integer bit.
unsigned getFPMantissaWidth();
}];
+
+ let extraTraitClassDeclaration = [{
+ /// DenseElementTypeInterface implementations for float types.
+ size_t getDenseElementBitSize() const {
+ return ::mlir::detail::getFloatTypeDenseElementBitSize($_type);
+ }
+ ::mlir::Attribute convertToAttribute(::llvm::ArrayRef<char> rawData) const {
+ return ::mlir::detail::convertFloatTypeToAttribute($_type, rawData);
+ }
+ ::llvm::LogicalResult
+ convertFromAttribute(::mlir::Attribute attr,
+ ::llvm::SmallVectorImpl<char> &result) const {
+ return ::mlir::detail::convertFloatTypeFromAttribute($_type, attr, result);
+ }
+ }];
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/IR/BuiltinTypes.td b/mlir/include/mlir/IR/BuiltinTypes.td
index 806064faeda00..e7d0a03a85e7d 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.td
+++ b/mlir/include/mlir/IR/BuiltinTypes.td
@@ -45,7 +45,10 @@ def ValueSemantics : NativeTypeTrait<"ValueSemantics"> {
// ComplexType
//===----------------------------------------------------------------------===//
-def Builtin_Complex : Builtin_Type<"Complex", "complex"> {
+def Builtin_Complex : Builtin_Type<"Complex", "complex",
+ [DeclareTypeInterfaceMethods<DenseElementTypeInterface,
+ ["getDenseElementBitSize", "convertToAttribute", "convertFromAttribute"]>
+ ]> {
let summary = "Complex number with a parameterized element type";
let description = [{
Syntax:
@@ -560,7 +563,9 @@ def Builtin_Graph : Builtin_FunctionLike<"Graph", "graph">;
//===----------------------------------------------------------------------===//
def Builtin_Index : Builtin_Type<"Index", "index",
- [VectorElementTypeInterface]> {
+ [DeclareTypeInterfaceMethods<DenseElementTypeInterface,
+ ["getDenseElementBitSize", "convertToAttribute", "convertFromAttribute"]>,
+ VectorElementTypeInterface]> {
let summary = "Integer-like type with unknown platform-dependent bit width";
let description = [{
Syntax:
@@ -591,7 +596,10 @@ def Builtin_Index : Builtin_Type<"Index", "index",
//===----------------------------------------------------------------------===//
def Builtin_Integer : Builtin_Type<"Integer", "integer",
- [VectorElementTypeInterface, QuantStorageTypeInterface]> {
+ [VectorElementTypeInterface, QuantStorageTypeInterface,
+ DeclareTypeInterfaceMethods<DenseElementTypeInterface, [
+ "getDenseElementBitSize", "convertToAttribute",
+ "convertFromAttribute"]>]> {
let summary = "Integer type with arbitrary precision up to a fixed limit";
let description = [{
Syntax:
diff --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp
index 5978a11d06bc9..dc9744a42b730 100644
--- a/mlir/lib/AsmParser/AttributeParser.cpp
+++ b/mlir/lib/AsmParser/AttributeParser.cpp
@@ -16,6 +16,7 @@
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/DialectResourceBlobManager.h"
#include "mlir/IR/IntegerSet.h"
@@ -953,6 +954,119 @@ Attribute Parser::parseDenseArrayAttr(Type attrType) {
return eltParser.getAttr();
}
+/// Try to parse a dense elements attribute with the type-first syntax.
+/// Syntax: dense<TYPE : [ATTR, ATTR, ...]>
+/// This syntax is used for types other than int, float, index and complex.
+///
+/// Returns:
+/// - "null" attribute if this is not the type-first syntax.
+/// - "failure" in case of a parse error.
+/// - A valid Attribute otherwise.
+static FailureOr<Attribute> parseDenseElementsAttrTyped(Parser &p, SMLoc loc) {
+ // Skip l_paren because "parseType" would try to parse it as a tuple/function
+ // type, but '(' starts a complex literal like in the literal-first syntax.
+ if (p.getToken().is(Token::l_paren))
+ return Attribute();
+
+ // Parse type and valdiate that it's a shaped type.
+ auto typeLoc = p.getToken().getLoc();
+ Type type;
+ OptionalParseResult typeResult = p.parseOptionalType(type);
+ if (!typeResult.has_value())
+ return Attribute(); // Not type-first syntax.
+ if (failed(*typeResult))
+ return failure(); // Type parse error.
+
+ auto shapedType = dyn_cast<ShapedType>(type);
+ if (!shapedType) {
+ p.emitError(typeLoc, "expected a shaped type for dense elements");
+ return failure();
+ }
+ if (!shapedType.hasStaticShape()) {
+ p.emitError(typeLoc, "dense elements type must have static shape");
+ return failure();
+ }
+
+ // Check that the element type implements DenseElementTypeInterface.
+ auto denseEltType = dyn_cast<DenseElementType>(shapedType.getElementType());
+ if (!denseEltType) {
+ p.emitError(typeLoc,
+ "element type must implement DenseElementTypeInterface "
+ "for type-first dense syntax");
+ return failure();
+ }
+
+ // Parse colon.
+ if (p.parseToken(Token::colon, "expected ':' after type in dense attribute"))
+ return failure();
+
+ // Parse the element attributes and convert to raw bytes.
+ SmallVector<char> rawData;
+
+ // Helper to parse a single element.
+ auto parseSingleElement = [&]() -> ParseResult {
+ Attribute elemAttr = p.parseAttribute();
+ if (!elemAttr)
+ return failure();
+ if (failed(denseEltType.convertFromAttribute(elemAttr, rawData))) {
+ p.emitError("incompatible attribute for element type");
+ return failure();
+ }
+ return success();
+ };
+
+ // Recursively parse elements matching the expected shape.
+ std::function<ParseResult(ArrayRef<int64_t>)> parseElements;
+ parseElements = [&](ArrayRef<int64_t> remainingShape) -> ParseResult {
+ // Leaf: parse a single element.
+ if (remainingShape.empty())
+ return parseSingleElement();
+
+ // Non-leaf: expect a list with the correct number of elements.
+ int64_t expectedCount = remainingShape.front();
+ ArrayRef<int64_t> innerShape = remainingShape.drop_front();
+ int64_t actualCount = 0;
+
+ auto parseOne = [&]() -> ParseResult {
+ if (parseElements(innerShape))
+ return failure();
+ ++actualCount;
+ return success();
+ };
+
+ if (p.parseCommaSeparatedList(Parser::Delimiter::Square, parseOne))
+ return failure();
+
+ if (actualCount != expectedCount) {
+ p.emitError() << "expected " << expectedCount
+ << " elements in dimension, got " << actualCount;
+ return failure();
+ }
+ return success();
+ };
+
+ // Parse elements.
+ if (!p.getToken().is(Token::l_square)) {
+ // Single element - parse as splat.
+ if (parseSingleElement())
+ return failure();
+ } else if (shapedType.getShape().empty()) {
+ // Scalar type shouldn't have a list.
+ p.emitError(loc, "expected single element for scalar type, got list");
+ return failure();
+ } else {
+ // Parse structured literal matching the shape.
+ if (parseElements(shapedType.getShape()))
+ return failure();
+ }
+
+ if (p.parseToken(Token::greater, "expected '>' to close dense attribute"))
+ return failure();
+
+ // Create the attribute from raw buffer.
+ return DenseElementsAttr::getFromRawBuffer(shapedType, rawData);
+}
+
/// Parse a dense elements attribute.
Attribute Parser::parseDenseElementsAttr(Type attrType) {
auto attribLoc = getToken().getLoc();
@@ -960,7 +1074,16 @@ Attribute Parser::parseDenseElementsAttr(Type attrType) {
if (parseToken(Token::less, "expected '<' after 'dense'"))
return nullptr;
- // Parse the literal data if necessary.
+ // Try to parse the type-first syntax: dense<TYPE : [ATTR, ...]>
+ FailureOr<Attribute> typedResult =
+ parseDenseElementsAttrTyped(*this, attribLoc);
+ if (failed(typedResult))
+ return nullptr;
+ if (*typedResult)
+ return *typedResult;
+
+ // Try to parse the literal-first syntax, which is the default format for
+ // int, float, index and complex element types.
TensorLiteralParser literalParser(*this);
if (!consumeIf(Token::greater)) {
if (literalParser.parse(/*allowHex=*/true) ||
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 81455699421cc..b3242f838fc1d 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -507,11 +507,18 @@ class AsmPrinter::Impl {
/// Print a dense string elements attribute.
void printDenseStringElementsAttr(DenseStringElementsAttr attr);
- /// Print a dense elements attribute. If 'allowHex' is true, a hex string is
- /// used instead of individual elements when the elements attr is large.
+ /// Print a dense elements attribute in the literal-first syntax. If
+ /// 'allowHex' is true, a hex string is used instead of individual elements
+ /// when the elements attr is large.
void printDenseIntOrFPElementsAttr(DenseIntOrFPElementsAttr attr,
bool allowHex);
+ /// Print a dense elements attribute using the type-first syntax and the
+ /// DenseElementTypeInterface, which provides the attribute printer for each
+ /// element.
+ void printTypeFirstDenseElementsAttr(DenseElementsAttr attr,
+ DenseElementType denseEltType);
+
/// Print a dense array attribute.
void printDenseArrayAttr(DenseArrayAttr attr);
@@ -2507,7 +2514,17 @@ void AsmPrinter::Impl::printAttributeImpl(Attribute attr,
printElidedElementsAttr(os);
} else {
os << "dense<";
- printDenseIntOrFPElementsAttr(intOrFpEltAttr, /*allowHex=*/true);
+ // Check if the element type implements DenseElementTypeInterface and is
+ // not a built-in type. Built-in types (int, float, index, complex) use
+ // the existing printing format for backwards compatibility.
+ Type eltType = intOrFpEltAttr.getElementType();
+ if (isa<FloatType, IntegerType, IndexType, ComplexType>(eltType)) {
+ printDenseIntOrFPElementsAttr(intOrFpEltAttr, /*allowHex=*/true);
+ } else {
+ printTypeFirstDenseElementsAttr(intOrFpEltAttr,
+ cast<DenseElementType>(eltType));
+ typeElision = AttrTypeElision::Must;
+ }
os << '>';
}
@@ -2705,6 +2722,27 @@ void AsmPrinter::Impl::printDenseStringElementsAttr(
printDenseElementsAttrImpl(attr.isSplat(), attr.getType(), os, printFn);
}
+void AsmPrinter::Impl::printTypeFirstDenseElementsAttr(
+ DenseElementsAttr attr, DenseElementType denseEltType) {
+ // Print the type first: dense<TYPE : [ELEMENTS]>
+ printType(attr.getType());
+ os << " : ";
+
+ ArrayRef<char> rawData = attr.getRawData();
+ // Storage is byte-aligned: align bit size up to next byte boundary.
+ size_t bitSize = denseEltType.getDenseElementBitSize();
+ size_t byteSize = llvm::divideCeil(bitSize, static_cast<size_t>(CHAR_BIT));
+
+ // Print elements: convert raw bytes to attribute, then print attribute.
+ printDenseElementsAttrImpl(
+ attr.isSplat(), attr.getType(), os, [&](unsigned index) {
+ size_t offset = attr.isSplat() ? 0 : index * byteSize;
+ ArrayRef<char> elemData = rawData.slice(offset, byteSize);
+ Attribute elemAttr = denseEltType.convertToAttribute(elemData);
+ printAttributeImpl(elemAttr);
+ });
+}
+
void AsmPrinter::Impl::printDenseArrayAttr(DenseArrayAttr attr) {
Type type = attr.getElementType();
unsigned bitwidth = type.isInteger(1) ? 8 : type.getIntOrFloatBitWidth();
diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h
index 1f268603cf37f..8505149afdd9c 100644
--- a/mlir/lib/IR/AttributeDetail.h
+++ b/mlir/lib/IR/AttributeDetail.h
@@ -16,6 +16,7 @@
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/AttributeSupport.h"
#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/IntegerSet.h"
#include "mlir/IR/MLIRContext.h"
@@ -32,12 +33,9 @@ namespace detail {
/// Return the bit width which DenseElementsAttr should use for this type.
inline size_t getDenseElementBitWidth(Type eltType) {
- // Align the width for complex to 8 to make storage and interpretation easier.
- if (ComplexType comp = llvm::dyn_cast<ComplexType>(eltType))
- return llvm::alignTo<8>(getDenseElementBitWidth(comp.getElementType())) * 2;
- if (eltType.isIndex())
- return IndexType::kInternalStorageBitWidth;
- return eltType.getIntOrFloatBitWidth();
+ if (auto denseEltType = llvm::dyn_cast<DenseElementType>(eltType))
+ return denseEltType.getDenseElementBitSize();
+ llvm_unreachable("unsupported element type");
}
/// An attribute representing a reference to a dense vector or tensor object.
diff --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index 1a29fc534b40f..bbbc9198a68ab 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/IR/BuiltinAttributes.cpp
@@ -10,6 +10,7 @@
#include "AttributeDetail.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/DialectResourceBlobManager.h"
#include "mlir/IR/IntegerSet.h"
@@ -527,7 +528,7 @@ static void copyArrayToAPIntForBEmachine(const char *inArray, size_t numBytes,
}
/// Writes value to the bit position `bitPos` in array `rawData`.
-static void writeBits(char *rawData, size_t bitPos, APInt value) {
+void mlir::detail::writeBits(char *rawData, size_t bitPos, APInt value) {
size_t bitWidth = value.getBitWidth();
// The bit position is guaranteed to be byte aligned.
@@ -549,7 +550,8 @@ static void writeBits(char *rawData, size_t bitPos, APInt value) {
/// Reads the next `bitWidth` bits from the bit position `bitPos` in array
/// `rawData`.
-static APInt readBits(const char *rawData, size_t bitPos, size_t bitWidth) {
+APInt mlir::detail::readBits(const char *rawData, size_t bitPos,
+ size_t bitWidth) {
// The bit position is guaranteed to be byte aligned.
assert((bitPos % CHAR_BIT) == 0 && "expected bitPos to be 8-bit aligned");
APInt result(bitWidth, 0);
@@ -595,39 +597,21 @@ DenseElementsAttr::AttributeElementIterator::AttributeElementIterator(
Attribute DenseElementsAttr::AttributeElementIterator::operator*() const {
auto owner = llvm::cast<DenseElementsAttr>(getFromOpaquePointer(base));
Type eltTy = owner.getElementType();
- if (llvm::dyn_cast<IntegerType>(eltTy))
- return IntegerAttr::get(eltTy, *IntElementIterator(owner, index));
- if (llvm::isa<IndexType>(eltTy))
- return IntegerAttr::get(eltTy, *IntElementIterator(owner, index));
- if (auto floatEltTy = llvm::dyn_cast<FloatType>(eltTy)) {
- IntElementIterator intIt(owner, index);
- FloatElementIterator floatIt(floatEltTy.getFloatSemantics(), intIt);
- return FloatAttr::get(eltTy, *floatIt);
- }
- if (auto complexTy = llvm::dyn_cast<ComplexType>(eltTy)) {
- auto complexEltTy = complexTy.getElementType();
- ComplexIntElementIterator complexIntIt(owner, index);
- if (llvm::isa<IntegerType>(complexEltTy)) {
- auto value = *complexIntIt;
- auto real = IntegerAttr::get(complexEltTy, value.real());
- auto imag = IntegerAttr::get(complexEltTy, value.imag());
- return ArrayAttr::get(complexTy.getContext(),
- ArrayRef<Attribute>{real, imag});
- }
- ComplexFloatElementIterator complexFloatIt(
- llvm::cast<FloatType>(complexEltTy).getFloatSemantics(), complexIntIt);
- auto value = *complexFloatIt;
- auto real = FloatAttr::get(complexEltTy, value.real());
- auto imag = FloatAttr::get(complexEltTy, value.imag());
- return ArrayAttr::get(complexTy.getContext(),
- ArrayRef<Attribute>{real, imag});
- }
+ // Handle strings specially.
if (llvm::isa<DenseStringElementsAttr>(owner)) {
ArrayRef<StringRef> vals = owner.getRawStringData();
return StringAttr::get(owner.isSplat() ? vals.front() : vals[index], eltTy);
}
- llvm_unreachable("unexpected element type");
+
+ // All other types should implement DenseElementTypeInterface.
+ auto denseEltTy = llvm::cast<DenseElementType>(eltTy);
+ ArrayRef<char> rawData = owner.getRawData();
+ // Storage is byte-aligned: align bit size up to next byte boundary.
+ size_t bitSize = denseEltTy.getDenseElementBitSize();
+ size_t byteSize = llvm::divideCeil(bitSize, CHAR_BIT);
+ size_t offset = owner.isSplat() ? 0 : index * byteSize;
+ return denseEltTy.convertToAttribute(rawData.slice(offset, byteSize));
}
//===----------------------------------------------------------------------===//
@@ -888,79 +872,28 @@ DenseElementsAttr DenseElementsAttr::get(ShapedType type,
assert(hasSameNumElementsOrSplat(type, values));
Type eltType = type.getElementType();
- // Take care complex type case first.
- if (auto complexType = llvm::dyn_cast<ComplexType>(eltType)) {
- if (complexType.getElementType().isIntOrIndex()) {
- SmallVector<std::complex<APInt>> complexValues;
- complexValues.reserve(values.size());
- for (Attribute attr : values) {
- assert(llvm::isa<ArrayAttr>(attr) && "expected ArrayAttr for complex");
- auto arrayAttr = llvm::cast<ArrayAttr>(attr);
- assert(arrayAttr.size() == 2 && "expected 2 element for complex");
- auto attr0 = arrayAttr[0];
- auto attr1 = arrayAttr[1];
- complexValues.push_back(
- std::complex<APInt>(llvm::cast<IntegerAttr>(attr0).getValue(),
- llvm::cast<IntegerAttr>(attr1).getValue()));
- }
- return DenseElementsAttr::get(type, complexValues);
- }
- // Must be float.
- SmallVector<std::complex<APFloat>> complexValues;
- complexValues.reserve(values.size());
- for (Attribute attr : values) {
- assert(llvm::isa<ArrayAttr>(attr) && "expected ArrayAttr for complex");
- auto arrayAttr = llvm::cast<ArrayAttr>(attr);
- assert(arrayAttr.size() == 2 && "expected 2 element for complex");
- auto attr0 = arrayAttr[0];
- auto attr1 = arrayAttr[1];
- complexValues.push_back(
- std::complex<APFloat>(llvm::cast<FloatAttr>(attr0).getValue(),
- llvm::cast<FloatAttr>(attr1).getValue()));
- }
- return DenseElementsAttr::get(type, complexValues);
- }
-
- // If the element type is not based on int/float/index, assume it is a string
- // type.
- if (!eltType.isIntOrIndexOrFloat()) {
+ // Handle strings specially.
+ if (!llvm::isa<DenseElementType>(eltType)) {
SmallVector<StringRef, 8> stringValues;
stringValues.reserve(values.size());
for (Attribute attr : values) {
assert(llvm::isa<StringAttr>(attr) &&
- "expected string value for non integer/index/float element");
+ "expected string value for non-DenseElementType element");
stringValues.push_back(llvm::cast<StringAttr>(attr).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);
-
- // Compress the attribute values into a character buffer.
- SmallVector<char, 8> data(
- llvm::divideCeil(storageBitWidth * values.size(), CHAR_BIT));
- APInt intVal;
- for (unsigned i = 0, e = values.size(); i < e; ++i) {
- if (auto floatAttr = llvm::dyn_cast<FloatAttr>(values[i])) {
- assert(floatAttr.getType() == eltType &&
- "expected float attribute type to equal element type");
- intVal = floatAttr.getValue().bitcastToAPInt();
- } else if (auto intAttr = llvm::dyn_cast<IntegerAttr>(values[i])) {
- assert(intAttr.getType() == eltType &&
- "expected integer attribute type to equal element type");
- intVal = intAttr.getValue();
- } else {
- // Unsupported attribute type.
+ // All other types go through DenseElementTypeInterface.
+ auto denseEltType = llvm::dyn_cast<DenseElementType>(eltType);
+ assert(denseEltType &&
+ "attempted to get DenseElementsAttr with unsupported element type");
+ SmallVector<char> data;
+ for (Attribute attr : values) {
+ LogicalResult result = denseEltType.convertFromAttribute(attr, data);
+ if (failed(result))
return {};
- }
-
- assert(intVal.getBitWidth() == bitWidth &&
- "expected value to have same bitwidth as element type");
- writeBits(data.data(), i * storageBitWidth, intVal);
}
-
return DenseIntOrFPElementsAttr::getRaw(type, data);
}
diff --git a/mlir/lib/IR/BuiltinTypeInterfaces.cpp b/mlir/lib/IR/BuiltinTypeInterfaces.cpp
index 2f063be3e7cd0..29303d95eb003 100644
--- a/mlir/lib/IR/BuiltinTypeInterfaces.cpp
+++ b/mlir/lib/IR/BuiltinTypeInterfaces.cpp
@@ -6,9 +6,12 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/Support/CheckedArithmetic.h"
+#include "llvm/Support/MathExtras.h"
+#include <climits>
using namespace mlir;
using namespace mlir::detail;
@@ -19,6 +22,37 @@ using namespace mlir::detail;
#include "mlir/IR/BuiltinTypeInterfaces.cpp.inc"
+//===----------------------------------------------------------------------===//
+// DenseElementTypeInterface implementations for float types
+//===----------------------------------------------------------------------===//
+
+size_t mlir::detail::getFloatTypeDenseElementBitSize(Type type) {
+ return cast<FloatType>(type).getWidth();
+}
+
+Attribute mlir::detail::convertFloatTypeToAttribute(Type type,
+ ArrayRef<char> rawData) {
+ auto floatType = cast<FloatType>(type);
+ APInt intVal = readBits(rawData.data(), /*bitPos=*/0, floatType.getWidth());
+ APFloat floatVal(floatType.getFloatSemantics(), intVal);
+ return FloatAttr::get(type, floatVal);
+}
+
+LogicalResult
+mlir::detail::convertFloatTypeFromAttribute(Type type, Attribute attr,
+ SmallVectorImpl<char> &result) {
+ auto floatType = cast<FloatType>(type);
+ auto floatAttr = dyn_cast<FloatAttr>(attr);
+ if (!floatAttr || floatAttr.getType() != type)
+ return failure();
+ size_t byteSize =
+ llvm::divideCeil(floatType.getWidth(), static_cast<unsigned>(CHAR_BIT));
+ size_t bitPos = result.size() * CHAR_BIT;
+ result.resize(result.size() + byteSize);
+ writeBits(result.data(), bitPos, floatAttr.getValue().bitcastToAPInt());
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// FloatType
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp
index 1e198043c590a..786c30851a071 100644
--- a/mlir/lib/IR/BuiltinTypes.cpp
+++ b/mlir/lib/IR/BuiltinTypes.cpp
@@ -12,14 +12,17 @@
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/TensorEncoding.h"
#include "mlir/IR/TypeUtilities.h"
#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/APInt.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/CheckedArithmetic.h"
+#include <cstring>
using namespace mlir;
using namespace mlir::detail;
@@ -58,6 +61,39 @@ LogicalResult ComplexType::verify(function_ref<InFlightDiagnostic()> emitError,
return success();
}
+size_t ComplexType::getDenseElementBitSize() const {
+ auto elemTy = cast<DenseElementType>(getElementType());
+ return llvm::alignTo<8>(elemTy.getDenseElementBitSize()) * 2;
+}
+
+Attribute ComplexType::convertToAttribute(ArrayRef<char> rawData) const {
+ auto elemTy = cast<DenseElementType>(getElementType());
+ size_t singleElementBytes =
+ llvm::alignTo<8>(elemTy.getDenseElementBitSize()) / 8;
+ Attribute real =
+ elemTy.convertToAttribute(rawData.take_front(singleElementBytes));
+ Attribute imag =
+ elemTy.convertToAttribute(rawData.take_back(singleElementBytes));
+ return ArrayAttr::get(getContext(), {real, imag});
+}
+
+LogicalResult
+ComplexType::convertFromAttribute(Attribute attr,
+ SmallVectorImpl<char> &result) const {
+ auto arrayAttr = dyn_cast<ArrayAttr>(attr);
+ if (!arrayAttr || arrayAttr.size() != 2)
+ return failure();
+ auto elemTy = cast<DenseElementType>(getElementType());
+ SmallVector<char> realData, imagData;
+ if (failed(elemTy.convertFromAttribute(arrayAttr[0], realData)))
+ return failure();
+ if (failed(elemTy.convertFromAttribute(arrayAttr[1], imagData)))
+ return failure();
+ result.append(realData);
+ result.append(imagData);
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// Integer Type
//===----------------------------------------------------------------------===//
@@ -85,6 +121,57 @@ IntegerType IntegerType::scaleElementBitwidth(unsigned scale) {
return IntegerType::get(getContext(), scale * getWidth(), getSignedness());
}
+size_t IntegerType::getDenseElementBitSize() const {
+ // Return the actual bit width. Storage alignment is handled separately.
+ return getWidth();
+}
+
+Attribute IntegerType::convertToAttribute(ArrayRef<char> rawData) const {
+ APInt value = detail::readBits(rawData.data(), /*bitPos=*/0, getWidth());
+ return IntegerAttr::get(*this, value);
+}
+
+static void writeAPIntToVector(APInt apInt, SmallVectorImpl<char> &result) {
+ size_t byteSize = llvm::divideCeil(apInt.getBitWidth(), CHAR_BIT);
+ size_t bitPos = result.size() * CHAR_BIT;
+ result.resize(result.size() + byteSize);
+ detail::writeBits(result.data(), bitPos, apInt);
+}
+
+LogicalResult
+IntegerType::convertFromAttribute(Attribute attr,
+ SmallVectorImpl<char> &result) const {
+ auto intAttr = dyn_cast<IntegerAttr>(attr);
+ if (!intAttr || intAttr.getType() != *this)
+ return failure();
+ writeAPIntToVector(intAttr.getValue(), result);
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
+// Index Type
+//===----------------------------------------------------------------------===//
+
+size_t IndexType::getDenseElementBitSize() const {
+ return kInternalStorageBitWidth;
+}
+
+Attribute IndexType::convertToAttribute(ArrayRef<char> rawData) const {
+ APInt value =
+ detail::readBits(rawData.data(), /*bitPos=*/0, kInternalStorageBitWidth);
+ return IntegerAttr::get(*this, value);
+}
+
+LogicalResult
+IndexType::convertFromAttribute(Attribute attr,
+ SmallVectorImpl<char> &result) const {
+ auto intAttr = dyn_cast<IntegerAttr>(attr);
+ if (!intAttr || intAttr.getType() != *this)
+ return failure();
+ writeAPIntToVector(intAttr.getValue(), result);
+ return success();
+}
+
//===----------------------------------------------------------------------===//
// Float Types
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/IR/dense-elements-type-interface.mlir b/mlir/test/IR/dense-elements-type-interface.mlir
new file mode 100644
index 0000000000000..8749e562087c2
--- /dev/null
+++ b/mlir/test/IR/dense-elements-type-interface.mlir
@@ -0,0 +1,83 @@
+// RUN: mlir-opt %s -verify-diagnostics -split-input-file | FileCheck %s
+
+// Test dense elements attribute with custom element type using DenseElementTypeInterface.
+// Uses the new type-first syntax: dense<TYPE : [ATTR, ...]>
+// Note: The type is embedded in the attribute, so it's not printed again at the end.
+
+// CHECK-LABEL: func @dense_custom_element_type
+func.func @dense_custom_element_type() {
+ // CHECK: "test.dummy"() {attr = dense<tensor<3x!test.dense_element> : [1 : i32, 2 : i32, 3 : i32]>}
+ "test.dummy"() {attr = dense<tensor<3x!test.dense_element> : [1 : i32, 2 : i32, 3 : i32]>} : () -> ()
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @dense_custom_element_type_2d
+func.func @dense_custom_element_type_2d() {
+ // CHECK: "test.dummy"() {attr = dense<tensor<2x2x!test.dense_element> : {{\[}}{{\[}}1 : i32, 2 : i32], [3 : i32, 4 : i32]]>}
+ "test.dummy"() {attr = dense<tensor<2x2x!test.dense_element> : [[1 : i32, 2 : i32], [3 : i32, 4 : i32]]>} : () -> ()
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @dense_custom_element_splat
+func.func @dense_custom_element_splat() {
+ // CHECK: "test.dummy"() {attr = dense<tensor<4x!test.dense_element> : 42 : i32>}
+ "test.dummy"() {attr = dense<tensor<4x!test.dense_element> : 42 : i32>} : () -> ()
+ return
+}
+
+// -----
+
+// CHECK-LABEL func @dense_i32_1d
+func.func @dense_i32_1d() {
+ // The default assembly format for int, index, float, complex element types is
+ // the literal-first syntax. Such a dense elements attribute can be parsed
+ // with the type-first syntax, but it will come back with the literal-first
+ // syntax.
+ // CHECK: "test.dummy"() {attr = dense<[1, 2, 3]> : tensor<3xi32>} : () -> ()
+ "test.dummy"() {attr = dense<tensor<3xi32> : [1 : i32, 2 : i32, 3 : i32]>} : () -> ()
+ return
+}
+
+// -----
+
+func.func @invalid_element() {
+ // expected-error @+1 {{expected attribute value}}
+ "test.dummy"() {attr = dense<tensor<3xi32> : [foo]>} : () -> ()
+ return
+}
+
+// -----
+
+func.func @incompatible_attribute() {
+ // expected-error @+1 {{incompatible attribute for element type}}
+ "test.dummy"() {attr = dense<tensor<3xi32> : ["foo"]>} : () -> ()
+ return
+}
+
+// -----
+
+func.func @shape_mismatch() {
+ // expected-error @+1 {{expected 3 elements in dimension, got 2}}
+ "test.dummy"() {attr = dense<tensor<3xi32> : [1 : i32, 2 : i32]>} : () -> ()
+ return
+}
+
+// -----
+
+func.func @dynamic_shape() {
+ // expected-error @+1 {{dense elements type must have static shape}}
+ "test.dummy"() {attr = dense<tensor<?xi32> : [1 : i32, 2 : i32, 3 : i32]>} : () -> ()
+ return
+}
+
+// -----
+
+func.func @invalid_type() {
+ // expected-error @+1 {{expected a shaped type for dense elements}}
+ "test.dummy"() {attr = dense<i32 : [1 : i32, 2 : i32, 3 : i32]>} : () -> ()
+ return
+}
diff --git a/mlir/test/lib/Dialect/Test/TestTypeDefs.td b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
index 964792ceebc07..08600ce713a17 100644
--- a/mlir/test/lib/Dialect/Test/TestTypeDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
@@ -18,6 +18,7 @@ include "TestDialect.td"
include "TestAttrDefs.td"
include "TestInterfaces.td"
include "mlir/IR/BuiltinTypes.td"
+include "mlir/IR/BuiltinTypeInterfaces.td"
include "mlir/Interfaces/DataLayoutInterfaces.td"
include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td"
@@ -512,4 +513,15 @@ def TestTypeNewlineAndIndent : Test_Type<"TestTypeNewlineAndIndent"> {
let hasCustomAssemblyFormat = 1;
}
+def TestTypeDenseElement : Test_Type<"TestDenseElement",
+ [DeclareTypeInterfaceMethods<DenseElementTypeInterface,
+ ["getDenseElementBitSize", "convertToAttribute", "convertFromAttribute"]>
+ ]> {
+ let mnemonic = "dense_element";
+ let description = [{
+ A test type that implements DenseElementTypeInterface to test dense
+ elements with custom element types. Elements are stored as 32-bit integers.
+ }];
+}
+
#endif // TEST_TYPEDEFS
diff --git a/mlir/test/lib/Dialect/Test/TestTypes.cpp b/mlir/test/lib/Dialect/Test/TestTypes.cpp
index 71dd25b0093e0..ef3396fc4f610 100644
--- a/mlir/test/lib/Dialect/Test/TestTypes.cpp
+++ b/mlir/test/lib/Dialect/Test/TestTypes.cpp
@@ -15,6 +15,7 @@
#include "TestDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/ExtensibleDialect.h"
#include "mlir/IR/Types.h"
@@ -22,6 +23,7 @@
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/TypeSize.h"
+#include <cstring>
#include <optional>
using namespace mlir;
@@ -605,3 +607,29 @@ void TestTypeNewlineAndIndentType::print(::mlir::AsmPrinter &printer) const {
printer.printNewline();
printer << ">";
}
+
+//===----------------------------------------------------------------------===//
+// TestDenseElementType - DenseElementTypeInterface Implementation
+//===----------------------------------------------------------------------===//
+
+// Elements are stored as 32-bit integers.
+size_t TestDenseElementType::getDenseElementBitSize() const { return 32; }
+
+Attribute
+TestDenseElementType::convertToAttribute(ArrayRef<char> rawData) const {
+ assert(rawData.size() == 4 && "expected 4 bytes for TestDenseElement");
+ int32_t value;
+ std::memcpy(&value, rawData.data(), sizeof(value));
+ return IntegerAttr::get(IntegerType::get(getContext(), 32), value);
+}
+
+LogicalResult TestDenseElementType::convertFromAttribute(
+ Attribute attr, SmallVectorImpl<char> &result) const {
+ auto intAttr = dyn_cast<IntegerAttr>(attr);
+ if (!intAttr || intAttr.getType().getIntOrFloatBitWidth() != 32)
+ return failure();
+ int32_t value = intAttr.getValue().getSExtValue();
+ result.append(reinterpret_cast<const char *>(&value),
+ reinterpret_cast<const char *>(&value) + sizeof(value));
+ return success();
+}
diff --git a/mlir/test/lib/Dialect/Test/TestTypes.h b/mlir/test/lib/Dialect/Test/TestTypes.h
index 6499a96f495d0..705fb86e9e9b3 100644
--- a/mlir/test/lib/Dialect/Test/TestTypes.h
+++ b/mlir/test/lib/Dialect/Test/TestTypes.h
@@ -19,6 +19,7 @@
#include "TestTraits.h"
#include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/DialectImplementation.h"
>From 638224b50da1b0310cd991c5e08d71483ae5812c Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Fri, 13 Feb 2026 16:58:41 +0000
Subject: [PATCH 2/2] [mlir][IR] Separate `DenseStringElementsAttr` from
`DenseElementsAttr`
---
mlir/include/mlir/IR/BuiltinAttributes.h | 37 ++----
mlir/include/mlir/IR/BuiltinAttributes.td | 110 +++++++++++++++---
mlir/include/mlir/IR/CommonAttrConstraints.td | 4 +-
mlir/lib/AsmParser/AttributeParser.cpp | 37 ++++--
mlir/lib/CAPI/IR/BuiltinAttributes.cpp | 22 ++--
mlir/lib/IR/BuiltinAttributes.cpp | 70 ++++-------
mlir/test/IR/parser.mlir | 6 +-
mlir/test/mlir-tblgen/openmp-clause-ops.td | 2 +-
mlir/unittests/IR/AttributeTest.cpp | 39 +++++--
9 files changed, 198 insertions(+), 129 deletions(-)
diff --git a/mlir/include/mlir/IR/BuiltinAttributes.h b/mlir/include/mlir/IR/BuiltinAttributes.h
index ee6a8f4e4d948..3ba943c7ccd41 100644
--- a/mlir/include/mlir/IR/BuiltinAttributes.h
+++ b/mlir/include/mlir/IR/BuiltinAttributes.h
@@ -152,10 +152,6 @@ class DenseElementsAttr : public Attribute {
/// Overload of the above 'get' method that is specialized for boolean values.
static DenseElementsAttr get(ShapedType type, ArrayRef<bool> values);
- /// Overload of the above 'get' method that is specialized for StringRef
- /// values.
- static DenseElementsAttr get(ShapedType type, ArrayRef<StringRef> values);
-
/// Constructs a dense integer elements attribute from an array of APInt
/// values. Each APInt value is expected to have the same bitwidth as the
/// element type of 'type'. 'type' must be a vector or tensor with static
@@ -223,7 +219,8 @@ class DenseElementsAttr : public Attribute {
decltype(std::declval<AttrT>().template getValues<T>());
/// A utility iterator that allows walking over the internal Attribute values
- /// of a DenseElementsAttr.
+ /// of a dense elements attribute (DenseElementsAttr or
+ /// DenseStringElementsAttr).
class AttributeElementIterator
: public llvm::indexed_accessor_iterator<AttributeElementIterator,
const void *, Attribute,
@@ -232,11 +229,9 @@ class DenseElementsAttr : public Attribute {
/// Accesses the Attribute value at this iterator position.
Attribute operator*() const;
- private:
- friend DenseElementsAttr;
-
- /// Constructs a new iterator.
- AttributeElementIterator(DenseElementsAttr attr, size_t index);
+ /// Constructs a new iterator. Accepts any attribute implementing
+ /// ElementsAttr (e.g. DenseElementsAttr, DenseStringElementsAttr).
+ AttributeElementIterator(Attribute attr, size_t index);
};
/// Iterator for walking raw element values of the specified type 'T', which
@@ -461,21 +456,6 @@ class DenseElementsAttr : public Attribute {
ElementIterator<T>(rawData, splat, getNumElements()));
}
- /// Try to get the held element values as a range of StringRef.
- template <typename T>
- using StringRefValueTemplateCheckT =
- std::enable_if_t<std::is_same<T, StringRef>::value>;
- template <typename T, typename = StringRefValueTemplateCheckT<T>>
- FailureOr<iterator_range_impl<ElementIterator<StringRef>>>
- tryGetValues() const {
- auto stringRefs = getRawStringData();
- const char *ptr = reinterpret_cast<const char *>(stringRefs.data());
- bool splat = isSplat();
- return iterator_range_impl<ElementIterator<StringRef>>(
- getType(), ElementIterator<StringRef>(ptr, splat, 0),
- ElementIterator<StringRef>(ptr, splat, getNumElements()));
- }
-
/// Try to get the held element values as a range of Attributes.
template <typename T>
using AttributeValueTemplateCheckT =
@@ -484,8 +464,8 @@ class DenseElementsAttr : public Attribute {
FailureOr<iterator_range_impl<AttributeElementIterator>>
tryGetValues() const {
return iterator_range_impl<AttributeElementIterator>(
- getType(), AttributeElementIterator(*this, 0),
- AttributeElementIterator(*this, getNumElements()));
+ getType(), AttributeElementIterator(Attribute(*this), 0),
+ AttributeElementIterator(Attribute(*this), getNumElements()));
}
/// Try to get the held element values a range of T, where T is a derived
@@ -578,9 +558,6 @@ class DenseElementsAttr : public Attribute {
/// form the user might expect.
ArrayRef<char> getRawData() const;
- /// Return the raw StringRef data held by this attribute.
- ArrayRef<StringRef> getRawStringData() const;
-
/// Return the type of this ElementsAttr, guaranteed to be a vector or tensor
/// with static shape.
ShapedType getType() const;
diff --git a/mlir/include/mlir/IR/BuiltinAttributes.td b/mlir/include/mlir/IR/BuiltinAttributes.td
index dced379d1f979..064783ae5f87a 100644
--- a/mlir/include/mlir/IR/BuiltinAttributes.td
+++ b/mlir/include/mlir/IR/BuiltinAttributes.td
@@ -395,8 +395,7 @@ def Builtin_DenseIntOrFPElementsAttr : Builtin_Attr<
//===----------------------------------------------------------------------===//
def Builtin_DenseStringElementsAttr : Builtin_Attr<
- "DenseStringElements", "dense_string_elements", [ElementsAttrInterface],
- "DenseElementsAttr"
+ "DenseStringElements", "dense_string_elements", [ElementsAttrInterface]
> {
let summary = "An Attribute containing a dense multi-dimensional array of "
"strings";
@@ -431,13 +430,97 @@ def Builtin_DenseStringElementsAttr : Builtin_Attr<
}]>,
];
let extraClassDeclaration = [{
- using DenseElementsAttr::empty;
- using DenseElementsAttr::getNumElements;
- using DenseElementsAttr::getElementType;
- using DenseElementsAttr::getValues;
- using DenseElementsAttr::isSplat;
- using DenseElementsAttr::size;
- using DenseElementsAttr::value_begin;
+ /// Iterator for walking StringRef element values.
+ class StringRefElementIterator
+ : public detail::DenseElementIndexedIteratorImpl<StringRefElementIterator,
+ const StringRef> {
+ public:
+ const StringRef &operator*() const {
+ return reinterpret_cast<const StringRef *>(this->getData())[this->getDataIndex()];
+ }
+ StringRefElementIterator(const char *data, bool isSplat, size_t dataIndex)
+ : detail::DenseElementIndexedIteratorImpl<StringRefElementIterator,
+ const StringRef>(
+ data, isSplat, dataIndex) {}
+ };
+
+ /// Iterator for walking element values as Attribute (StringAttr).
+ class StringAttributeElementIterator
+ : public llvm::indexed_accessor_iterator<StringAttributeElementIterator,
+ const void *, Attribute,
+ Attribute, Attribute> {
+ public:
+ Attribute operator*() const;
+ StringAttributeElementIterator(const DenseStringElementsAttr *attr,
+ size_t index)
+ : llvm::indexed_accessor_iterator<StringAttributeElementIterator,
+ const void *, Attribute,
+ Attribute, Attribute>(
+ attr->getAsOpaquePointer(), index) {}
+ };
+
+ /// Return the type of this attribute (vector or tensor with static shape).
+ ShapedType getType() const;
+
+ /// Helper methods for ElementsAttr interface.
+ bool empty() const { return getNumElements() == 0; }
+ int64_t getNumElements() const { return getType().getNumElements(); }
+ Type getElementType() const { return getType().getElementType(); }
+ bool isSplat() const { return getRawStringData().size() == 1; }
+ int64_t size() const { return getNumElements(); }
+
+ /// Return the raw StringRef data held by this attribute.
+ ArrayRef<StringRef> getRawStringData() const;
+
+ /// Try to get the held element values as a range of StringRef.
+ template <typename T>
+ using StringRefValueTemplateCheckT =
+ std::enable_if_t<std::is_same<T, StringRef>::value>;
+ template <typename T, typename = StringRefValueTemplateCheckT<T>>
+ FailureOr<detail::ElementsAttrRange<StringRefElementIterator>>
+ tryGetValues() const {
+ auto stringRefs = getRawStringData();
+ const char *ptr = reinterpret_cast<const char *>(stringRefs.data());
+ bool splat = isSplat();
+ return detail::ElementsAttrRange<StringRefElementIterator>(
+ getType(), StringRefElementIterator(ptr, splat, 0),
+ StringRefElementIterator(ptr, splat, getNumElements()));
+ }
+
+ /// Try to get the held element values as a range of Attributes.
+ template <typename T>
+ using AttributeValueTemplateCheckT =
+ std::enable_if_t<std::is_same<T, Attribute>::value>;
+ template <typename T, typename = AttributeValueTemplateCheckT<T>>
+ FailureOr<detail::ElementsAttrRange<StringAttributeElementIterator>>
+ tryGetValues() const {
+ return detail::ElementsAttrRange<StringAttributeElementIterator>(
+ getType(), StringAttributeElementIterator(this, 0),
+ StringAttributeElementIterator(this, getNumElements()));
+ }
+
+ template <typename T>
+ auto getValues() const {
+ auto range = tryGetValues<T>();
+ assert(succeeded(range) && "element type cannot be iterated");
+ return std::move(*range);
+ }
+ template <typename T>
+ auto value_begin() const { return getValues<T>().begin(); }
+ template <typename T>
+ auto value_end() const { return getValues<T>().end(); }
+ /// Return the splat value. Asserts that the attribute is a splat.
+ template <typename T>
+ auto getSplatValue() const {
+ assert(isSplat() && "expected the attribute to be a splat");
+ return *value_begin<T>();
+ }
+ template <typename T>
+ auto try_value_begin() const {
+ auto range = tryGetValues<T>();
+ using iterator = decltype(range->begin());
+ return failed(range) ? FailureOr<iterator>(failure()) : range->begin();
+ }
/// The set of data types that can be iterated by this attribute.
using ContiguousIterableTypesT = std::tuple<StringRef>;
@@ -449,11 +532,6 @@ def Builtin_DenseStringElementsAttr : Builtin_Attr<
auto try_value_begin_impl(OverloadToken<T>) const {
return try_value_begin<T>();
}
-
- protected:
- friend DenseElementsAttr;
-
- public:
}];
let genAccessors = 0;
let genStorageClass = 0;
@@ -931,9 +1009,7 @@ def Builtin_SparseElementsAttr : Builtin_Attr<
std::complex<int16_t>, std::complex<int32_t>, std::complex<int64_t>,
// Float types.
APFloat, float, double,
- std::complex<APFloat>, std::complex<float>, std::complex<double>,
- // String types.
- StringRef
+ std::complex<APFloat>, std::complex<float>, std::complex<double>
>;
using ElementsAttr::Trait<SparseElementsAttr>::getValues;
using ElementsAttr::Trait<SparseElementsAttr>::value_begin;
diff --git a/mlir/include/mlir/IR/CommonAttrConstraints.td b/mlir/include/mlir/IR/CommonAttrConstraints.td
index ba6cf55a8fb9e..634881f5813f3 100644
--- a/mlir/include/mlir/IR/CommonAttrConstraints.td
+++ b/mlir/include/mlir/IR/CommonAttrConstraints.td
@@ -565,8 +565,8 @@ def StringElementsAttr : ElementsAttrBase<
CPred<"::llvm::isa<::mlir::DenseStringElementsAttr>($_self)" >,
"string elements attribute"> {
- let storageType = [{ ::mlir::DenseElementsAttr }];
- let returnType = [{ ::mlir::DenseElementsAttr }];
+ let storageType = [{ ::mlir::DenseStringElementsAttr }];
+ let returnType = [{ ::mlir::DenseStringElementsAttr }];
let convertFromStorage = "$_self";
}
diff --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp
index dc9744a42b730..17f036b8455f3 100644
--- a/mlir/lib/AsmParser/AttributeParser.cpp
+++ b/mlir/lib/AsmParser/AttributeParser.cpp
@@ -472,8 +472,8 @@ class TensorLiteralParser {
ParseResult parse(bool allowHex);
/// Build a dense attribute instance with the parsed elements and the given
- /// shaped type.
- DenseElementsAttr getAttr(SMLoc loc, ShapedType type);
+ /// shaped type. Returns DenseElementsAttr or DenseStringElementsAttr.
+ Attribute getAttr(SMLoc loc, ShapedType type);
ArrayRef<int64_t> getShape() const { return shape; }
@@ -487,7 +487,7 @@ class TensorLiteralParser {
std::vector<APFloat> &floatValues);
/// Build a Dense String attribute for the given type.
- DenseElementsAttr getStringAttr(SMLoc loc, ShapedType type, Type eltTy);
+ DenseStringElementsAttr getStringAttr(SMLoc loc, ShapedType type, Type eltTy);
/// Build a Dense attribute with hex data for the given type.
DenseElementsAttr getHexAttr(SMLoc loc, ShapedType type);
@@ -539,7 +539,7 @@ ParseResult TensorLiteralParser::parse(bool allowHex) {
/// Build a dense attribute instance with the parsed elements and the given
/// shaped type.
-DenseElementsAttr TensorLiteralParser::getAttr(SMLoc loc, ShapedType type) {
+Attribute TensorLiteralParser::getAttr(SMLoc loc, ShapedType type) {
Type eltType = type.getElementType();
// Check to see if we parse the literal from a hex string.
@@ -679,8 +679,8 @@ TensorLiteralParser::getFloatAttrElements(SMLoc loc, FloatType eltTy,
}
/// Build a Dense String attribute for the given type.
-DenseElementsAttr TensorLiteralParser::getStringAttr(SMLoc loc, ShapedType type,
- Type eltTy) {
+DenseStringElementsAttr
+TensorLiteralParser::getStringAttr(SMLoc loc, ShapedType type, Type eltTy) {
if (hexStorage.has_value()) {
auto stringValue = hexStorage->getStringValue();
return DenseStringElementsAttr::get(type, {stringValue});
@@ -1174,6 +1174,13 @@ Attribute Parser::parseSparseElementsAttr(Type attrType) {
if (!type)
return nullptr;
+ // SparseElementsAttr only supports int/float element types.
+ if (!type.getElementType().isIntOrIndexOrFloat()) {
+ emitError(loc) << "sparse elements attribute does not support string "
+ "element type";
+ return nullptr;
+ }
+
// Construct the sparse elements attr using zero element indice/value
// attributes.
ShapedType indicesType =
@@ -1219,9 +1226,10 @@ Attribute Parser::parseSparseElementsAttr(Type attrType) {
// Otherwise, set the shape to the one parsed by the literal parser.
indicesType = RankedTensorType::get(indiceParser.getShape(), indiceEltType);
}
- auto indices = indiceParser.getAttr(indicesLoc, indicesType);
- if (!indices)
+ auto indicesAttr = indiceParser.getAttr(indicesLoc, indicesType);
+ if (!indicesAttr)
return nullptr;
+ auto indices = llvm::cast<DenseIntElementsAttr>(indicesAttr);
// If the values are a splat, set the shape explicitly based on the number of
// indices. The number of indices is encoded in the first dimension of the
@@ -1231,10 +1239,19 @@ Attribute Parser::parseSparseElementsAttr(Type attrType) {
valuesParser.getShape().empty()
? RankedTensorType::get({indicesType.getDimSize(0)}, valuesEltType)
: RankedTensorType::get(valuesParser.getShape(), valuesEltType);
- auto values = valuesParser.getAttr(valuesLoc, valuesType);
- if (!values)
+ auto valuesAttr = valuesParser.getAttr(valuesLoc, valuesType);
+ if (!valuesAttr)
return nullptr;
+ // SparseElementsAttr only supports DenseElementsAttr for values (not string).
+ auto values = llvm::dyn_cast<DenseElementsAttr>(valuesAttr);
+ if (!values) {
+ emitError(valuesLoc)
+ << "sparse elements attribute requires dense int/float values (string "
+ "element type not supported)";
+ return nullptr;
+ }
+
// Build the sparse elements attribute by the indices and values.
return getChecked<SparseElementsAttr>(loc, type, indices, values);
}
diff --git a/mlir/lib/CAPI/IR/BuiltinAttributes.cpp b/mlir/lib/CAPI/IR/BuiltinAttributes.cpp
index 44a3deaf57db5..7325179c047c5 100644
--- a/mlir/lib/CAPI/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/CAPI/IR/BuiltinAttributes.cpp
@@ -728,8 +728,8 @@ MlirAttribute mlirDenseElementsAttrStringGet(MlirType shapedType,
for (intptr_t i = 0; i < numElements; ++i)
values.push_back(unwrap(strs[i]));
- return wrap(DenseElementsAttr::get(llvm::cast<ShapedType>(unwrap(shapedType)),
- values));
+ return wrap(DenseStringElementsAttr::get(
+ llvm::cast<ShapedType>(unwrap(shapedType)), values));
}
MlirAttribute mlirDenseElementsAttrReshapeGet(MlirAttribute attr,
@@ -743,12 +743,18 @@ MlirAttribute mlirDenseElementsAttrReshapeGet(MlirAttribute attr,
//===----------------------------------------------------------------------===//
bool mlirDenseElementsAttrIsSplat(MlirAttribute attr) {
- return llvm::cast<DenseElementsAttr>(unwrap(attr)).isSplat();
+ Attribute a = unwrap(attr);
+ if (auto strAttr = llvm::dyn_cast<DenseStringElementsAttr>(a))
+ return strAttr.isSplat();
+ return llvm::cast<DenseElementsAttr>(a).isSplat();
}
MlirAttribute mlirDenseElementsAttrGetSplatValue(MlirAttribute attr) {
+ mlir::Attribute a = unwrap(attr);
+ if (auto strAttr = llvm::dyn_cast<DenseStringElementsAttr>(a))
+ return wrap(strAttr.getSplatValue<mlir::Attribute>());
return wrap(
- llvm::cast<DenseElementsAttr>(unwrap(attr)).getSplatValue<Attribute>());
+ llvm::cast<DenseElementsAttr>(a).getSplatValue<mlir::Attribute>());
}
int mlirDenseElementsAttrGetBoolSplatValue(MlirAttribute attr) {
return llvm::cast<DenseElementsAttr>(unwrap(attr)).getSplatValue<bool>();
@@ -778,8 +784,8 @@ double mlirDenseElementsAttrGetDoubleSplatValue(MlirAttribute attr) {
return llvm::cast<DenseElementsAttr>(unwrap(attr)).getSplatValue<double>();
}
MlirStringRef mlirDenseElementsAttrGetStringSplatValue(MlirAttribute attr) {
- return wrap(
- llvm::cast<DenseElementsAttr>(unwrap(attr)).getSplatValue<StringRef>());
+ return wrap(llvm::cast<DenseStringElementsAttr>(unwrap(attr))
+ .getSplatValue<llvm::StringRef>());
}
//===----------------------------------------------------------------------===//
@@ -824,8 +830,8 @@ double mlirDenseElementsAttrGetDoubleValue(MlirAttribute attr, intptr_t pos) {
}
MlirStringRef mlirDenseElementsAttrGetStringValue(MlirAttribute attr,
intptr_t pos) {
- return wrap(
- llvm::cast<DenseElementsAttr>(unwrap(attr)).getValues<StringRef>()[pos]);
+ return wrap(llvm::cast<DenseStringElementsAttr>(unwrap(attr))
+ .getValues<StringRef>()[pos]);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index bbbc9198a68ab..e288be3271fab 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/IR/BuiltinAttributes.cpp
@@ -589,23 +589,14 @@ static bool hasSameNumElementsOrSplat(ShapedType type, const Values &values) {
//===----------------------------------------------------------------------===//
DenseElementsAttr::AttributeElementIterator::AttributeElementIterator(
- DenseElementsAttr attr, size_t index)
+ Attribute attr, size_t index)
: llvm::indexed_accessor_iterator<AttributeElementIterator, const void *,
Attribute, Attribute, Attribute>(
attr.getAsOpaquePointer(), index) {}
Attribute DenseElementsAttr::AttributeElementIterator::operator*() const {
auto owner = llvm::cast<DenseElementsAttr>(getFromOpaquePointer(base));
- Type eltTy = owner.getElementType();
-
- // Handle strings specially.
- if (llvm::isa<DenseStringElementsAttr>(owner)) {
- ArrayRef<StringRef> vals = owner.getRawStringData();
- return StringAttr::get(owner.isSplat() ? vals.front() : vals[index], eltTy);
- }
-
- // All other types should implement DenseElementTypeInterface.
- auto denseEltTy = llvm::cast<DenseElementType>(eltTy);
+ auto denseEltTy = llvm::cast<DenseElementType>(owner.getElementType());
ArrayRef<char> rawData = owner.getRawData();
// Storage is byte-aligned: align bit size up to next byte boundary.
size_t bitSize = denseEltTy.getDenseElementBitSize();
@@ -864,28 +855,13 @@ template class DenseArrayAttrImpl<double>;
/// Method for support type inquiry through isa, cast and dyn_cast.
bool DenseElementsAttr::classof(Attribute attr) {
- return llvm::isa<DenseIntOrFPElementsAttr, DenseStringElementsAttr>(attr);
+ return llvm::isa<DenseIntOrFPElementsAttr>(attr);
}
DenseElementsAttr DenseElementsAttr::get(ShapedType type,
ArrayRef<Attribute> values) {
assert(hasSameNumElementsOrSplat(type, values));
- Type eltType = type.getElementType();
-
- // Handle strings specially.
- if (!llvm::isa<DenseElementType>(eltType)) {
- SmallVector<StringRef, 8> stringValues;
- stringValues.reserve(values.size());
- for (Attribute attr : values) {
- assert(llvm::isa<StringAttr>(attr) &&
- "expected string value for non-DenseElementType element");
- stringValues.push_back(llvm::cast<StringAttr>(attr).getValue());
- }
- return get(type, stringValues);
- }
-
- // All other types go through DenseElementTypeInterface.
- auto denseEltType = llvm::dyn_cast<DenseElementType>(eltType);
+ auto denseEltType = llvm::dyn_cast<DenseElementType>(type.getElementType());
assert(denseEltType &&
"attempted to get DenseElementsAttr with unsupported element type");
SmallVector<char> data;
@@ -906,12 +882,6 @@ DenseElementsAttr DenseElementsAttr::get(ShapedType type,
values.size()));
}
-DenseElementsAttr DenseElementsAttr::get(ShapedType type,
- ArrayRef<StringRef> values) {
- assert(!type.getElementType().isIntOrFloat());
- return DenseStringElementsAttr::get(type, values);
-}
-
/// Constructs a dense integer elements attribute from an array of APInt
/// values. Each APInt value is expected to have the same bitwidth as the
/// element type of 'type'.
@@ -1048,9 +1018,6 @@ bool DenseElementsAttr::isValidComplex(int64_t dataEltSize, bool isInt,
/// values are the same.
bool DenseElementsAttr::isSplat() const {
// Splat iff the data array has exactly one element.
- if (isa<DenseStringElementsAttr>(*this))
- return getRawStringData().size() == 1;
- // FP/Int case.
size_t storageSize = llvm::divideCeil(
getDenseElementBitWidth(getType().getElementType()), CHAR_BIT);
return getRawData().size() == storageSize;
@@ -1100,10 +1067,6 @@ ArrayRef<char> DenseElementsAttr::getRawData() const {
return static_cast<DenseIntOrFPElementsAttrStorage *>(impl)->data;
}
-ArrayRef<StringRef> DenseElementsAttr::getRawStringData() const {
- return static_cast<DenseStringElementsAttrStorage *>(impl)->data;
-}
-
/// Return a new DenseElementsAttr that has the same data as the current
/// attribute, but has been reshaped to 'newType'. The new type must have the
/// same total number of elements as well as element type.
@@ -1390,6 +1353,27 @@ bool DenseIntElementsAttr::classof(Attribute attr) {
return false;
}
+//===----------------------------------------------------------------------===//
+// DenseStringElementsAttr
+//===----------------------------------------------------------------------===//
+
+ShapedType DenseStringElementsAttr::getType() const {
+ return static_cast<const DenseStringElementsAttrStorage *>(impl)->type;
+}
+
+ArrayRef<StringRef> DenseStringElementsAttr::getRawStringData() const {
+ return static_cast<const DenseStringElementsAttrStorage *>(impl)->data;
+}
+
+Attribute
+DenseStringElementsAttr::StringAttributeElementIterator::operator*() const {
+ auto attr = llvm::cast<DenseStringElementsAttr>(
+ Attribute::getFromOpaquePointer(this->base));
+ auto data = attr.getRawStringData();
+ return StringAttr::get(attr.isSplat() ? data.front() : data[this->index],
+ attr.getElementType());
+}
+
//===----------------------------------------------------------------------===//
// DenseResourceElementsAttr
//===----------------------------------------------------------------------===//
@@ -1557,10 +1541,6 @@ Attribute SparseElementsAttr::getZeroAttr() const {
ArrayRef<Attribute>{zero, zero});
}
- // Handle string type.
- if (llvm::isa<DenseStringElementsAttr>(getValues()))
- return StringAttr::get("", eltType);
-
// Otherwise, this is an integer.
return IntegerAttr::get(eltType, 0);
}
diff --git a/mlir/test/IR/parser.mlir b/mlir/test/IR/parser.mlir
index 3bb6e38b4d613..7972d1da88e83 100644
--- a/mlir/test/IR/parser.mlir
+++ b/mlir/test/IR/parser.mlir
@@ -797,10 +797,8 @@ func.func @sparsetensorattr() -> () {
// CHECK: "foof321"() {bar = sparse<> : tensor<f32>} : () -> ()
"foof321"(){bar = sparse<> : tensor<f32>} : () -> ()
-// CHECK: "foostr"() {bar = sparse<0, "foo"> : tensor<1x1x1x!unknown<>>} : () -> ()
- "foostr"(){bar = sparse<0, "foo"> : tensor<1x1x1x!unknown<>>} : () -> ()
-// CHECK: "foostr"() {bar = sparse<{{\[\[}}1, 1, 0], {{\[}}0, 1, 0], {{\[}}0, 0, 1]], {{\[}}"a", "b", "c"]> : tensor<2x2x2x!unknown<>>} : () -> ()
- "foostr"(){bar = sparse<[[1, 1, 0], [0, 1, 0], [0, 0, 1]], ["a", "b", "c"]> : tensor<2x2x2x!unknown<>>} : () -> ()
+// Sparse elements with string element type is not supported (dense string uses
+// DenseStringElementsAttr). Ops foostr with sparse string are omitted.
return
}
diff --git a/mlir/test/mlir-tblgen/openmp-clause-ops.td b/mlir/test/mlir-tblgen/openmp-clause-ops.td
index 3e5896a00182b..c502b21c3baf8 100644
--- a/mlir/test/mlir-tblgen/openmp-clause-ops.td
+++ b/mlir/test/mlir-tblgen/openmp-clause-ops.td
@@ -59,7 +59,7 @@ def OpenMP_MyFirstClause : OpenMP_Clause<
// CHECK-NEXT: ::mlir::IntegerAttr complexOptIntAttr;
// CHECK-NEXT: ::mlir::ElementsAttr elementsAttr;
-// CHECK-NEXT: ::mlir::DenseElementsAttr stringElementsAttr;
+// CHECK-NEXT: ::mlir::DenseStringElementsAttr stringElementsAttr;
// CHECK-NEXT: }
def OpenMP_MySecondClause : OpenMP_Clause<
diff --git a/mlir/unittests/IR/AttributeTest.cpp b/mlir/unittests/IR/AttributeTest.cpp
index 404aa8c0dcf3d..79ca28908ccca 100644
--- a/mlir/unittests/IR/AttributeTest.cpp
+++ b/mlir/unittests/IR/AttributeTest.cpp
@@ -38,6 +38,21 @@ static void testSplat(Type eltType, const EltTy &splatElt) {
EXPECT_TRUE(newValue == splatElt);
}
+template <>
+void testSplat<StringRef>(Type eltType, const StringRef &splatElt) {
+ RankedTensorType shape = RankedTensorType::get({2, 1}, eltType);
+
+ DenseStringElementsAttr splat = DenseStringElementsAttr::get(shape, splatElt);
+ EXPECT_TRUE(splat.isSplat());
+
+ auto detectedSplat =
+ DenseStringElementsAttr::get(shape, llvm::ArrayRef({splatElt, splatElt}));
+ EXPECT_EQ(detectedSplat, splat);
+
+ for (auto newValue : detectedSplat.getValues<StringRef>())
+ EXPECT_TRUE(newValue == splatElt);
+}
+
namespace {
TEST(DenseSplatTest, BoolSplat) {
MLIRContext context;
@@ -184,8 +199,16 @@ TEST(DenseSplatTest, StringAttrSplat) {
context.allowUnregisteredDialects();
Type stringType =
OpaqueType::get(StringAttr::get(&context, "test"), "string");
+ RankedTensorType shape = RankedTensorType::get({2, 1}, stringType);
Attribute stringAttr = StringAttr::get("test-string", stringType);
- testSplat(stringType, stringAttr);
+ StringRef value = llvm::cast<StringAttr>(stringAttr).getValue();
+ DenseStringElementsAttr splat = DenseStringElementsAttr::get(shape, value);
+ EXPECT_TRUE(splat.isSplat());
+ auto detectedSplat =
+ DenseStringElementsAttr::get(shape, llvm::ArrayRef({value, value}));
+ EXPECT_EQ(detectedSplat, splat);
+ for (auto newValue : detectedSplat.getValues<StringRef>())
+ EXPECT_TRUE(newValue == value);
}
TEST(DenseComplexTest, ComplexFloatSplat) {
@@ -396,11 +419,9 @@ TEST(SparseElementsAttrTest, GetZero) {
IntegerType intTy = IntegerType::get(&context, 32);
FloatType floatTy = Float32Type::get(&context);
- Type stringTy = OpaqueType::get(StringAttr::get(&context, "test"), "string");
ShapedType tensorI32 = RankedTensorType::get({2, 2}, intTy);
ShapedType tensorF32 = RankedTensorType::get({2, 2}, floatTy);
- ShapedType tensorString = RankedTensorType::get({2, 2}, stringTy);
auto indicesType =
RankedTensorType::get({1, 2}, IntegerType::get(&context, 64));
@@ -413,13 +434,8 @@ TEST(SparseElementsAttrTest, GetZero) {
RankedTensorType floatValueTy = RankedTensorType::get({1}, floatTy);
auto floatValue = DenseFPElementsAttr::get(floatValueTy, {1.0f});
- RankedTensorType stringValueTy = RankedTensorType::get({1}, stringTy);
- auto stringValue = DenseElementsAttr::get(stringValueTy, {StringRef("foo")});
-
auto sparseInt = SparseElementsAttr::get(tensorI32, indices, intValue);
auto sparseFloat = SparseElementsAttr::get(tensorF32, indices, floatValue);
- auto sparseString =
- SparseElementsAttr::get(tensorString, indices, stringValue);
// Only index (0, 0) contains an element, others are supposed to return
// the zero/empty value.
@@ -433,10 +449,9 @@ TEST(SparseElementsAttrTest, GetZero) {
EXPECT_EQ(zeroFloatValue.getValueAsDouble(), 0.0f);
EXPECT_TRUE(zeroFloatValue.getType() == floatTy);
- auto zeroStringValue =
- cast<StringAttr>(sparseString.getValues<Attribute>()[{1, 1}]);
- EXPECT_TRUE(zeroStringValue.empty());
- EXPECT_TRUE(zeroStringValue.getType() == stringTy);
+ // Note: SparseElementsAttr does not support string element type (values must
+ // be DenseElementsAttr). Use DenseStringElementsAttr for dense string
+ // tensors.
}
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list