[RFC] Always printing type aliases for certain types

[mshockwave]

TL;DR

Our current way of printing LLVMStructType in the LLVM dialect doesn’t scale well, in terms of memory consumption and printing speed, when it comes to complicate structs (multiple layers of nested structs). We are proposing a new way to print LLVMStructType, consisting of type aliasing and shorthand syntax.
Note that this proposal not only changes LLVM dialect but also affects how types alias is implemented.

Problem definition

Here is what struct types look like in LLVM dialect:

!llvm.struct<"simple", (i32, i32)>
!llvm.struct<"nested", (ptr<struct<"simple", (i32, i32)>>, i16, i8)>
!llvm.struct<"self_reference", (ptr<struct<"self_reference">>, i16, i32)>
!llvm.struct<"nested_sibling", (ptr<struct<"simple", (i32, i32)>>,
                                ptr<struct<"simple", (i32, i32)>>)>

From “nested” we saw that LLVM dialect always spells out the body of nested struct types (i.e. the inner “simple”), unless it’s referencing itself, like “self_reference” struct above. In which case it uses a shorthand syntax, !llvm.struct<"struct_name">.
Note that shorthand syntax is only applicable to types that had appeared in the parent hierarchy, so “nested_sibling” still spells out “simple” twice in sibling element types.

Imagining we have a struct type that is really wide (has large number of child element types) and deep (has pointer to another struct that also has pointer to another different struct…) and there are many usages of this struct type among all the operations (references of struct type in operations also spell out the struct body), it’s really easy to spend significant amount of memory for those (spelled-out) struct bodies strings and take quite a long time to print them out as well. From our previous trials using real-world code (SQLite3), it took more than 5 minutes just to print MLIR out before it was killed due to OOM (on a 32GB memory machine).

Proposed solution

The solution is consisting of two parts. First, the printer always prints the shorthand syntax for struct types whose body had been printed before in the same module. In other words, we memorized the struct bodies we printed before globally and avoid printing them again.

llvm.func @foo(%arg0: !llvm.ptr<i8>) {
  // "hello" 's body is spelled out here
  %0 = llvm.bitcast %arg0: !llvm.ptr<i8> to !llvm.ptr<struct<"hello", (i32, i32)>>
  // "hello" is printed in shorthand syntax here
  %1 = llvm.load %1: !llvm.ptr<struct<"hello">>
  ...
}

// "hello" is printed in shorthand syntax here
llvm.func @bar(%arg0: !llvm.ptr<struct<"hello">>) {
  ...
}

However, it might be a little hard for developers to tracing back to find the body of a struct when it is printed in shorthand syntax, because the body might be printed in arbitrary position in previous code.

To aid this issue, the second part of this proposal creates type aliases for struct types:

// "hello" 's body is always spelled out at top level
!llvmStruct_hello = !llvm.struct<"hello", (i32, i32)>
llvm.func @foo(%arg0: !llvm.ptr<i8>) {
  %0 = llvm.bitcast %arg0: !llvm.ptr<i8> to !llvm.ptr<!llvmStruct_hello>
  %1 = llvm.load %1: !llvm.ptr<!llvmStruct_hello>
  ...
}

llvm.func @bar(%arg0: !llvm.ptr<!llvmStruct_hello>) {
  ...
}

Whether we should only create aliases for expensive struct types or every struct types is opened for discussion. The point is, developers can quickly navigate to the struct body since type alias definitions always appear at the top.

Let’s look an example of these type aliases:

!alias_a = !llvm.struct<"nested_sibling", (ptr<struct<"simple", (i32, i32)>>,
                                           ptr<struct<"simple">>)>
!alias_b = ...
!alias_c = ...
...
!alias_z = !llvm.struct<"simple">

Now when we’re tracing the type using !alias_z, a similar problem we discussed above emerged again: We would like to know the body of “simple”, but that piece of information is located somewhere in the previous type alias definitions, possibly buried in another large struct body.
A simple re-ordering can solve this issue though:

!alias_z = !llvm.struct<"simple", (i32, i32)>
!alias_a = !llvm.struct<"nested_sibling", (ptr<struct<"simple">>,
                                           ptr<struct<"simple">>)>

