[RFC] Allow Custom Element Types in `DenseElementAttr`

MLIR
1

tldr: DenseElementsAttr is currently hard-coded to certain builtin types. Make DenseElementsAttr extensible, so that it can be used with custom element types.

Background

DenseElementsAttr (e.g., dense<[1, 2, 3, 4]>: tensor<4xi32>) is an attribute that stores elements of shaped values in a dense storage format: instead of storing one Attribute per element, it stores the data in an ArrayRef<char>. This is more efficient for large shaped values.

DenseElementsAttr is currently hard-coded to IntegerType, IndexType, FloatType, ComplexType and strings. (Strings are a bit special, let’s not get into that…)

Here are some examples:

Proposal

Users have expressed interest in a more extensible DenseElementsAttr infrastructure. Most recently on this PR. This RFC proposes to:

  1. Add a new TensorElementTypeInterface, so that custom tensor element types are possible. This interface is similar to the VectorElementTypeInterface.
  2. Add a new DenseElementTypeInterface for element types that are supported by DenseElementsAttr. The interface has the following interface methods:
    2.1. size_t getDenseElementByteSize(): Return the storage size of a single element. For verification purposes and to compute where a certain element begins and ends (needed for iterators + random access).
    2.2. Attribute convertToAttribute(ArrayRef<char>): Deserialize the bytes of a single element into an MLIR attribute.
    2.3. LogicalResult convertFromAttribute(Attribute, SmallVectorImpl<char> &): Serialize an MLIR attribute into bytes.
  3. Update the assembly format of DenseElementsAttr: the shaped type should come first, so that you know what you’re about to parse. (The existing format can also remain in place.)

The conversion functions from/to MLIR attributes allow users to iterate over a DenseElementsAttr in terms of Attribute. For efficiency reasons, there should also be a way to iterate over elements in a “dense” way, i.e., without constructing an Attribute for each scalar. This is not shown here, but in essence, users will receive a void * / char * for each element and can then interpret the data as needed.

Another benefit of the conversion from/to attributes is that the existing parser/printer can be utilized for the assembly format of the DenseElementsAttr. Note: This assumes that there is a corresponding attribute for each supported element type. E.g., IntegerType -> IntegerAttr or MyCustomIntegerType -> IntegerAttr. (Two different types can share the same attribute class.)

Assembly Format

Current assembly format: dense<[[1, 3], [0, 7]]> : tensor<2x2xi32>

The existing assembly format does not work for custom DenseElementsAttr element types. In order to parse an element, we first have to know what we are going to parse.

Custom element types must be written in this format: dense<tensor<2x2x!test.my_int_type> : [[1 : i32, 2 : i32], [3 : i32, 4 : i32]]>. Here’s a breakdown of what happens when parsing this attribute:

  1. The ranked tensor type with element type !test.my_int_type is parsed.
  2. 2x2 element attributes are parsed. These could be arbitrary attributes.
  3. For each element attribute, MyIntType::convertFromAttribute is called. This interface function expects IntegerAttr and will fail when other attributes are passed, triggering a parser error. The serialized bytes are stored in the DenseElementsAttr.

Both the current assembly format and the new assembly format can exist side-by-side. However, for consistency reasons, it would be great if we could eventually deprecated the current assembly format.

Future Extensions

Parsing/printing is a bit inefficient because it materializes temporary attributes for each element. This could be avoided by adding additional interface methods for parsing/printing the entire dense elements attr data. This would also allow for denser notations such as dense<tensor<2x2x!test.my_int_type> : [[1, 2], [3, 4]]>.

Note that while the “parsing/printing via scalar attribute” approach as presented here is inefficient, it does not slow down attribute creation / access via the C++ API.

Prototype

Here’s a Cursor-generated prototype: [mlir][IR] Generalize`DenseElementsAttr` to custom element types by matthias-springer · Pull Request #179122 · llvm/llvm-project · GitHub. I’d like to gather some feedback before cleaning up that PR and splitting it up into smaller pieces.

2

Can’t this be done by just moving the current dense parser to the appropriate types?

How would it work with the hex syntax dense<"0x...">?

There’s a cost to this during debugging, development and prototyping that makes some things prohibitive, so parser/printer is not a small problem.

3

I very much want this! It will also have downstream benefits like simplifying constant folding for elementwise ops with custom types.

4

I feel like these are related. Printing everything as hex as soon as you run over a certain amount of elements seems like the simplest way to me, and is already done by the existing attribute.

In fact, this was all so reasonable to me that I did this 3 years ago, but with an ad-hoc attribute exposing the ElementsAttrInterface, as there is almost no boilerplate involved to get this functionality for a known type:

While I support this proposal, I feel like it is quite complicated (with the new interfaces and all), considering users desiring this functionality can already get it with relative ease. Custom data types are likely semantically mismatched with upstream consumers of the dense attribute, so I never felt that rolling your own gets in your way here. E.g., arith.constant isn’t going to know how to materialize !dialect.my_weird_int into LLVM, so why bother?

