[Mlir-commits] [mlir] [mlir][WIP] `DenseElementsAttr` generalized (PR #179122)
Matthias Springer
llvmlistbot at llvm.org
Sun Feb 1 09:42:23 PST 2026
https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/179122
None
>From e1076a82ab0ba24a33aafddb4456a3fbb7112b24 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] [mlir][WIP] `DenseElementsAttr` generalized
---
mlir/include/mlir/IR/BuiltinTypeInterfaces.td | 76 +++++++++
mlir/include/mlir/IR/BuiltinTypes.td | 7 +
mlir/lib/AsmParser/AttributeParser.cpp | 156 +++++++++++++++++-
mlir/lib/AsmParser/Parser.h | 5 +
mlir/lib/IR/AsmPrinter.cpp | 72 ++++++--
mlir/lib/IR/AttributeDetail.h | 4 +
mlir/lib/IR/BuiltinTypes.cpp | 2 +
.../IR/dense-elements-type-interface.mlir | 28 ++++
mlir/test/lib/Dialect/Test/TestTypeDefs.td | 16 ++
mlir/test/lib/Dialect/Test/TestTypes.cpp | 28 ++++
mlir/test/lib/Dialect/Test/TestTypes.h | 1 +
11 files changed, 379 insertions(+), 16 deletions(-)
create mode 100644 mlir/test/IR/dense-elements-type-interface.mlir
diff --git a/mlir/include/mlir/IR/BuiltinTypeInterfaces.td b/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
index 9ef08b7020b99..8e75f5da52de2 100644
--- a/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
+++ b/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
@@ -110,6 +110,25 @@ def MemRefElementTypeInterface : TypeInterface<"MemRefElementTypeInterface"> {
}];
}
+//===----------------------------------------------------------------------===//
+// TensorElementTypeInterface
+//===----------------------------------------------------------------------===//
+
+def TensorElementTypeInterface : TypeInterface<"TensorElementTypeInterface"> {
+ let cppNamespace = "::mlir";
+ let description = [{
+ Indication that this type can be used as element in tensor types.
+
+ Implementing this interface establishes a contract between this type and
+ tensor types (RankedTensorType, UnrankedTensorType), indicating that this
+ type can be used as an element type of tensors.
+
+ The interface currently has no methods and is used by types to opt into
+ being tensor elements. This may change in the future, in particular to
+ require types to provide their size given a data layout.
+ }];
+}
+
//===----------------------------------------------------------------------===//
// PtrLikeTypeInterface
//===----------------------------------------------------------------------===//
@@ -338,4 +357,61 @@ def ShapedTypeInterface : TypeInterface<"ShapedType"> {
}];
}
+//===----------------------------------------------------------------------===//
+// 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 byte size for element storage
+ 2. How to convert between raw bytes and the corresponding Attribute
+
+ For example, a custom `!my_dialect.fixed_point` type might map to a
+ `my_dialect.fixed_point_attr` attribute. The interface methods handle
+ conversion between the attribute and raw storage bytes.
+
+ When this interface is implemented by an element type, DenseElementsAttr
+ uses the new type-first syntax:
+ `dense<tensor<2x!my.type> : [#my.val<1>, #my.val<2>]>`
+ instead of the traditional literal-first syntax:
+ `dense<[1, 2, 3]> : tensor<3xi32>`
+ }];
+
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the number of bytes required to store one element in dense
+ storage. This must be a compile-time constant for the type.
+ }],
+ /*retTy=*/"size_t",
+ /*methodName=*/"getDenseElementByteSize",
+ /*args=*/(ins)
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Convert raw storage bytes to an attribute representing this element
+ value. The `rawData` array contains exactly `getDenseElementByteSize()`
+ bytes.
+ }],
+ /*retTy=*/"::mlir::Attribute",
+ /*methodName=*/"convertToAttribute",
+ /*args=*/(ins "::llvm::ArrayRef<char>":$rawData)
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Convert an attribute to raw storage bytes. Appends exactly
+ `getDenseElementByteSize()` bytes to `result`. Returns failure if the
+ attribute is incompatible with this element type.
+ }],
+ /*retTy=*/"::llvm::LogicalResult",
+ /*methodName=*/"convertFromAttribute",
+ /*args=*/(ins "::mlir::Attribute":$attr,
+ "::llvm::SmallVectorImpl<char>&":$result)
+ >,
+ ];
+}
+
#endif // MLIR_IR_BUILTINTYPEINTERFACES_TD_
diff --git a/mlir/include/mlir/IR/BuiltinTypes.td b/mlir/include/mlir/IR/BuiltinTypes.td
index 08847dd11c685..5ae5493982538 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.td
+++ b/mlir/include/mlir/IR/BuiltinTypes.td
@@ -1294,6 +1294,13 @@ def Builtin_UnrankedTensor : Builtin_Type<"UnrankedTensor", "unranked_tensor", [
// VectorType
//===----------------------------------------------------------------------===//
+// TensorElementTypeInterface-based constraint for tensor element types.
+// This allows custom types implementing TensorElementTypeInterface to be used
+// as tensor elements.
+def Builtin_TensorTypeElementType : AnyTypeOf<[TensorElementTypeInterface]> {
+ let cppFunctionName = "isValidTensorTypeElementType";
+}
+
// Note: VectorType uses this type constraint instead of a plain
// VectorElementTypeInterface, so that methods with mlir::Type are generated.
// We may want to drop this in future and require VectorElementTypeInterface
diff --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp
index 374471fd3ed41..bd70df0a232c1 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"
@@ -954,6 +955,42 @@ Attribute Parser::parseDenseArrayAttr(Type attrType) {
return eltParser.getAttr();
}
+/// Check if the current token can start a type (for type-first dense syntax).
+static bool canStartType(Token tok) {
+ switch (tok.getKind()) {
+ case Token::kw_memref:
+ case Token::kw_tensor:
+ case Token::kw_complex:
+ case Token::kw_tuple:
+ case Token::kw_vector:
+ case Token::inttype:
+ case Token::kw_f4E2M1FN:
+ case Token::kw_f6E2M3FN:
+ case Token::kw_f6E3M2FN:
+ case Token::kw_f8E5M2:
+ case Token::kw_f8E4M3:
+ case Token::kw_f8E4M3FN:
+ case Token::kw_f8E5M2FNUZ:
+ case Token::kw_f8E4M3FNUZ:
+ case Token::kw_f8E4M3B11FNUZ:
+ case Token::kw_f8E3M4:
+ case Token::kw_f8E8M0FNU:
+ case Token::kw_bf16:
+ case Token::kw_f16:
+ case Token::kw_tf32:
+ case Token::kw_f32:
+ case Token::kw_f64:
+ case Token::kw_f80:
+ case Token::kw_f128:
+ case Token::kw_index:
+ case Token::kw_none:
+ case Token::exclamation_identifier:
+ return true;
+ default:
+ return false;
+ }
+}
+
/// Parse a dense elements attribute.
Attribute Parser::parseDenseElementsAttr(Type attrType) {
auto attribLoc = getToken().getLoc();
@@ -961,7 +998,12 @@ Attribute Parser::parseDenseElementsAttr(Type attrType) {
if (parseToken(Token::less, "expected '<' after 'dense'"))
return nullptr;
- // Parse the literal data if necessary.
+ // Check for the new type-first syntax: dense<TYPE : [ATTR, ...]>
+ // This is used for element types implementing DenseElementTypeInterface.
+ if (canStartType(getToken()))
+ return parseDenseElementsAttrTyped(attribLoc);
+
+ // Parse the literal data if necessary (old syntax: dense<LITERAL> : TYPE).
TensorLiteralParser literalParser(*this);
if (!consumeIf(Token::greater)) {
if (literalParser.parse(/*allowHex=*/true) ||
@@ -975,6 +1017,118 @@ Attribute Parser::parseDenseElementsAttr(Type attrType) {
return literalParser.getAttr(attribLoc, type);
}
+/// Parse a dense elements attribute with the type-first syntax.
+/// Syntax: dense<TYPE : [ATTR, ATTR, ...]>
+/// This is used for element types implementing DenseElementTypeInterface.
+Attribute Parser::parseDenseElementsAttrTyped(SMLoc loc) {
+ // Parse the shaped type.
+ Type type = parseType();
+ if (!type)
+ return nullptr;
+
+ auto shapedType = dyn_cast<ShapedType>(type);
+ if (!shapedType) {
+ emitError(loc, "expected a shaped type for dense elements");
+ return nullptr;
+ }
+
+ if (!shapedType.hasStaticShape()) {
+ emitError(loc, "dense elements type must have static shape");
+ return nullptr;
+ }
+
+ // Check that the element type implements DenseElementTypeInterface.
+ Type eltType = shapedType.getElementType();
+ auto denseEltType = dyn_cast<DenseElementType>(eltType);
+ if (!denseEltType) {
+ emitError(loc, "element type must implement DenseElementTypeInterface for "
+ "type-first dense syntax");
+ return nullptr;
+ }
+
+ if (parseToken(Token::colon, "expected ':' after type in dense attribute"))
+ return nullptr;
+
+ // Parse the element attributes and convert to raw bytes.
+ SmallVector<char> rawData;
+ size_t byteSize = denseEltType.getDenseElementByteSize();
+ int64_t numElements = shapedType.getNumElements();
+ SmallVector<int64_t, 4> shape;
+
+ // Helper to parse a single element or nested list.
+ std::function<ParseResult(SmallVectorImpl<int64_t> &)> parseElements;
+ parseElements = [&](SmallVectorImpl<int64_t> &dims) -> ParseResult {
+ // Check for nested list.
+ if (getToken().is(Token::l_square)) {
+ SmallVector<int64_t, 4> prevDims;
+ bool first = true;
+ unsigned size = 0;
+ auto parseOneElement = [&]() -> ParseResult {
+ SmallVector<int64_t, 4> thisDims;
+ if (parseElements(thisDims))
+ return failure();
+ ++size;
+ if (!first && prevDims != thisDims) {
+ emitError("tensor literal has inconsistent ranks between elements");
+ return failure();
+ }
+ prevDims = thisDims;
+ first = false;
+ return success();
+ };
+ if (parseCommaSeparatedList(Delimiter::Square, parseOneElement))
+ return failure();
+ dims.clear();
+ dims.push_back(size);
+ dims.append(prevDims.begin(), prevDims.end());
+ return success();
+ }
+
+ // Parse a single element attribute.
+ Attribute elemAttr = parseAttribute();
+ if (!elemAttr)
+ return failure();
+
+ // Convert attribute to raw bytes using the interface.
+ if (failed(denseEltType.convertFromAttribute(elemAttr, rawData))) {
+ emitError("incompatible attribute for element type");
+ return failure();
+ }
+
+ // Single element has empty dims (will be combined by parent).
+ dims.clear();
+ return success();
+ };
+
+ // Parse the elements.
+ if (parseElements(shape))
+ return nullptr;
+
+ // Verify the shape matches (if non-empty shape was inferred).
+ if (!shape.empty() && shape != ArrayRef(shapedType.getShape())) {
+ emitError(loc) << "inferred shape of elements literal ([" << shape
+ << "]) does not match type ([" << shapedType.getShape()
+ << "])";
+ return nullptr;
+ }
+
+ // Handle splat detection: if only one element was parsed, it's a splat.
+ bool isSplat = (rawData.size() == byteSize && numElements != 1);
+ if (isSplat) {
+ // For splat, we only store one element.
+ } else if (rawData.size() != byteSize * numElements) {
+ emitError(loc) << "parsed " << (rawData.size() / byteSize)
+ << " elements, but type expects " << numElements;
+ return nullptr;
+ }
+
+ if (parseToken(Token::greater, "expected '>' to close dense attribute"))
+ return nullptr;
+
+ // Create the attribute from raw buffer.
+ return DenseElementsAttr::getFromRawBuffer(shapedType, rawData);
+}
+
Attribute Parser::parseDenseResourceElementsAttr(Type attrType) {
auto loc = getToken().getLoc();
consumeToken(Token::kw_dense_resource);
diff --git a/mlir/lib/AsmParser/Parser.h b/mlir/lib/AsmParser/Parser.h
index ecc128cf767b3..8c6434c99b28c 100644
--- a/mlir/lib/AsmParser/Parser.h
+++ b/mlir/lib/AsmParser/Parser.h
@@ -290,6 +290,11 @@ class Parser {
Attribute parseDenseElementsAttr(Type attrType);
ShapedType parseElementsLiteralType(SMLoc loc, Type type);
+ /// Parse a dense elements attribute with type-first syntax.
+ /// This is used for element types that implement DenseElementTypeInterface.
+ /// Syntax: dense<TYPE : [ATTR, ATTR, ...]>
+ Attribute parseDenseElementsAttrTyped(SMLoc loc);
+
/// Parse a dense resource elements attribute.
Attribute parseDenseResourceElementsAttr(Type attrType);
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 81455699421cc..1733f42c6b333 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -512,6 +512,11 @@ class AsmPrinter::Impl {
void printDenseIntOrFPElementsAttr(DenseIntOrFPElementsAttr attr,
bool allowHex);
+ /// Print a dense elements attribute using DenseElementTypeInterface.
+ /// Uses the type-first syntax: dense<TYPE : [ATTR, ...]>
+ void printDenseElementsAttrWithInterface(DenseElementsAttr attr,
+ DenseElementType denseEltType);
+
/// Print a dense array attribute.
void printDenseArrayAttr(DenseArrayAttr attr);
@@ -2501,23 +2506,41 @@ void AsmPrinter::Impl::printAttributeImpl(Attribute attr,
printSymbolReference(nestedRef.getValue(), os);
}
- } else if (auto intOrFpEltAttr =
- llvm::dyn_cast<DenseIntOrFPElementsAttr>(attr)) {
- if (printerFlags.shouldElideElementsAttr(intOrFpEltAttr)) {
- printElidedElementsAttr(os);
- } else {
- os << "dense<";
- printDenseIntOrFPElementsAttr(intOrFpEltAttr, /*allowHex=*/true);
- os << '>';
+ } else if (auto denseEltAttr = llvm::dyn_cast<DenseElementsAttr>(attr)) {
+ // Check if the element type implements DenseElementTypeInterface.
+ // If so, use the type-first syntax which embeds the type in the attribute.
+ Type eltType = denseEltAttr.getElementType();
+ if (auto denseEltType = llvm::dyn_cast<DenseElementType>(eltType)) {
+ if (printerFlags.shouldElideElementsAttr(denseEltAttr)) {
+ printElidedElementsAttr(os);
+ } else {
+ os << "dense<";
+ printDenseElementsAttrWithInterface(denseEltAttr, denseEltType);
+ os << '>';
+ }
+ // Type is embedded in the syntax, don't print it again.
+ return;
}
- } else if (auto strEltAttr = llvm::dyn_cast<DenseStringElementsAttr>(attr)) {
- if (printerFlags.shouldElideElementsAttr(strEltAttr)) {
- printElidedElementsAttr(os);
- } else {
- os << "dense<";
- printDenseStringElementsAttr(strEltAttr);
- os << '>';
+ // Fall back to existing printing for built-in element types.
+ if (auto intOrFpEltAttr =
+ llvm::dyn_cast<DenseIntOrFPElementsAttr>(denseEltAttr)) {
+ if (printerFlags.shouldElideElementsAttr(intOrFpEltAttr)) {
+ printElidedElementsAttr(os);
+ } else {
+ os << "dense<";
+ printDenseIntOrFPElementsAttr(intOrFpEltAttr, /*allowHex=*/true);
+ os << '>';
+ }
+ } else if (auto strEltAttr =
+ llvm::dyn_cast<DenseStringElementsAttr>(denseEltAttr)) {
+ if (printerFlags.shouldElideElementsAttr(strEltAttr)) {
+ printElidedElementsAttr(os);
+ } else {
+ os << "dense<";
+ printDenseStringElementsAttr(strEltAttr);
+ os << '>';
+ }
}
} else if (auto sparseEltAttr = llvm::dyn_cast<SparseElementsAttr>(attr)) {
@@ -2705,6 +2728,25 @@ void AsmPrinter::Impl::printDenseStringElementsAttr(
printDenseElementsAttrImpl(attr.isSplat(), attr.getType(), os, printFn);
}
+void AsmPrinter::Impl::printDenseElementsAttrWithInterface(
+ DenseElementsAttr attr, DenseElementType denseEltType) {
+ // Print the type first: dense<TYPE : [ELEMENTS]>
+ printType(attr.getType());
+ os << " : ";
+
+ ArrayRef<char> rawData = attr.getRawData();
+ size_t byteSize = denseEltType.getDenseElementByteSize();
+
+ // 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 cb9d21bf3e611..7a502c2699ff6 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,6 +33,9 @@ namespace detail {
/// Return the bit width which DenseElementsAttr should use for this type.
inline size_t getDenseElementBitWidth(Type eltType) {
+ // Check for DenseElementTypeInterface first.
+ if (auto denseEltType = llvm::dyn_cast<DenseElementType>(eltType))
+ return denseEltType.getDenseElementByteSize() * CHAR_BIT;
// 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;
diff --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp
index 1e198043c590a..cd60544e7a4f1 100644
--- a/mlir/lib/IR/BuiltinTypes.cpp
+++ b/mlir/lib/IR/BuiltinTypes.cpp
@@ -12,6 +12,7 @@
#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"
@@ -349,6 +350,7 @@ bool TensorType::isValidElementType(Type type) {
// element type within that dialect.
return llvm::isa<ComplexType, FloatType, IntegerType, OpaqueType, VectorType,
IndexType>(type) ||
+ llvm::isa<TensorElementTypeInterface>(type) ||
!llvm::isa<BuiltinDialect>(type.getDialect());
}
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..8aa1386cfeb73
--- /dev/null
+++ b/mlir/test/IR/dense-elements-type-interface.mlir
@@ -0,0 +1,28 @@
+// RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect | 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() {
+ // The type is embedded in the dense attribute syntax, not printed separately.
+ // CHECK: "unregistered_op"() {attr = dense<tensor<3x!test.dense_element> : [1 : i32, 2 : i32, 3 : i32]>}
+ "unregistered_op"() {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: "unregistered_op"() {attr = dense<tensor<2x2x!test.dense_element> : {{\[}}{{\[}}1 : i32, 2 : i32], [3 : i32, 4 : i32]]>}
+ "unregistered_op"() {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() {
+ // A splat should be detected and stored efficiently
+ // CHECK: "unregistered_op"() {attr = dense<tensor<4x!test.dense_element> : 42 : i32>}
+ "unregistered_op"() {attr = dense<tensor<4x!test.dense_element> : 42 : i32>} : () -> ()
+ return
+}
diff --git a/mlir/test/lib/Dialect/Test/TestTypeDefs.td b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
index d840339a09266..244e55965a009 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"
@@ -482,4 +483,19 @@ def TestTypeNewlineAndIndent : Test_Type<"TestTypeNewlineAndIndent"> {
let hasCustomAssemblyFormat = 1;
}
+//===----------------------------------------------------------------------===//
+// Test type for DenseElementTypeInterface
+//===----------------------------------------------------------------------===//
+
+def TestTypeDenseElement : Test_Type<"TestDenseElement",
+ [DeclareTypeInterfaceMethods<DenseElementTypeInterface>,
+ TensorElementTypeInterface]> {
+ 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.
+ Also implements TensorElementTypeInterface to be usable as tensor elements.
+ }];
+}
+
#endif // TEST_TYPEDEFS
diff --git a/mlir/test/lib/Dialect/Test/TestTypes.cpp b/mlir/test/lib/Dialect/Test/TestTypes.cpp
index 71dd25b0093e0..b6da167e48787 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::getDenseElementByteSize() const { return 4; }
+
+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"
More information about the Mlir-commits
mailing list