Basically, when encountering a nested struct, it will be better to have its type alias definition hoisted before the current one. It looks more naturally and it’s also easier for developers to track the struct body from a type alias reference.
Such ordering policy is the last piece of change we would like to propose here.

Your comments on this proposal are much appreciated!

Q & A

Q: Are type alias definitions ordered now?
A: Yes, they’re organized in lexical order by their alias names (I don’t know why).

Q: Can we use previously defined struct type aliases in another struct type alias definition appears later, instead of shorthand syntax?
A: I tried this before, but a MLIR type alias always needs to appear before its reference so it’s not applicable for struct types that have mutual references. We can mix these two strategies together though (using type alias for most of the time and only use shorthand syntax for mutual referencing structs). It’s opened for discussion.

[ftynse]

I am happy to see some progress in this direction! FWIW, this is rather similar to what was originally discussed in [RFC] First-Party Support for LLVM Types under “region-scoped type definitions”. There are a couple of broad concerns that I remember as important:

• Extensibility – this already proposes to change the core parser, so it is worth discussing how this will work for types other than LLVMStructType.
• Structure – the proposal makes the assumption of a “module->function->blocks” structure of the IR, similar to what LLVM IR has, that does not hold of all of MLIR. We can and actually do use nested modules with the LLVM dialect in them. Should the aliases span across modules or be defined in each module? What should happen if we were to introduce an LLVMModuleOp and use it instead of the builtin (the top-level entity in the IR is a Block, not an Op)? How to deal with repeated names / shadowing? Inlining a module/region with conflicting alias names?
• Genericity – does this roundtrip correctly when the LLVM dialect is not registered (and the ops are printed in the generic form)?

Some more minor comments below.

@mshockwave

in terms of memory consumption

A small clarification that this only concerns the memory used while printing or to store the textual form. It doesn’t look like this proposal anyhow affects the in-memory IR, which only stores pointers to nested types anyway.

@mshockwave

First, the printer always prints the shorthand syntax for struct types whose body had been printed before in the same module.

A big -1 on this. The first occurrence is not necessarily the most obvious one, it might be within an attribute attached to an op from a different dialect deeply nested in some region constructs from other unrelated dialects, potentially in the generic form. This wouldn’t make the IR any more readable IMO. Using explicit aliases consistently sounds significantly better, and also makes it clear that an alias is being used and its definition must be looked up, instead of confusing the reader who might erroneously assume they are facing an opaque struct type.

@mshockwave

A simple re-ordering can solve this issue though:

Ordering cannot solve issues with circular dependencies:

!a = !llvm.struct<"a", (ptr<struct<"b", (ptr<struct<"a">>)>>)>
!b = !llvm.struct<"b", (ptr<struct<"a", (ptr<struct<"b">>)>>)>

Allowing “forward-use” of aliases can, however:

// Here, !b is used as "forward declaration", similarly to symbols.
!a = !llvm.struct<"a", (ptr<!b>)>
!b = !llvm.struct<"b", (ptr<!a>)>

Instead of enshrining the “print on first use” approach, which I consider problematic, I would rather suggest always replacing some types with aliases, as opt-in type printer feature. We do that for affine maps, and we currently do the inverse for types and always expand the aliases.

[Mogball]

@ftynse

I would rather suggest always replacing some types with aliases, as opt-in type printer feature. We do that for affine maps, and we currently do the inverse for types and always expand the aliases.

Yeah, I think it might be possible to get away with just letting type printers print aliases.

[mshockwave]

@ftynse

A big -1 on this. The first occurrence is not necessarily the most obvious one, it might be within an attribute attached to an op from a different dialect deeply nested in some region constructs from other unrelated dialects, potentially in the generic form. This wouldn’t make the IR any more readable IMO.

Yep! it was also a pain for me to read, which was the motivation behind using type alias in the second half of the solution.

@ftynse

Ordering cannot solve issues with circular dependencies:

!a = !llvm.struct<"a", (ptr<struct<"b", (ptr<struct<"a">>)>>)>
!b = !llvm.struct<"b", (ptr<struct<"a", (ptr<struct<"b">>)>>)>

