[RFC] Add a builtin ComplexAttr

Motivation:

Complex constants are already supported inside DenseElementsAttr, but DenseElementsAttr::AttributeElementIterator::operator*() runs into an llvm_unreachable(“unexpected type”). This operator is (indirectly) used by the getSplatValue() method which in turn is needed by the FoldSplatConstant pass in linalg. If we have a builtin ComplexAttr, we could implement the missing functionality in this operator and add support for folding complex splat constants.

Suggested syntax:

Ideally we would use the same syntax for complex values inside DenseElementsAttr and for complex values for ComplexAttr. Complex values inside DenseElementsAttr are represented as “(-1.0, 1.0)”, i.e. the two float or int constants are inside parentheses and separated by a comma. However if we want to reuse that for ComplexAttr, the parser would have to somehow distinguish between reading a ComplexAttr or a TypeAttr that has a function type (e.g. “type = (tensor<2xf32>) → tensor<2xf32>” inside a attribute dictionary). Maybe this is possible if TypeAttr can only occur in certain contexts, but even then it would mean that ComplexAttr could never be used in the same context as a TypeAttr.
Therefore the suggestion is to change the syntax which is used for complex values inside DenseElementsAttr and use this new syntax also for ComplexAttr.
Here are some ideas how the syntax could look like:
complex(1.0, -1.0) : complex<f32>
complex(1.0, -1.0) : f32
pair(-1.0, -1.0) : complex<f32>