However, I’m always for new interfaces, as long as IRDL will let me finally implement them :smiley:

5

I’m also a bit confused by this statement. You have all the machinery you need. Do you mean, it can’t be used to directly yield a new kind of Attribute?

Also 3 years ago, I made this iterator that is used by an attribute that wraps two arbitrary ElementsAttr, one of i1, one of anything else T, to iterate over std::optional<T>:

The corresponding attribute definition has some more experimental features to delegate constant materialization to different dialects. However, it still allowed me to get the desired functionality for arbitrary T quite easily while adapting any (built-in) aggregate attribute:

6

In the general case, you cannot parse elements without knowing what you’re going to parse. That’s why the current syntax (dense<[1, 2, 3, 4]> : tensor<4xi32>) no longer works. This syntax works only as long as you’re parsing integers and floats.

We can move the int/float dense parsing/printing to the appropriate types by doing what I mentioned under “Future Extensions”: add an additional interface method SmallVector<char> parseElements(), same for the printer. However, that will still require a new assembly format that starts with the shaped type.

Changing the assembly format is a large breaking change, so I’d suggest to keep both formats side-by-side, so that users have time to migrate (if we choose to deprecate the current one).

No changes are needed for this syntax, although, for consistency it could also start with the shaped type instead of the bits. In fact, if you don’t care about pretty-printing and data access via Attribute, you don’t need much of the TensorElementTypeInterface at all. Then you just need to know the size of each element, so that the attribute can provide an iterator over the dense data. If you don’t even need the iterator and are fine with just querying the blob of data, you have more or less that @KFAF suggested with Bit_DenseBitsAttr.

Maybe this also raises the question, what’s the difference between DenseElementsAttr and the various DenseI32ArrayAttr etc. The RFC that introduced that attribute also mentioned potentially removing/replacing DenseElementsAttr entirely.

7

There are two aspects to this:

  1. You can indeed not iterate over custom attributes because of the hard-coded list of supported types here.
  2. We have a quite long list of iterators, getters and other auxiliary functions for specific hard-coded types inside of DenseElementsAttr. It’s not part of the RFC, but I’m trying to remove this hard-coded list of types. I feel like a type interface is the first step towards that because it allows you to query the storage size in a generic way.

Correct, this will likely require a dedicated lowering pattern for arith.constant that picks up all ops that are not handled by the “default” pattern.

8

I see, thank you! I didn’t know of this iterator, since I always used the ElementsAttr interface to obtain underlying data directly without going through attributes. Precisely because it never seems like a good idea to instantiate attributes in MLIR if you can avoid it (see operation properties).

I think this should be changed, then. But couldn’t it also lazily construct the attributes from the underlying untyped bytes / hex blob? I think this would mean the DenseElementTypeInterface needs more methods, either an attribute factory or some sort of “bitcast” from llvm::APInt as the most generic way to store a bit vector.

9

The attribute factory is already part of the proposed interface: Attribute DenseElementTypeInterface::convertToAttribute(ArrayRef<char>)

Precisely because it never seems like a good idea to instantiate attributes in MLIR if you can avoid it (see operation properties).

I agree, I think this is useful mostly for debugging purposes / lit tests, where you may want to write down the elements in a human-readable form. All transformations/analyses should ideally directly operate on the bytes for performance reasons.

10

Ahh, I didn’t even parse that because of char. That’s on me though, I wasn’t thinking, clearly. Since C++17, I’m expecting to see std::byte for aliasing memory and just skipped over that.

I notice that LLVM widely uses char for that purpose still, but I also haven’t seen anyone make an argument for why new APIs shouldn’t be allowed to switch. Just out of curiosity, do you have an opinion on this matter?

11

Not a C++ expert… I think std::byte is pretty new? I just followed the existing design here, which is ArrayRef<char>, maybe std::byte would be better.

12

In the generic case, I think that std::byte communicates the intent better (blob) over char (blob or text), but in the context of LLVM / MLIR / ADT, we’ve been using Array<char> for passing blobs around and I don’t think it’s warranted to make an exception to that in just one API.

I’d rather stick to char and potentially do a larger cleanup to std::byte in the future, keeping the overall style consistent.

13

I noticed one downside though. While the current behavior of DenseElementsAttr does in fact mean that the minimum size a single element of llvm::APInt payload can occupy within the data array is 1 byte, the builtin BoolElementIterator actually iterates over individual bits. If I understand correctly, this behavior can’t be replicated with the proposed interface.

I don’t think this is a problem in general, but if eventual consolidation of all the possible data types of the DenseElementsAttr is intended, this will not be possible for i1 data, or cause it to take up much more space. Right?

14

Good observation, there is special handling for types with bitwidth 1.

// Storage width of 1 is special as it is packed by the bit.
...
// All other types are 8-bit aligned, so we can just check the buffer width
// to know if only a single initializer element was passed in.