Allowing “forward-use” of aliases can, however:

// Here, !b is used as "forward declaration", similarly to symbols.
!a = !llvm.struct<"a", (ptr<!b>)>
!b = !llvm.struct<"b", (ptr<!a>)>

Yes, this syntax is ideal, I also tried to implement it but type alias doesn’t support forward-use right now, which brings the discussions to…

@ftynse

Structure – the proposal makes the assumption of a “module->function->blocks” structure of the IR, similar to what LLVM IR has, that does not hold of all of MLIR. We can and actually do use nested modules with the LLVM dialect in them. Should the aliases span across modules or be defined in each module

I like the idea of region-scoped type definitions and I feel like we can do it on top of type alias (i.e. reusing the type alias syntax instead of creating a new one). The key problem is how we, as you mentioned in the previous proposal, “…provide a sequence point (first block or end-of-region) to check if all forward references were defined”. In a long term I firmly believe we should have a region-scope type alias, and possibly shadowing aliases in inner regions.

Unfortunately right now I don’t have the energy to make such radical changes. I think we should have a transitional approach to implement forward-use type aliases first. My hand-waving thoughts would be either of the following:

1. All type alias definitions that appear before an (top-level) operation need to be resolved.
2. Within the LLVM dialect, whenever it encounters an unresolved type alias, it LLVMStructType::getIdentified an empty struct type using the alias name. Then, we need to create a new hook for the parser to notify the dialect whenever a new type alias is added. Such that the dialect can resolve the forward-used alias name. This, however, requires changes in the LLVMStructTypeStorage regarding how a struct type is identified: both by its struct name (if there is any) and the alias name.

These are just some rough ideas to implement forward-use type aliases without significant changes.

[mshockwave]

@mshockwave

My hand-waving thoughts would be either of the following:

1. All type alias definitions that appear before an (top-level) operation need to be resolved.
2. Within the LLVM dialect, whenever it encounters an unresolved type alias, it LLVMStructType::getIdentified an empty struct type using the alias name. Then, we need to create a new hook for the parser to notify the dialect whenever a new type alias is added. Such that the dialect can resolve the forward-used alias name. This, however, requires changes in the LLVMStructTypeStorage regarding how a struct type is identified: both by its struct name (if there is any) and the alias name.

These are just some rough ideas to implement forward-use type aliases without significant changes.

After some studies I found that supporting forward-use for types in general is not so easy. First, we need to ask the question of what placeholder type do you put for a forward-use type before it got resolved.
For LLVMStructType, we can use an empty struct as placeholder , but the key is that you need to fetch the right Type (singleton) object in the first place. LLVMStructType is indexed / hashed by its struct name if it’s identified struct (let’s ignore literal struct here). In the ideal case, we can use the alias name to LLVMStructType::getIdentified an empty struct. Unfortunately, alias names are always sanitized, which might be different from the original / real struct name.

!hello = !llvm.struct<"hello", (ptr<!struct2EB>)>
// "struct.B" is changed into "struct2EB"
!struct2EB = !llvm.struct<"struct.B", ...>

So I think a potential solution will be creating a name that can be demangled:

!hello = !llvm.struct<"hello", (ptr<!S6structH22ES1B>)>
!S6structH22ES1B = !llvm.struct<"struct.B", ...>

In this simple scheme, “S<N>” will take the next N char literally, “H<N>” will interpret the next N char as hex and further translated them back to a string.
Note that this new mangle / demangle scheme can implemented by each dialect. No core changes are needed here.

After that, as I briefly mentioned previously, we might need to add a new way for the parser to call back to the dialect whenever it parses a new type alias definition. For LLVM dialect, upon receiving that callback, it can finally set the right struct body for the corresponding forward-use type.

The biggest concern I have for this approach is that not every dialect has types that can be identified by string, or something that can be serialized to a string. Also, not every type can be mutated after its creation (e.g. !llvm.ptr uses the element type as its identifier so there is no way to mutate it afterward). But at the same time, supporting forward-use type is inherently a difficult problem, it’s fine that not every types can be forward-used (even in LLVM, only struct types can have recursive dependencies). Plus, the only core change is the interface to call back from parser to a dialect whenever there is a new type alias definition being parsed, in which IMHO, this new interface is considered general / benign.

