Compilation to categories

[awf]

From @simonpj :

I believe [Elliot-style AD] is all working nicely now.

To run it, use Ksc.Pipeline.demoCL “ex10”, or choose your own example. It adds the Prelude automatically.

• Compare with
  Ksc.Pipeline.demoF BasicAD ["src/runtime/prelude.ks","test/ksc/ex10.ks"]
• All of the code is in CatLang.hs Remarkably brief considering it includes the language, its pretty printer,, conversions to and from KSC, and AD both forward and reverse.
• Incidentally the forward mode is tupled.
• ex10 demonstrates a major advantage of the Elliot approach: in rev$g, in reverse mode BasicAD we currently get two calls to rev$f. They have different arguments and cannot be CSE’d. But with the Elliott approach we can the argument before calling rev$f, so we get just one call.

I’d be interested in any observations you might have.

Next: splitting for reverse mode.

Some other points:

• I realise that ANF uses a monad to generate uniques; but that doesn’t guarantee uniqueness, depending on what names are already in use. I changed it to use the form “t$nn”, to generate names of a form that only it uses.
• I realise that the primitive ‘add’ function (see Prim.hs) claims to add T+tangent(T). But actually it’s only ever used to add tangent(T)+tangent(T)! Do you agree? When we are adding T’s in ksc code we use add@ff etc.

[awf]

I get a failure on generated_deltavecs:

Prelude AD> Ksc.Pipeline.demoCL "generated_deltavecs"
...
----------------------------
toCLDefs
----------------------------

