Polymorphic functions and SSA values in types

MLIR
1

Hi,
We’re experimenting with using MLIR in the Verona compiler. One of the things we are looking into is using MLIR as early as possible in the pipeline, in particular before monomorphization of generics. All type checking would be operating on the MLIR.

This means we need to encode type parameters in MLIR types. For example, from the following Verona function, we need a way to define the type parameter X, and a way to reference it in types.

fun id[X](a: X): X { return a; }

The first approach I can think of is using indices for type variables, and an attribute to specify the number of type parameters. The type X becomes !verona.var<0>.

fun @id(%a: !verona.var<0>) -> !verona.var<0> attributes { type_parameters = 1 } {
  return %a :  !verona.var<0>
}

I think its a bit unwieldy to manipulate, especially when having multiple levels of generics (eg. a generic class with a generic method), so I’m looking for a better solution.

One of these is to allow SSA values in types, allowing type parameters to be represented as regular parameters. Here, !verona.type is the “kind of simple types” (also known as * in the Haskell world).

fun @id(%X: !verona.type, %a: !verona.var<%X>) -> !verona.var<%X> {
  return %a: !verona.var<%X>
}

This second approach probably requires deep changes to MLIR itself, but could have, I believe, a multitude of use cases.

For example, an other concept in Verona is the support for “regions” (not the MLIR kind). A Verona region is, to some very broad approximation, an arena of objects. All objects allocated in that region are deallocated together. In the IR, we are interested in tracking the region of each pointer, for both type-checking and code generation purposes. If we had SSA values in types, we could express this as follows, !verona.pointer<%x> is a pointer to an object in the region %x.

fun @create(%x: !verona.region) -> !verona.pointer<%x> {
  %z = !verona.allocate %x : !verona.region -> !verona.pointer<%x>
  return %z
}

Unlike the polymorphic functions use case, !verona.region values do have a runtime value, and thus cannot be encoded with just attributes.

There’s of course many questions raised by this, including how to represent the types without giving up on immutability and uniquing, how substitution works, how to determine type equality, …

I’m curious whether such an extension to MLIR types has been considered before, whether this would be useful to anyone else, and whether there is an alternative approach I could take instead.

1 Like
2

Hello,

we are working on representing a functional pattern-based IR called RISE that we have earlier this year presented at an MLIR design meeting (slides and recording).

For this we are definitely interested in representing polymorphism directly in the IR!
In fact, we are also interested in representing different forms of dependent types where types might depend on values.

In RISE we already represent lambda expressions explicitly as an MLIR op.
We would like to model polymorphism quite naturally as another kind of a lambda expression that introduces a type parameter into the scope (and not a value parameter like ordinary lambda expression do).

So where currently the identity lambda expression for floating point values and applying it looks like:

%f32id = rise.lambda (%x: !rise.scalar<f32>) -> !rise.scalar<f32> {
  rise.return %x: !rise.scalar<f32>
}
%res = rise.apply %f32id %zero

A polymorphic identity lambda expression and its application could look like:

%tid = rise.typeLambda (%t: !rise.type) -> !rise.fun<rise.var<%t> -> rise.var<%t>> {
  %id = rise.lambda (%x: !rise.var<%t>) -> !rise.var<%t> {
    rise.return %x: !rise.var<%t>
  }
  rise.return %id
}
%f32id = rise.typeApply %tid #rise.scalar<f32>
%res = rise.apply %f32id %zero

@plietar One thing that I want to point out is that for your @id function besides needing to allow SSA values to appear in the type, you will also need types to be wrapped as SSA values so that they can be passed into the %X parameter.

3

Thanks, I’d found your paper already but not the presentation.

Having distinct type- and value- abstraction/application operations would work for us as well if it makes implementation easier, but I think supporting them as a single operation would be nicer.

@plietar One thing that I want to point out is that for your @id function besides needing to allow SSA values to appear in the type, you will also need types to be wrapped as SSA values so that they can be passed into the %X parameter.

Indeed. What I have in mind is a type_constant T operation that lifts a type attribute into an SSA value.
The value has type !verona.type<T>, in order to track that the type is known. This is in contrast with !verona.type which denotes an unknown type (more complicated notation could be used for eg. bounded polymorphism).

Given %t : !verona.type<T>, types T and !verona.var<%t> would be considered equivalent.

For example:

%t = verona.type_constant !i32 : !verona.type<i32>
%x = constant 12 : i32
%y = verona.call @id(%t, %x) : !verona.type<i32> -> i32 -> i32

The details of this could be implemented directly in the dialect (maybe with shared infrastructure if useful), with only the basic functionality in core MLIR.

4

Thinking about this a bit more, I think for polymorphic types you might not need to represent types as SSA values and still improve over your !verona.var<0> approach (that sounds a lot like De Bruijn indices).

Referring to my earlier example in the functional style you could do something like this:

%tid = rise.typeLambda #rise.tvar<T> : !rise.fun<rise.tvar<T> -> rise.tvar<T>> {
  %id = rise.lambda (%x: !rise.tvar<T>) -> !rise.tvar<T> {
    rise.return %x: !rise.var<T>
  }
  rise.return %id
} : !rise.forall<rise.tvar<T> . (rise.fun<rise.tvar<T> -> rise.tvar<T>>)>
%f32id = rise.typeApply %tid #rise.scalar<f32>
%res = rise.apply %f32id %zero

