Adding polymorphism to ksc

[simonpj]

The big changes are these:

• A Type can now contain a type variable

• A top-level Def binds type variables, def_qvars :: [TyVar]

• The concrete syntax and parser supports polymorphic
  definitions. E.g. here is the identity function:
  (def id [a] a ( x : a ) x)

• In a Call, the TFun has a field tf_targs :: [Type], which
  are the types at which the function is instantiated. Of course
  tf_targs is usually empty.

• A BaseUserFun is contains an ArgTypeDescriptor, rather than
  a mere Type (as previosly):
  data ArgTypeDescriptor = Mono Type | Poly

  So we can can have BaseUserFuns like [f Int], [f Bool],
  but also [f poly].

• The GblSymTab can contain

  • either a polymorphic definition
  • or multiple monomorphic ones

  Its structure reflects this:
  type GblSymTab = M.Map (DerivedFun BaseUserFunName)
  FunBindings
  data FunBindings = PolyBind TDef
  | MonoBinds (M.Map Type TDef)

Everything else is driven by these changes in the data types.

[simonpj]

@toelli-msft, we have polymorphism!

Would you like to take a look? I should write a Note, but I don't think you'll find any surprises.

There is a tiny test in poly1.ks, but perhaps you can try others?

I have done nothing about back ends.

[toelli-msft]

Great! It would be good to elaborate the benefits. Off the top of my head:

1. In general writing polymorphic programs is a useful thing to be able to (reduces code duplication)
2. Many prims can become just normal functions (with polymorphic types) (which? we should try to list them)
3. We can implement program transformations that require polymorphism (such as AD transformations -- although I think the ones that we have recently been considering do not require polymorphism in an essential way)

[simonpj]

It would be good to elaborate the benefits

Yes indeed. Where would be a good place to enumerate these benefits? Buried in an MR comment stream doesn't seem the right place.

We can implement program transformations that require polymorphism

For me this is a driver -- it'll let us implement reverse-mode AD via the "backpropagator" method, which we've been discussing with Neel and Faustyna. My plan is to do a proof-of-concept AD transformation using this idea.

[awf]

4.GNN can process polymorphic definitions, instead of running over many different implementations.

[toelli-msft]

Where would be a good place to enumerate these benefits? Buried in an MR comment stream doesn't seem the right place.

I think starting in an MR comment stream is fine. Pertinent points get transferred somewhere more permanent such as the OneNote. Alternatively the OneNote is a good place to start.

it'll let us implement reverse-mode AD via the "backpropagator" method, which we've been discussing with Neel and Faustyna. My plan is to do a proof-of-concept AD transformation using this idea.

Just as a point of information, I don't think polymorphism is required to get that AD transformation to work. It could be done with a transformation that takes a type as a parameter. On the other hand it's much easier to do with polymorphism otherwise one has to conjure up names for functions that differ only in that type parameter.

[toelli-msft]

(And @simonpj I'm sorry that I haven't had a chance to look at this yet, but thanks for the walkthrough yesterday and will get to it when I can.)

[ryotatomioka]

[This may be just my ignorance] I have seen challenge in AD due to polymorphism in the rank of inputs to ONNX's MatMul Op. That is, basically you can call MatMul on two n-dimensional tensors (n >= 2) and you can get broadcasting automatically.

See OnnxRuntime's AD implementation here:
https://github.com/microsoft/onnxruntime/blob/16ca7677e6aa4c4eed1fe0fc1761bae55e212cc4/orttraining/orttraining/core/graph/gradient_builder.cc#L127

I am curious if polymorphism in KSC will cause similar headache for us or we have a different plan (for example, we have a better type inference or other mechanism to resolve type before AD) so we are ok.

More concretely, the backward for

(def add (Vec Float) ((a : (Vec Float)) (b : Float)) (build (size a) (lam (i : Integer) (add (index a i) b))))

would be

(def rev$add (Tuple (Vec Float) Float) ((a : (Vec Float)) (b : Float) (dout : (Vec Float)))
  (tuple dout (sum dout)))

which is different from the backward for

(def add (Vec Float) ((a : (Vec Float)) (b : (Vec Float)))...)

[simonpj]

I think that your two add functions, namely

(def add (Vec Float) ((a : (Vec Float)) (b : Float)) ...)
(def add (Vec Float) ((a : (Vec Float)) (b : (Vec Float)))...)

are simply two entirely separate functions, just shorthand for

(def [add (Vec Float,Float)]          (Vec Float) ((a : (Vec Float)) (b : Float)) ...)
(def (add (Vec Float, Vec Float)]  (Vec Float) ((a : (Vec Float)) (b : (Vec Float)))...)

As a programming convenience we allow you to call them both "add", but that's rapidly resolved into two completely distinct functions [add (Vec Float,Float)] and [add (Vec Float,Vec Float)]

No polymorphism here!

[awf]

2. Many prims can become just normal functions (with polymorphic types) (which? we should try to list them)

@simonpj - can you take a look at Prim.hs and see what actually gets removed?