def matrix_multiply Vec Float (_t1 : (Vec Vec Float, Vec Float))
  = let { [w] = get$1$2 . get$1$1 }
    let { [i, m, a, g, e] = get$2$2 . get$2$2 }
    let { [o] = size . get$2$3 }
    let { [k] =
            size
            . index
            . Prune[3] [ 0, get$1$1 ] }
    Prune[1, 2, 3, 4]
      (build get$2$4
             (\oi. sumbuild*** Exception: toCLExpr
  (lam (ki : Integer)
       mul@ff(index(ki, index(oi, w)), index(ki, image)))
CallStack (from HasCallStack):
  error, called at src/ksc\Lang.hs:935:20 in main:Lang
  pprPanic, called at src/ksc\CatLang.hs:136:30 in main:CatLang

[toelli-msft]

This is great! One important thing to do is to make sure it works all the way through Cgen. Another thing is that there are lots of small administrative changes that should be split off and submitted in separate commits (if not a separate PR). I'm happy to take action on those.

[awf]

And re tangent types:

[Monday 2:31 PM] Tom Ellis
Regarding add (called ts_add on master, I believe), it is used at type (T, tangent(T)) in gmm.ks to add wishart to dwishart (the two are not the same).
Arguably having separate functions for (T, T) -> T and (T, tangent(T)) -> T is more mathematically faithful, but seems redundant.

[Monday 5:05 PM] Andrew Fitzgibbon
Re ts_add -- check #199 , and feel free to comment on the issue if you think it needs expanding...

[Monday 5:09 PM] Simon Peyton Jones
Yes, but we also need (tangent(T),tangent(T)) -> tangent(T). I'm a bit uncomfortable with using the same function for both. But maybe no urgent need to change.

[Monday 5:18 PM] Tom Ellis
Simon Peyton Jones tangent(tangent(T)) is always the same as tangent(T), so I think we are safe.

[Monday 8:58 PM] Andrew Fitzgibbon
Someone will want a proof that "tangent(tangent(T)) is always the same as tangent(T)" which we can't quite assemble because we were vague about what dT is. But if we talk about it as "the type of a small change in a value of T", then we get that
a small change in a small change must be a small change...
Or "the space dT" is already a vector space,
Or "stuff that's linear in dx must already be linear in ddx,
Or.... id(x +(dx + ddx)) = id(x) + id'(x)*(dx + ddx) = yawn....
Or...
Let's wait til we find a counterexample, which I'm sure we won't.

[simonpj]

Let's wait til we find a counterexample, which I'm sure we won't.

OK. But it would really help to have a Note with this code in Prim.hs

      ("add"      , TypeTuple [t, dt]                    ) -> if dt == tangentType t
                                                                then Just t
                                                                else Nothing

making some of these points. I tripped up on this when thinking "is it ok to add two tangent types with 'add'"?

[awf]

making some of these points. I tripped up on this when thinking "is it ok to add two tangent types with 'add'"?

Aha, and that's exactly where we are precise about tangentType -- we can prove there that ``tangentType . tangentType = tangentType`

[simonpj]

I get a failure on generated_deltavecs:

Yes, it knows about build, but not about sumbuild. In the old AD setup we had

-- I'm not very happy about this rule, which effectively
-- undoes sum (build e) --> sumbuild e
gradE adp s (Call f (Tuple [n, body]))
  | f `isThePrimFun` "sumbuild"
  = gradE adp s (pSum (pBuild n body))

I suppose I can add that to the conversion from ksc -> categories

[toelli-msft]

Hi @simonpj, I'm seeing a strange error. Sometimes variable names come out as numbers without any alphabetic prefix which is invalid in our C++ and Futhark backends, as one might expect. I don't have any idea how to track this down. Do you have a suggestion of where to look?

ghci -isrc/ksc -e 'doall 999 ["src/runtime/prelude"] "test/ksc/generated_deltavecs"' src/ksc/Ksc/Pipeline.hs

let { 19 : Vec Vec Vec Float = get$1$2(_t1) }
let { 20 : Vec Vec Float = index(1, 19) }

[toelli-msft]

The command line to reproduce is

ghci -isrc/ksc -e 'doall 999 ["src/runtime/prelude"] "test/ksc/generated_deltavecs"' src/ksc/Ksc/Pipeline.hs

[simonpj]

Sometimes variable names come out as numbers without any alphabetic prefix

Oh dear! Apparently I introduced a bug in notInScope. I've pushed a fix to spj/conal:

commit 9caa7dcce4de42c491b2af5b9253fd91c515d4f0 (HEAD -> spj/conal, origin/spj/conal)
Author: Simon Peyton Jones <simonpj@microsoft.com>
Date:   Tue Apr 7 23:50:57 2020 +0100

    Fix a nasty bug in 'notInScope'
    
    This bug resulted in us getting variable names
    like "18", which was is very confusing.

I hope that unsticks you.

[toelli-msft]

Thanks Simon. The change in behaviour to notInScope is that it now does not have an underscore between variable name and number, whereas previously it had one. Is this for cosmetic reasons or does that underscore interact badly with other parts of the code? Either way is fine by me, but if it's the latter reason we ought to document it thoroughly.

[awf]

Can we use something noisier e.g. ANF$<N>?

[simonpj]

Is this for cosmetic reasons or does that underscore interact badly with other parts of the code?

Purely cosmetic! I just got tired of seeing lots of underscores :-). No more than that.

[toelli-msft]

Thanks Simon, good to know.

[awf]

Sorry, I said ANF$ instead of ksc$anf -- I had thought to do argVar as Prim$argVar

[toelli-msft]

@awf, did you mean the variables in the below? If so, do you prefer "ANF$..." or "ksc$..."

knossos-ksc/src/ksc/ANF.hs

Lines 170 to 176 in 8328a32

	v = mkVar ("t$" ++ show u)
	instance GenBndr Typed where
	mkNewVar u e = (tv, tv)
	where
	-- See Note [Fresh names in ANF]
	tv = mkTVar (typeof e) ("t$" ++ show u)

[toelli-msft]

@simonpj: I'm curious about this change. Is it intended to make the pretty-printed syntax less noisy? The change actually breaks the Futhark code generator (because it uses the pretty-printing machinery here). Do you mind if I revert this change?

[simonpj]

Hmm. Maybe I made a mistake. Where the precedence is equal we probably do want parens. But what about

  let x = e in let y = e2 in ...

we don't really want parens around that inner let.

What's the example that breaks? Just so we can see.

We can revert and lets see if strange things happen.

[toelli-msft]

One thing that breaks is the Futhark code generated for the below code from prelude.ks. With > in the precedence the condition doesn't fire and the Futhark comes out looking like

neg div (mul x1 d_ddiv) (mul x2 x2)

instead of

neg (div (mul x1 d_ddiv) (mul x2 x2))

https://github.com/microsoft/knossos-ksc/blob/master/src/runtime/prelude.ks#L87-L88

[toelli-msft]

I'm hitting this a lot in the test suite. Is it something that is yet to be implemented?

[simonpj]

On what expressions?

[simonpj]

The Let case looks odd. I think it should be

to_cl_expr Pruned env (Let tv rhs body)
  = CLLet (tVarName tv)
          (toCLExpr env rhs)
          (toCLExpr (tv:env) body)

[toelli-msft]

On what expressions?

On any assert or any lam that's not part of a build (i.e. the ones that are part of sumbuild or fold).

[toelli-msft]

OK, I fixed the let and pushed, along with a few other small changes that shouldn't interfere with anything you're doing.

[simonpj]

pushed where? spj/conal?

[simonpj]

On any assert or any lam that's not part of a build (i.e. the ones that are part of sumbuild or fold).

Ah yes, well, we can't deal with lambdas in general. I have not yet done sumbuild or fold. For sumbuild in AD.hs we just expand it back to sum of build; I suppose we can do that here too.

Fold is still to-do.

I have not thought about assert. Not sure if it's worth trying to round trip it through CL. Do you have a view?

[toelli-msft]

pushed where? spj/conal?

Yes

[toelli-msft]

Thanks for pushing the rebased version, Simon. For the record the old branch was at 62f3156.

[toelli-msft]

Is this change actually needed? It doesn't seem so, given the change to pSel above. pSel itself now checks whether it's being applied at a "tuple size" of 1 and if so assumes it was given the actual value, not a singleton tuple containing the value.

If we can do without this change I would prefer it. I'm rather wary of treating T and TypeTuple [T] as the same thing.

[toelli-msft]

@simonpj: demoCL "test3" gives an fsAdExpr error that I don't understand. Is this expected?

[toelli-msft]

I've succeeded in splitting this up into smaller pieces so that I can understand it. It's on origin/toelli/elliot-rebased. @simonpj: you might want to reset your branch to that one and force push. The two branches are identical, with the exception of some small new fixes on mine.

I'm comfortable with everything that's going on in the PR except treating 1-tuples as though they are the same as the thing they wrap. I worry that this may cause us headaches down the line. In particular, our backends do not make this conflation. Is it really necessary? It's the kind of thing I would be keen not to invoke unless there's no alternative.

[awf]

I'm comfortable with everything that's going on in the PR except treating 1-tuples as though they are the same as the thing they wrap. I worry that this may cause us headaches down the line. In particular, our backends do not make this conflation. Is it really necessary? It's the kind of thing I would be keen not to invoke unless there's no alternative.

Agree that there's an important question here. I believe that it will probably force us to say that semantically x == (tuple x), and that that pretty much means we will end up having to 'flatten' all tuples, so (tuple (tuple 1 2) (tuple 3 4)) == (tuple 1 2 3 4).
This has no speed implication -- in Knossos we always reserve the right to lay things out in memory as we wish, but it feels as if it may have some other infelicity.

So yes, it would be nice to see if we can work with x != (tuple x) unless it's really messy. IMO it should keep things clean.

[simonpj]

I'm comfortable with everything that's going on in the PR except treating 1-tuples as though they are the same as the thing they wrap.

I think it's entirely unnecessary. I think at one point I explored that route and then discarded it but failed to remove the code.

Thanks for splitting up. I'll adopt your patches.

[awf]

Where are we on this?

[simonpj]

branch spj/catlang has a pretty solid implementation. It converts to and from from CatLang, and does various forms of fwd and reverse AD. It's the only implementation we have that does "split reverse" AD, which is for the first time fully compositional. It's very nice!

Still needed if we want to push this further

• Sums (for conditionals)
• Fold (I think Tom is going to add this). We have build already.
• Recursion
• Higher order

Tom has yet a new approach (with explicit fork and join constructs) which may be better still, so it's not clear to me whether we should push the the categorical approach further or purse Tom's, or both in parallel.