mlscript icon indicating copy to clipboard operation
mlscript copied to clipboard

User guided specialisation stage

Open Oli-Ar opened this issue 1 year ago • 3 comments

Draft PR for implementing a user-guided specialisation pass in the mlscript compiler, very much still WIP

Oli-Ar avatar Dec 17 '24 05:12 Oli-Ar

TODO: Write specification/architecture of algorithm first.

You need to attach the flows of normal parameters to the flows of specialized ones.

Example:

fun foo(spec x: ?a, p1: ?p1, p2: ?p2) = ...

fun bar(x: ?b, arg: ?arg) =
  foo(x, arg, 123)
// At this point, we know we'll need to specialize this call to `foo` for each new concrete LB coming into ?b
// Each of these specializations will need to create the flows `?arg <: ?p1` and `123 <: ?p2`.

Another way of thinking about this: specialized functions are polymorphic, but the polymorphic instances are indexed not by the individual call sites but by the list of concrete shapes flowing into their specialized parameters.

fun foo(spec x, spec y, p1, p2) = ... constraint ?p1 <: ?p2 -> ?result
foo: (x, y) => forall p1, p2, result where { p1 <: p2 -> result }. T

In this approach, we remember whatever flows into (p1, p2) by attaching it (indexing it) to what flows into (x, y). Only when the (x, y) flow becomes concrete do we look up the specialization table: then, if there is not yet an entry, we instantiate the polymorphic type scheme; afterwards, we connect the stashes (p1, p2) flows with the corresponding instantiated signature.

Note: Similarly, ClassInfos themselves should be indexed on their spec params.

LPTK avatar Feb 26 '25 03:02 LPTK

Another interesting thing: the calls to specialized functions happening in non-specialized contexts may have to be turned into pattern matches.

Reusing the previous example:

fun foo(spec x, p1, p2) = ...

fun bar(x, arg) =
  foo(x, arg, 123)

bar(Int, 0)
bar(Str, true)

This should basically specialize into:

fun foo_Int(x, p1, p2) = ...
fun foo_Str(x, p1, p2) = ...

fun bar(x, arg) =
  if x is
    Int then foo_Int(x, arg, 123)
    Str then foo_Str(x, arg, 123)

bar(Int, 0)
bar(Str, true)

In fact, at least conceptually, we should probably always generate pattern matching by default and possibly optimize it later. So even a direct call like

foo(42, true, 0)

should be turned into

let $arg = 42
if $arg is
  Int then foo_Int($arg, true, 0)

which will be optimized the the expected

let $arg = 42
foo_Int($arg, true, 0)

So this looks pretty similar to the earlier defunctionalization/demethodization project!

LPTK avatar Feb 26 '25 03:02 LPTK

The approach I took to specialisation works by taking the SimpleType that a parameter would be and then wrapping it in a SpecialisableType SimpleType, which stores the underlying SimpleType, along with a SpecPoint instance, which stores more details about the specialisation.

Specifically I store any known concrete types, any variables that flow into it, and any flows out of it into other SpecialisableTypes. With this, I can do one pass of the program, and constrain the types, and while I constrain I fill in the rest of the information about the specialised parameter. Then with all of the SpecialisableTypes I can figure out what each function should be specialised to. Then on the next pass I case use that table to actually specialise the functions and the applications of the functions.

The reason I did it this way over the linked list was just that I was struggling to get the linked list method working, and I had this idea so tried it quickly and it worked better than anything I'd managed with the first approach so I decided to keep at it for a bit longer.

Oli-Ar avatar Mar 28 '25 17:03 Oli-Ar

This is unlikely to progress further, so I marked it to revisit later and will close it for now.

LPTK avatar Aug 12 '25 04:08 LPTK