Here I am using attributes to introduce a type variable rise.tvar<T> (line 1) and apply it (line 7). I think with custom parsing you could use a name (such as T here).

The disadvantage of not using an SSA value for the type variable is that now you have to be careful with scoping yourself: nobody is ensuring that rise.tvar<T> is only used where it is in scope - you would get this for free when using an SSA value!

It is also worthwhile to think about what the type of the polymorphic @id function is that you provided as an example, as the type of this function now depends on the %X parameter that is a SSA value and not in scope outside the function.
In the example I provide here in this post, I introduce the rise.forall universally quantified type that captures the introduction and occurrence of the type variable in the function type explicitly. This is the standard way how once would model this in the polymorphic lambda calculus (a.k.a. SystemF).

5

@mehdi_amini @River707

Any idea on the complexity of making types more like values in MLIR, to allow type parametrization above?

6

I’d recommend taking a look at how the Swift SIL type system works. It uses the same structural representation approach as MLIR (uniqued types as “pointers”, etc). The way it does this is through type variables.

@rjmccall do you know of a good reference in the Swift docs that describes how this works?

7

Incidentally, Verona looks really interested. If you’re interested in linear types, it is worth researching what Swift is doing in this space and you might also find it interesting to learn how core features like definitive initialization, law of exclusivity, and other things like value semantics work together.

8

A lot of the SIL documentation is focused on the formal textual language rather than the in-memory representation.

The quick summary is that we use different types for polymorphic types in the abstract (e.g. in the expression of a polymorphic signature) vs. in the concrete (e.g. in the body of a polymorphic function). This allows the system to more easily detect substitution bugs where types are used in the wrong context, which are extremely common. Having types be able to point back to their generic context is also useful for this.

In general, polymorphism needs a lot of very careful design. MLIR might be in a hard position here if it wants to dictate the representation centrally because it will likely need to support the most general conceivable things, which will mean a lot of core properties of types will become more awkward to use. For example, type equality almost certainly needs to be a contextual predicate — representing two types with exactly the same type object means that they’re the same in all contexts, but representing them with different objects doesn’t mean they might not be the same in the right context. By “contextually” here I mean at least down to the basic block and maybe down to the instruction — you want to be able to refine generic types dynamically. This is something that SIL currently cannot do.

I have to say that relying on canonical uniqueness is probably the wrong approach in polymorphic type systems. We do this in Swift, and I feel it causes a lot of unnecessary problems. Treat it as a representational optimization rather than something necessary for correctness.

Giving types use-def chains is tempting. Often it’s important to be able to introduce new type variables in a function, if you have existentials or some similar feature. At the very least you’ll want to represent the dependencies of instructions that use that type on the instruction that introduced it. SIL does this with implicit extra use edges that can be carried by any instruction which uses a type other than as the type of an operand. It works relatively well but doesn’t feel clean. Making types actually values, I dunno.

9

Thanks for your input.

we use different types for polymorphic types in the abstract (e.g. in the expression of a polymorphic signature) vs. in the concrete (e.g. in the body of a polymorphic function).

This is something I’d be inclined to do as well. Inside a method body, the type parameter is represented as var<%t>. However, the type of the id function would be along the lines of forall<X. fun<var<X> -> var<X>>>. In memory, such abstract typ parameter could be represented with debuijn indices (or similar), so it would actually be forall<fun<var<0> -> var<0>>>. This is similar to the locally nameless approach found in a lot of literrature.

The abstract representation can already be expressed today quite easily, with some dialect specific type application operation. The only thing missing is, I think, a way to enforce that types are closed, but that is reasonably easy.

I’ll have a look at SIL. In the mean time, a representation I’ve had in mind is to separate the “type scheme” from its application to SSA values. In other words, a type foo<%x, %y, %z> would be represented by a pair (foo<#0, #1, #2>, [%x, %y, %z]). The first element of the pair is uniqued, but not the second element. The list of SSA values could be stored in the Operation/Block that defines the value with that type.

This gives the MLIR infrastructure visibility into the SSA values used by the type, and allows for easy replacement of one value by another. For instance it makes the following IR valid. When checking that !foo<%a> is the same type as !foo<%c>, we compare the scheme part using pointer equality and check that the applications match, by replacing %c by %a.

br ^bb(%a: !type, %b: !foo<%a>)

^bb(%c: !type, %d: !foo<%c>): ...

It does not however allow an SSA value to be replaced by a given type, that would be the responsibility of the dialect.

I have to say that relying on canonical uniqueness is probably the wrong approach in polymorphic type systems. We do this in Swift, and I feel it causes a lot of unnecessary problems. Treat it as a representational optimization rather than something necessary for correctness.

In Verona, we already need to support subtyping, so most operations use a custom verifier to compare types. Syntactic type equality, which is determined through pointer equality, is a special case of this subtyping. To support polymorphic types, we would have an additonal subtyping rule, that I vaguely referenced earlier, that looks like:

      %x : !verona.type<T>
---------------------------------------
     !verona.var<%x> <: T