[ftynse]

@mshockwave

The biggest concern I have for this approach is that not every dialect has types that can be identified by string, or something that can be serialized to a string.

I think that forward references, if we introduce them, should be resolved by the parser itself. We don’t want to bring the text format-specific information to the IR itself, similarly to how we don’t care about value names. We can parse graph regions with values used before they are defined. So the questions of placeholder types, identifiers and mutations shouldn’t even appear.

Now I appreciate this is not a simple modification, FWIW, I haven’t done it myself back at the time. So if we are looking for a mostly low-cost quality of life improvement, I would suggest dropping the forward reference idea and focusing on just aliases that have been already printed. In the mutually-referencing struct cast, the first struct will have to remain in the fully spelt out form, but the second struct can use the alias. Any IR below can use both aliases, so this should significantly improve on the current overly verbose solution.

Also, I would recommend renaming the thread to something like “always printing type aliases for certain types”, otherwise people who may be interested in this can ignore the thread because it only refers to a specific LLVM dialect type.

[mshockwave]

@ftynse

similarly to how we don’t care about value names. We can parse graph regions with values used before they are defined. So the questions of placeholder types, identifiers and mutations shouldn’t even appear.

I think the story on parsing forward-used values (and blocks) are different: they have a use list so we can always create an empty mlir::Value as placeholder first and simply RAUW it after the real definitions are parsed later. But MLIR Type doesn’t have such convenient property. Every references to a type are holding the unique type pointer directly, as long as the placeholder type is different from the real type definition, the pointer is different. Thus, I think it’s not easy to have a generalized way (to parse forward-used types) for different types in different dialects.

@ftynse

I would suggest dropping the forward reference idea and focusing on just aliases that have been already printed. In the mutually-referencing struct cast, the first struct will have to remain in the fully spelt out form, but the second struct can use the alias. Any IR below can use both aliases, so this should significantly improve on the current overly verbose solution.

This is a good idea and super easy to do. I just tried it and it scaled well in medium-sized use cases…but not the large ones. I also implemented an advanced sorting algorithm on the ordering of type alias definitions to reduce the total length of strings in a group of mutually referencing type aliases, but that didn’t work out either.

@ftynse

Also, I would recommend renaming the thread to something like “always printing type aliases for certain types”, otherwise people who may be interested in this can ignore the thread because it only refers to a specific LLVM dialect type.

Will do!

[ftynse]

@mshockwave

I think the story on parsing forward-used values (and blocks) are different: they have a use list so we can always create an empty mlir::Value as placeholder first and simply RAUW it after the real definitions are parsed later. But MLIR Type doesn’t have such convenient property. Every references to a type are holding the unique type pointer directly, as long as the placeholder type is different from the real type definition, the pointer is different. Thus, I think it’s not easy to have a generalized way (to parse forward-used types) for different types in different dialects.

It is possible to have an extra layer of indirection where the Type pointer points to another address, which in turn points to either the actual storage (or is directly used as such if the storage fits into the pointer size minus some flag bit) or some placeholder token. The outer pointer remains stable. This has a cost beyond the parsing though, and I’d like to keep this to the parser as much as possible.

The generalized way is to introduce some TypeAsmTypeInterface by analogy with OpAsmOpInterface, we will need that anyway for types to declare they want a different aliasing behavior. It’s the methods of the interface that sound tricky. It sounds like we need some mechanism of postponing type creation until the entire connected component of mutually dependent types is parsed. And we don’t have a generic way of creating a type, unlike operations/blocks. There was a question on the chat about adding one though.

@mshockwave

This is a good idea and super easy to do. I just tried it and it scaled well in medium-sized use cases…but not the large ones. I also implemented an advanced sorting algorithm on the ordering of type alias definitions to reduce the total length of strings in a group of mutually referencing type aliases, but that didn’t work out either.