The idea with pair would be if we want to support a generalization of ComplexAttr, but without a need for a pair of values in a different context, this would not make much sense.
For complex values inside DenseElementsAttr we would still omit the type, because the type of elements can be inferred from the type of DenseElementsAttr.
Edit: I just realized that even with the keyword complex in front of the parenthesis, we would have difficulty with distinguishing it from a type attribute with type complex<f32>. We would need to look ahead whether the complex keyword is followed by a ‘(’ or a ‘<’. Suggestions for different syntax is welcome :slight_smile:

Details on the specification for ComplexAttr:

ComplexAttr could be specified as a pair of attributes, with constraints on which attributes we allow as first and second element (custom Verifier). As StorageClass, it would be easiest to use std::pair<Attribute, Attribute>, because std::pair already has DenseMapInfo support. If we were to use std::complex<Attribute>, we would have to add DenseMapInfo support for it. With extra class definitions we could still add a real() and imag() method to ComplexAttr class.

Further work:

Once we have ComplexAttr, we can support complex constants in the standard dialect and in the LLVMIR dialect. For the LLVMIR dialect, the ModuleTranslation would have to be changed to add limited support for struct type (i.e. only supporting struct type if the attribute is a ComplexAttr). I already have a prototype where this is working end-to-end with lowering tf.Reciprocal (i.e. 1 / x) for complex types.

Why not the same as DenseElementsAttr with 2 elements storage wise? A pair of Attributes would indeed allow two different types (same as std::complex) which the former wouldn’t (so no need to verify as one can’t have different types) and I don’t think there is storage benefit to pair format and computationally you’d want to operate on std::complex<float> or some such rather than Attribute.

Can you please expand a bit about the 2 elements storage? I am not familiar with DenseElementsAttr, but it looks to me it is just enforcing the “same type” constraint through the get() method, while using ArrayRef<char> as storage. Using ArrayRef<char> as storage seems complicated, and there are iterators to be able to retrieve values.
ArrayRef<Attribute> would definitely work as well, with an additional constraint that we want just 2 elements. But I guess this is not what you had in mind?

Right. Attributes have optional types, so you should be able to make one of this and give it a “complex” type. I don’t think we need to create a new attribute kind for this.

-Chris

I think that the point is that ArrayRef<Attribute> (or pair<Attribute, Attribute>) has an indirection with respect to the value stored, while DenseElementAttr does not have this. It is the difference between std::pair<float *, float *> and std:pair<float, float> conceptually: the latter seems better to store std::complex<float>.

Thanks for the explanation :slight_smile: Yes, that makes sense.

So assuming we reuse DenseElementsAttr with 2 elements instead of introducing a ComplexAttr, I guess we would still want to use a different syntax in the IR? Or would dense<1.0, -1.0> : complex<f32> be ok? I fear it would make the IR a bit confusing if a constant of type tensor<complex<f32>> is “dense<(1.0, -1.0)>” but a constant with type complex<f32> is “dense<1.0, -1.0>”

Regarding implementation: I think Jacques may have meant to still add a ComplexAttr, but set baseCppClass to DenseElementsAttr and set genStorageClass = 0 to reuse the one from DenseElementsAttr?
As far as I can tell, technically both (reusing existing DenseIntOrFPElementsAttr or adding a new attribute derived from DenseElementsAttr) would still allow to introduce a custom syntax and allow ComplexType.

Edit: after playing around with this for a while and further reading of the ElementsAttr code, I think anything deriving from ElementsAttr needs to be a ShapedType. So I guess that means we would have to use e.g. tensor<2xf32> instead of complex<f32> as type. Then I think the only way to still know it is a complex constant would be if we introduce a new attribute for it. But maybe I am wrong, as I said I am not familiar with DenseElementsAttr, I have just read the code.

So what should be the way forward here? Can I add a builtin ComplexAttr? Should it be derived from DenseElementsAttr, and if yes, how do we teach the parser to recognize ComplexAttr (my guess is, we would have to use a different keyword than dense, or change the type constraints stuff for ElementsAttr to allow ShapedType and ComplexType so that ComplexAttr can be recognized by the type).

I think it makes sense to reuse either dense or array. That would give you dense<1.0, -1.0> : complex<f32> or [1.0, -1.0] : complex<f32>.

-Chris

While we can make both of these options stored internally the same way in a ComplexAttr, parsing the array syntax will have some overhead compared to the dense I think (I don’t know if this is significant enough that we should case, probably not).

Parsing dense<1.0, -1.0> currently requires the type to be a tensor. We can “hack” the TensorLiteralParser to make a special case to handle type complex, but I guess then we should possibly also rename TensorLiteralParser to something else? Also, the type of ComplexAttr would still need to be a ShapedType if it is derived from DenseElementsAttr, but we could let the AsmPrinter print complex<f32> as type for ComplexAttr.
But maybe I am missing something :slight_smile:

Ah ok, then I’d go with the array literal syntax. Mehdi, what overhead are you thinking of here?

Parsing an array attribute means that for each element you create a new FloatAttr, then assemble them both in a new ArrayAttr, while we’d just to deconstruct it immediately to compress the two floats into a dense pair for the ComplexAttr.

That said it would only affect a “standalone” single ComplexAttr and not a collection, parsing a large dense tensor of complex would go through the TensorLiteralParser, so I assume we’d accept something like: dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2xcomplex<f32>> ?

(in which case the array syntax would be the same for a standalone ComplexAttr and for the dense case).

Yeah, agreed it would not inline the elements, that’s a good point – you’re right there is a difference in storage. That said, I don’t think it is a very significant memory size impact though, given the elements are themselves uniqued.

I would assume this is only a temporary array during parsing and internally we’d store a ComplexAttr densely. Which means this overhead is only paid for parsing textual assembly, not when using ComplexAttr programmatically.

I fear if we go through the ArrayAttr parsing, there will be problems with parsing the individual elements as the right type, unless each element has the type attached. The advantage of going through the TensorLiteralParser would be that it already knows how to store the tokens so that we can process them once we know the type of the elements. We could of course do the same for array parsing, but right now, it looks to me that ArrayAttr does not even have a type in IR format (instead it seems to rely on the attributes inside the array to have a type if needed).
So I guess if we want to store ComplexAttr densely, we would still want to reuse existing code for DenseElementsAttr? So we could set DenseElementsAttr as a base class, and use tensor<2xf32> as storage type. The AsmPrinter would however need to print it in a way so that the parser knows it is a ComplexAttr. If we use array syntax, we could change the ArrayAttr parsing logic to accept an optional type. And then either also print the type for each element in the array (e.g. [1.0 : f32, 2.0 : f32] : complex<f32>), or try to refactor TensorLiteralParser so that it can also be reused for ArrayAttr parsing?

What do you mean by this? What are you trying to achieve here?

I understand that the concern was that storing ComplexAttr as a std::pair<Attribute, Attribute> (as it was done in my prototype) has negative performance implications, and instead it should be stored densely, similar to how two values of the same type are stored in DenseElementsAttr. I am not attached to a particular idea or solution, I just want to get an agreement on a solution that works :slight_smile:
And from what I understand, technically it would be easiest to store a ComplexAttr the same way as a tensor<2xf32>. The difficulty with this is just that somehow we still need to be able to parse this and know that it is a ComplexAttr and not a regular DenseElementsAttr. Alternatively, I guess we could actually use DenseElementsAttr for complex constants (and not create a ComplexAttr), but let the constant have a complex type, and when printing the constant, we would print the complex type as well. Right now, it seems the type of the constant is only printed if the attribute is a SymbolRefAttr (printing the attribute from the constant would already print the type, anyway). So it would potentially look like this:
dense<1.0, -1.0> : tensor<2xf32> : complex<f32>
Which is kind of close to what you already suggested before (maybe it would be possible to adjust the printer and parser to avoid printing the tensor<2xf32> type, but it looks to me this could be tricky).
Maybe that solution isn’t actually so bad, even though it means that we would print complex values differently whether they are part of a complex tensor constant, or a complex scalar constant.

This is what I’m asking about. What negative performance implications are you worried about here? All these things are uniqued. Are you looking to save a couple bytes on complex half float values or something? Are you actually worried about the memory usage of the compiler being a few bytes more?

-Chris

This is a fairly long discussion by now and at the risk of derailing it even further, here is my take.

DenseElementsAttr is designed for attributes that have a shape, so it does not really apply to complex scalars. The only commonality is that both have more than one element that make up their value.

We have specialized attributes for Float and Integer. Why should we not have one for Complex? I think the potential reuse of code here does not outweigh the added complexity of making parsing etc. work.

I think this might even be defined using TableGen. Or is that path discouraged?