The general question is how to deal with sub-byte types. I’m thinking for now to just change the interface method of DenseElementTypeInterface to return the bitwidth instead of the byte size. And to keep the storage format as is. Storing other types (e.g., f4E2M1FN) in a dense way is a breaking change and can be considered later.

15

I guess you would otherwise need to change the interface:

Attribute DenseElementTypeInterface::convertToAttribute(ArrayRef<char>, unsigned /*bitOffset*/)
LogicalResult convertFromAttribute(Attribute, SmallVectorImpl<char> &, unsigned /*bitOffset*/)

Or rather, keep the old interface and, for sub-byte types, just eat the cost of using a temporary buffer in which you align the bits to a byte? While slow, we established that this is the slow path anyways, and it would prevent bad implementations from clobbering the wrong values. Then we could get sub-byte types at reasonable cost?

16

I cleaned up the implementation: [mlir][IR] Generalize`DenseElementsAttr` to custom element types by matthias-springer · Pull Request #179122 · llvm/llvm-project · GitHub. Feel free to leave comments/reviews.

Looking for feedback especially for the new type-first assembly format and this observation:
The new syntax is needed because we have to know the number of elements that we are about to parse. (Note: We could avoid the new syntax if we could assume that the assembly format of an element attribute does not start with [.)

The existing literal-first syntax is ambiguous the moment you allow attributes as elements:

dense< [  [      ...
         ^^^
  Does this LSquare indicate a new dimension of the shaped type
  or the beginning of an attribute's assembly format?

(The more I think about it, to keep the PR small, maybe start without a new syntax…)

edit: The new syntax cannot be avoided easily. When seeing a digit, the TensorLiteralParser wouldn’t know if it’s parsing an integer or an attribute in its assembly format.

17

Was this discussed?

What is the overhead of now using interface here?

Please yes, the current format is much easier to read.

18

There are two kinds of overhead:

  1. Performance (running time) overhead due to type interface / interface method calls instead of a hard-coded list of types in a switch-case. You will see this overhead when dense elements are converted to/from Attribute: parsing/printing IR, DenseElementsAttr::get(ArrayRef<Atttribute>), AttributeElementIterator. In performance-critical code, users are expected to avoid this API and operate on raw bytes, so I there should be no performance overhead. (These APIs are used mainly for debugging and test cases. If you never use these APIs (and don’t write Lit tests), you could probably use element types that don’t implement the interface.)
  2. Code complexity: There is a new type interface, which allows us to remove hard-coding for supported types in a few places, e.g., when constructing a DenseElementsAttr.

Before:

DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                         ArrayRef<Attribute> values) {
  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()) {
      // ...
    }
    // Must be float.
    // ...
  }

  // If the element type is not based on int/float/index, assume it is a string
  // type.
  if (!eltType.isIntOrIndexOrFloat()) { /* ... */ }

  // Compress the attribute values into a character buffer.
  for (unsigned i = 0, e = values.size(); i < e; ++i) {
    if (auto floatAttr = llvm::dyn_cast<FloatAttr>(values[i])) {
      // ...
    } else if (auto intAttr = llvm::dyn_cast<IntegerAttr>(values[i])) {
      // ...
    } else {
      // Unsupported attribute type.
      // NOTE: This limitation is the main reason for this RFC.
      return {};
    }
  }
}

After:

DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                         ArrayRef<Attribute> values) {
  assert(hasSameNumElementsOrSplat(type, values));
  Type eltType = type.getElementType();

  // Handle strings specially.
  if (!llvm::isa<DenseElementType>(eltType)) { /* ... */ }

  // All other types go through DenseElementTypeInterface.
  auto denseEltType = llvm::cast<DenseElementType>(eltType);
  SmallVector<char> data;
  for (Attribute attr : values) {
    LogicalResult result = denseEltType.convertFromAttribute(attr, data);
    if (failed(result))
      return {};
  }
  return DenseIntOrFPElementsAttr::getRaw(type, data);
}

I have another PR in preparation that removes the special case for dense string elements attributes, so that the getter turns from the current 80 LoC implementation into a single loop with a call to DenseElementType::convertFromAttribute.

So overall, I’d say that this reduces complexity.

This is mainly due to repetitive type literals in the assembly format, and there are ways to improve that (see “Future Extensions” section in the RFC.) Having to maintain two parsers/printers is a bit unfortunate, so I wouldn’t want to write this idea off quite yet. At the same time, it’s a large change that affects thousands of test cases, so it may not be worth it. We can revisit that later.

19

For the current status of this RFC and PR: I merged the PR a while ago, but then had to revert due to a broken Windows build bot. There is a problem with the C++ code that we auto-generate from TableGen. The issue seems to be Windows-only.

Assertion failed: implDenseElementType && "`::mlir::FloatType` expected its base interface `::mlir::DenseElementType` to be registered"

This PR exercises an ODS feature that is not widely used yet: interface inheritance. I couldn’t find the bug in ODS code generator yet…