Interesting… The situations sounds quite specific, there seems to be a lot of mutually dependent and fairly large types (template instantiations?) that aren’t used much, otherwise aliases should have helped. Could you share such an example or characterize it somehow?

@mshockwave

Will do!

Thanks! I’m particularly interested in what @River707 has to say on this point.
Also @math-fehr for the generic type constructor plans if any.

[math-fehr]

Also @math-fehr for the generic type constructor plans if any.

Sadly, I do not think there is any plans to add generic type/attribute constructors yet.
I’m happy to discuss it, but it seems that for now the cost of adding them would be quite big compared to the reward (Not enough people showed interest at least).

To summarize, the main problem with adding generic constructors is that we need to find a way to handle the “base” c++ attribute/type parameters. One way of doing this would be to add multiple attributes to handle all the cases we need (I’m not sure how many we would need), and then add them in the “AttrOrTypeParameter” definition in TableGen.

If you want to discuss it, I’m happy to take the discussion somewhere else. But unless we fix this problem (or make it a smaller change), I’m not sure the generic attributes/types interface will be planned.

[Mogball]

@math-fehr

To summarize, the main problem with adding generic constructors is that we need to find a way to handle the “base” c++ attribute/type parameters. One way of doing this would be to add multiple attributes to handle all the cases we need (I’m not sure how many we would need), and then add them in the “AttrOrTypeParameter” definition in TableGen.

I don’t see why you couldn’t introduce a subset of attributes/types that are “generic” (e.g. contains only Attribute parameters). There was some level of discussion around this when discussing deprecating StructAttr, which is pretty close to a “generic attribute”.

[mshockwave]

@ftynse

Could you share such an example or characterize it somehow?

Here is the (somewhat) reduced LLVM IR: sqlite3.O2.sqlite3DeleteTrigger.ll · GitHub

The problematic type is struct.sqlite3. Let’s take one of its element types, struct.Vdbe, as example. Both of them are pretty heavy weight, simply reordering them won’t make a big difference. More generally speaking, if we depict these type as a graph, many (strongly) connected components are nearly fully connected, thus I think it’s hard to get a good ordering without using complex sorting algorithm.

[mshockwave]

@mshockwave

After that, as I briefly mentioned previously, we might need to add a new way for the parser to call back to the dialect whenever it parses a new type alias definition. For LLVM dialect, upon receiving that callback, it can finally set the right struct body for the corresponding forward-use type.

Update: there was actually a blind spot I overlooked in this statement, which was the origin of most of the following discussions: For LLVMStructType, as long as we use the correct struct name to retrieved an empty, identified LLVMStructType object as a placeholder, later on when that very LLVMStructType is defined, it will grab the same LLVMStructType object (since MLIR types are unique and LLVMStructType is indexed by its struct name) and populate the correct struct body for us.

At least in the case of LLVM dialect, there is no need to add a hook for parser to call back into dialect to handle alias definitions. This does require the name mangling / demangling scheme I mentioned in the previous comment. AsmParser also needs some modifications to return the alias name if it finds a forward-using type alias. But both of them are parsers (core and dialect parsers) only changes so I think we’re good here. I also tried to implement in this way and it worked like a charm for our large (LLVM IR) code.

That said, this is still not a generic solution for other dialects. We probably still need a generic type as previous discussions suggested at some point…

[River707]

Sorry missed this thread originally.

@mshockwave

At least in the case of LLVM dialect, there is no need to add a hook for parser to call back into dialect to handle alias definitions. This does require the name mangling / demangling scheme I mentioned in the previous comment. AsmParser also needs some modifications to return the alias name if it finds a forward-using type alias. But both of them are parsers (core and dialect parsers) only changes so I think we’re good here. I also tried to implement in this way and it worked like a charm for our large (LLVM IR) code.

That said, this is still not a generic solution for other dialects. We probably still need a generic type as previous discussions suggested at some point…

I would strongly avoid anything that makes alias names load bearing. We make no guarantees on their structure, and I don’t think we should.

