Synthesizing `RuntimeRep`-indexed type class instances in a GHC plugin to simulate monomorphization

Learn
1

Motiviation

GHC Haskell has a nice runtime representation system, centered around RuntimeRep and TYPE :: RuntimeRep -> * in GHC.Exts, that enables representation-polymorphic abstractions. Unfortunately, this control over representations came late in Haskell’s life, so most of the existing abstractions in the ecosystem (e.g., almost all of the standard type classes in base) work only over TYPE LiftedRep (i.e., *). I’m interested in creating an alternative foundational library (most likely largely incompatible with the existing ecosystem) to reduce the pain of using unlifted types, for two reasons:

  • To achieve better performance (e.g., avoiding allocation with unboxed types).
  • To program with strict evaluation (controversial, I know). (Why not just use strictness annotations (or Strict{,Data})? Because, unlike switching from TYPE LiftedRep to TYPE UnliftedRep, you can’t apply strictness annotations to categorically rule-out bottoms for data types, so certain laws on standard type classes don’t hold. E.g., that’s why Data.Strict.Maybe from the strict package isn’t an applicative functor; if Applicative somehow supported functors with domain TYPE UnliftedRep, and if Data.Strict.Maybe abstracted over TYPE UnliftedRep (instead of over TYPE LiftedRep with a strict annotation on the Just constructor), then we could have Applicative Data.Strict.Maybe.)

Problem

There are a number of obstacles to creating representation-polymorphic abstractions in Haskell, but the one I want to focus on here is that representations must be “concrete” in certain circumstances (details in Downen et al., “Kinds are Calling Conventions”). I think the main two restrictions are as follows:

  • You can’t form a lambda whose parameter has non-concrete representation. E.g., this doesn’t compile:
    foo :: forall (r :: RuntimeRep) (a :: TYPE r). a -> Int
foo _ = 0
  • You can’t form an application whose argument has non-concrete representation. E.g., this doesn’t compile:
    bar :: forall (r :: RuntimeRep) (a :: TYPE r) b. (a -> b) -> (() -> a) -> b
bar f g = f (g ())

Solution Idea

Inspired by Edward Kmett’s unboxed experiment and Conal Elliott’s work on “Compiling to Categories”, I’m trying to get around the concreteness restrictions on representation polymorphism with RuntimeRep-indexed type classes.

For example, in conjunction with DefaultSignatures, Edward applies this idea to define representation-polymorphic type classes with default method implementations. Look at how his version of Eq utilizes an extra RuntimeRep-indexed class EqRep: the naive default implementation of (/=) as a /= b = not (a == b) does not satisfy the concreteness constraint on lambdas, because a and b have non-concrete representation.

The downside of Edward’s approach is that we must implement each RuntimeRep-indexed class (EqRep, etc.) for each desired RuntimeRep in advance. Since there are infinitely-many RuntimeReps (thanks to the recursive TupleRep and SumRep constructors), this is a losing battle. (Of course, you can automate the instance generation, as Edward does for unboxed tuples and sums of arity 2. See the cabal file.)

It would be much nicer if the compiler would automatically generate the required instances of RuntimeRep-indexed classes on demand. So we would, for instance, (1) write Eq and EqRep exactly as Edward has, (2) write a RuntimeRep-uniform implementation “schema” for EqRep once, then (3) rely on the compiler to synthesize an instance of EqRep from the schema whenever EqRep r arises as a constraint with concrete r :: RuntimeRep.

I think this is possible with a GHC type checker plugin. I’m not very familiar with GHC internals in general, but I think it could work something like this:

  • User POV:

    1. The user writes a type class MyClass and marks it as “RuntimeRep-monomorphizable” with an annotation (ANN pragma). (If it helps, we could restrict the form of these type classes somewhat; e.g., class MyClass (r :: RuntimeRep) where ... probably covers most cases.)
    2. The user writes an implementation schema for each MyClass method. I’m not sure how this should work; it’s probably the hard part.
    3. The user freely uses MyClass r constraints.

  • Plugin POV:

    1. The plugin looks for classes annotated as RuntimeRep-monomorphizable and finds MyClass (in this example, anyway; in general, there can be many).
    2. Each time GHC can’t itself solve a MyClass r constraint, it consults the plugin. If r is a concrete representation, the plugin synthesizes a instance of MyClass r and returns it as evidence. (Hopefully, GHC only consults type checker plugins for constraints it can’t solve itself; otherwise, the plugin has to check if a suitable instance already exists.) (If r is not a concrete representation, the plugin does nothing.)

Request for Comments/Hints

Is this all a terrible idea? Is there a better way to do it?

If the suggested solution isn’t too bad, any hints (existing plugins to study, relevant parts of the GHC API to apply, etc.) would be greatly appreciated.

5 Likes
2

A good background paper is Levity polymorphism. Worth reading if you are digging into this.

I can’t comment about compiler support until I understand the proposed solution, the one you attribute to Ed K. The problem is a fundamental one. If we have

f :: Num a => a -> a
f x = ...lots of code...

Now if we want to apply f to an unboxed int, thus (f @Double# 19882.234#) we are stuck. The gobs of compiled matchine code for f manipulates boxed pointers, but we want compiled code that manipulates unboxed floats. Different machine instructions!

The only solutions I know are

  • Box the Double# before calling f – but that negates the whole point, which is to manipulate unboxed values exclusively.
  • Make a copy of the source code for f, specifically for Double#, and compile that. Now we have two versions of f. That’s fine but it’s hard to know which types to specialise f for (there are an infinite number of unboxed types). .NET does this, but GHC currently makes no attempt to do so.

It’d be worth having a clear design before thinking about implementation strategies.

3

I’ve been thinking about this recently and I am exploring this option. Instead of generating all possible options in advance, we could generate the code only at call sides.

A further problem, then, is that the function may be called at the same type in different places. You’d want to share specializations. I want to see if we could do two passes:

  1. Compile (maybe just up to type checking?) the package while keeping the runtime rep polymorphic functions as holes (in the sense of backpack). Also collect info on all calls.
  2. Generate a module with all specialisations for the calls we found and instantiate the hole with that module.

This does only work if we can rule out existential quantification and polymorphic recursion. To address fhat, I am also thinking about extending the type system to let us specify that a certain type variable must be monomorphised. That could be used to rule out problematic polymorphism.

This also could be nice for certain other type classes like PrimMonad or even Monad in general, which often see huge benefits from specialization, so you want to specialise those even if it leads to larger binaries and longer compile times.

4

It is possible to do partial specialisation. See for example Sixten.

A representation-polymorphic function accepts a witness of the representation, such as a log2 size (alignment, &c.), and uses that to deal with any values of that representation at runtime. This is a middle ground that allows separate compilation and polymorphic recursion without giving up unboxing altogether.

Likewise, in an existential, you would store this witness, in much the same way as a typeclass dictionary. Perhaps it could even reuse a lot of the same machinery. For example, when you see forall (r :: RuntimeRep), introduce a constraint IsRuntimeRep# r, whose dictionary is a Word# holding the size.

5

we could generate the code only at call sides.

type classes like PrimMonad or even Monad in general

This sounds like staging to me!

6

Are there staging systems that can do the code sharing I’m proposing?