User guided specialisation stage #252

Oli-Ar · 2024-12-17T05:48:10Z

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

LPTK · 2024-12-18T03:46:08Z

hkmc2/shared/src/test/mlscript/specialisation/Basic.mls

@@ -0,0 +1,87 @@
+
+
+:spt


It would be helpful to also show the JS code, which is much easier to read.

Co-authored-by: Anson Yeung <[email protected]>

Ported specialisation keyword from previous commits to new HEAD, as it was easier than rebasing, as other components had significantly changed.

As part of moving specialiser code to the new head rewrote logic to pass the specialisation flags through the compiler - now has a much smaller footprint.

LPTK · 2025-02-26T02:52:33Z

hkmc2/shared/src/main/scala/hkmc2/semantics/Specialiser.scala

+    case Primitive(name: Str)
+    case Function(lhs: SimpleType, rhs: SimpleType)
+    case Record(fields: Ls[(Str, SimpleType)])
+    case ClassType(sym: ClassSymbol, supers: Ls[SimpleType], members: Map[Str, SimpleType])


This is more like ClassInfo (which is not a SimpleType). The types of actual class instances should be a different constructor (which points back to the ClassInfo).

LPTK · 2025-02-26T03:12:04Z

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 · 2025-02-26T03:49:11Z

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!

I'm not sure this is desired behaviour but I added this in case a user specialises a parameter which takes in a lambda, which currently at least, requires checking if a value is of type function

Added small features: specialisation in specialised functions being more specific, lambda support Fixed many bugs including not properly hoisting functions, issues with application specialisation, issues in the constraining of some types

Oli-Ar · 2025-03-28T17:24:28Z

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 · 2025-03-28T17:26:38Z

hkmc2/shared/src/main/scala/hkmc2/codegen/js/JSBuilder.scala

@@ -307,6 +307,7 @@ class JSBuilder(using TL, State, Ctx) extends CodeBuilder:
          case Elaborator.ctx.builtins.Num => doc"typeof $sd === 'number'"
          case Elaborator.ctx.builtins.Bool => doc"typeof $sd === 'boolean'"
          case Elaborator.ctx.builtins.Int => doc"globalThis.Number.isInteger($sd)"
+          case Elaborator.ctx.builtins.Function => doc"typeof $sd === 'function'"


I added this as I had a test where I specialise a parameter which I pass a lambda to, so when I generate the pattern matching for the application I needed to be able to match the lambda, but I'd just like to check that this is okay to do since I'm not very familiar with the codegen side of mlscript.

hkmc2/shared/src/test/mlscript/specialiser/Basic.mls

LPTK · 2025-03-29T04:41:52Z

hkmc2DiffTests/src/test/scala/hkmc2/MLsDiffMaker.scala

-    val checker = new semantics.SimpleSub(curCtx, stl)
-    val spt = checker.analyzeTermTypes(e).asInstanceOf[semantics.Term.Blk]
+    val checker = new semantics.Specialiser(curCtx, stl)
+    val spt = checker.topLevel(e).asInstanceOf[semantics.Term.Blk]


Why is the cast needed here but not in MLsCompiler? And if we really need a Blk, couldn't you just make topLevel take and return one?

LPTK reviewed Dec 18, 2024

View reviewed changes

Oli-Ar force-pushed the hkmc2-specialisation-wip branch from 84db949 to 2a42dd2 Compare January 24, 2025 06:30

CAG2Mark and others added 5 commits February 14, 2025 05:42

Miscellaneous Fixes (hkust-taco#280)

322b68b

Co-authored-by: Anson Yeung <[email protected]>

Ported specialisation keyword to HEAD

740dacb

Ported specialisation keyword from previous commits to new HEAD, as it was easier than rebasing, as other components had significantly changed.

Added logic to pass specialisation flags through elaborator

1796d7d

As part of moving specialiser code to the new head rewrote logic to pass the specialisation flags through the compiler - now has a much smaller footprint.

Brought back specialiser logic

af9ac51

Pushed majority of logic for the simple sub algorithm

b91aa15

Oli-Ar force-pushed the hkmc2-specialisation-wip branch from 2a42dd2 to b91aa15 Compare February 26, 2025 01:25

LPTK reviewed Feb 26, 2025

View reviewed changes

Oli-Ar added 11 commits March 6, 2025 14:52

Pulled class info out from SimpleType to new class

423fafe

Fixed type inference for classes and some types of functions

7189d2c

Fixed relationship

3211cec

Code cleanup and minor fixes

138ecca

Added type inference for statements

9fa8b52

Delayed processing of TermDefs fixing some class of inference bugs

7ae6090

Specialisation inference

89ca14c

Generating specialised functions

f6bd7bb

Arguments to specialised parameters are now extracted

4aa2131

Cleaned out tests

d0ae2e6

Specialiser type pattern matching added

73b350a

Oli-Ar force-pushed the hkmc2-specialisation-wip branch from 231334e to 73b350a Compare March 14, 2025 05:14

Oli-Ar added 6 commits March 29, 2025 00:52

Specialiser type pattern matching added

7f4b2a5

Renaming

2b235bb

Handling specialisation inside specialised functions better

952e08f

Improved specificity of specialiser

e57dfcd

Added typeof function check support

ed681da

I'm not sure this is desired behaviour but I added this in case a user specialises a parameter which takes in a lambda, which currently at least, requires checking if a value is of type function

Small features & Bug Fixes

0a795a7

Added small features: specialisation in specialised functions being more specific, lambda support Fixed many bugs including not properly hoisting functions, issues with application specialisation, issues in the constraining of some types

Oli-Ar commented Mar 28, 2025

View reviewed changes

LPTK reviewed Mar 29, 2025

View reviewed changes

Oli-Ar added 2 commits March 29, 2025 18:38

Re-added removed test

c08a954

WIP: Changed specilisation to work on specialised parameter vectors

ae4dfb1

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

User guided specialisation stage #252

User guided specialisation stage #252

Oli-Ar commented Dec 17, 2024

LPTK Dec 18, 2024 •

edited

Loading

LPTK Feb 26, 2025

LPTK commented Feb 26, 2025 •

edited

Loading

LPTK commented Feb 26, 2025 •

edited

Loading

Oli-Ar commented Mar 28, 2025

Oli-Ar Mar 28, 2025

LPTK Mar 29, 2025

User guided specialisation stage #252

Are you sure you want to change the base?

User guided specialisation stage #252

Conversation

Oli-Ar commented Dec 17, 2024

LPTK Dec 18, 2024 • edited Loading

Choose a reason for hiding this comment

LPTK Feb 26, 2025

Choose a reason for hiding this comment

LPTK commented Feb 26, 2025 • edited Loading

LPTK commented Feb 26, 2025 • edited Loading

Oli-Ar commented Mar 28, 2025

Oli-Ar Mar 28, 2025

Choose a reason for hiding this comment

LPTK Mar 29, 2025

Choose a reason for hiding this comment

LPTK Dec 18, 2024 •

edited

Loading

LPTK commented Feb 26, 2025 •

edited

Loading

LPTK commented Feb 26, 2025 •

edited

Loading