For aliases in general, they are ordered by name because for the longest time we didn’t even have a way of introspecting nested elements (we do now though). It seems like one simple fix would be to always visit the elements of an attribute/type when building aliases, which would ensure that elements have aliases defined before the user. We can order aliases effectively however we want, and could even have name-ordering when there are no-dependencies and visit-ordering when there are.

For recursive things, we could bite the bullet and have some API that is specific to mutable attributes/types (i.e. the ones that can change their storage post construction) somewhere that is like printNonMutable. During alias construction it’s fairly easily to detect recursion, in which case a recursive attribute/type would print an alias for the non-mutable instance first, and then later on the full instance. In the case of LLVMStructType, that would first print an alias with just the name (i.e. the non-mutable bit) and then later the alias with the full instance (name+body). We will end up with something similar to this for the binary/bitcode MLIR format anyways, so it seems fine to have a textual counterpart.

[mshockwave]

@River707

It seems like one simple fix would be to always visit the elements of an attribute/type when building aliases, which would ensure that elements have aliases defined before the user.

Right, I’ve been doing this in my prototype, including some extends of sorting on type alias definitions. But the problem / performance bottleneck is always types with recursive references. In one of the previous comments I stated that unless we’re using some really advanced algorithms, types with high degree of interconnections are pretty hard to relaxed via sorting, IMHO.

@River707

in which case a recursive attribute/type would print an alias for the non-mutable instance first, and then later on the full instance. In the case of LLVMStructType, that would first print an alias with just the name (i.e. the non-mutable bit) and then later the alias with the full instance (name+body).

Just want to make sure I understand your idea correctly:

!alias0      = !llvm.struct<"foo">
!alias1      = !llvm.struct<"bar", (i32, !alias0)>
!alias0_real = !llvm.struct<"foo", (i16, !alias1)>
!alias2      = !llvm.struct<"other", (i8, !alias0_real)>

Does the snippet above match your idea? (I’m using different alias names for “foo” because I guess we don’t want to change the redefinition rule for type aliases)

[River707]

@mshockwave

Right, I’ve been doing this in my prototype, including some extends of sorting on type alias definitions. But the problem / performance bottleneck is always types with recursive references. In one of the previous comments I stated that unless we’re using some really advanced algorithms, types with high degree of interconnections are pretty hard to relaxed via sorting, IMHO.

Do we really need a complicated algorithm here? Visitation order should give us a basic relative priority, of which we could also have very simple weighting heuristics. E.g. we could bias “mutable” things to be later given that everything beforehand can use the much simpler “forward ref” alias instead of the one with the full definition.

@mshockwave

!alias0      = !llvm.struct<"foo">
!alias1      = !llvm.struct<"bar", (i32, !alias0)>
!alias0_real = !llvm.struct<"foo", (i16, !alias1)>
!alias2      = !llvm.struct<"other", (i8, !alias0_real)>

Does the snippet above match your idea? (I’m using different alias names for “foo” because I guess we don’t want to change the redefinition rule for type aliases)

Yeah, that’s essentially what I had in mind.

[mshockwave]

@River707

we could also have very simple weighting heuristics. E.g. we could bias “mutable” things to be later given that everything beforehand can use the much simpler “forward ref” alias instead of the one with the full definition.

Oh, sorry about the confusion, I was simply talking about one of the previous attempts[1][2] to solve the problem by only using some sort of ordering on the alias definitions. More specifically, we print recursion-free types first, then, for types with mutual references, we always spell out the full body of types that come first.

Back to your proposed idea, I think it’s doable and personally I’m fine with the syntax – it’s similar to forward function declaration in C-like languages.
@ftynse and @Mogball , any thoughts on this solution?

[Mogball]

The direction sounds good to me. Though I’m not sure about the syntax for forward declarations of aliases.

[ftynse]

This is aligned with what I had in mind, i.e., having an interface to handle these things. Care must be taken in defining the interface and making sure it works only for aliases (we don’t want the incomplete type to appear in the middle of IR for readability reasons).

Regarding the syntax, I’d rather reuse the alias name if possible instead of having !alias0 and !alias0_real, which makes them look different on a quick skim. Presumably, the type interface can say whether the parsed type is complete or not, and we can allow re-using the name for the complete version.