Add ChainRules

[oxinabox]

This replaces #291

The bits from that OP that still matter

Step 1) Change Zygote to check for chainrules before doing its normal stuff,
and adapt the stuff it gets back from chainrules to play nice with Zygote's expectations

Step 2) adapt Zygote more deeply, so it can take full advantage of thunks etc.

This PR is Step 1.

TODO: workout why this seems to segfault for me.

[oxinabox]

@MikeInnes any idea why this segfaults?

[MikeInnes]

It could be a Zygote compiler issue, but given that you haven't touched that, the second most likely option is that it's just a stackoverflow.

[oxinabox]

What makes you think stackoverflow?
I don't see anything in the logs pointing towards that

[MikeInnes]

IME, Julia tends to segfault due to stack overflows. Doesn't happen every time, but often.

[oxinabox]

It works!

1 of the blackslists can be removed immediately with JuliaDiff/ChainRules.jl#124

[MikeInnes]

I'm concerned about the performance implications of this, in its current form. In particular we have to be pretty careful to manage the call stack; recursing many times into _pullback will likely trigger recursion limiting heuristics that cause type inference to give up. (This is particularly pernicious because it doesn't show up in @code_typed.)

Also, presumably many of Zygote's current adjoints are now redundant, given this PR. It'd be nice to remove those to show the overall impact on Zygote's codebase (hopefully a significant net improvement).

[oxinabox]

I'm concerned about the performance implications of this, in its current form. In particular we have to be pretty careful to manage the call stack; recursing many times into _pullback will likely trigger recursion limiting heuristics that cause type inference to give up. (This is particularly pernicious because it doesn't show up in @code_typed.)

Which bits do you think are adding levels?
We don't recurse into _pullback anywhere in this PR.

How every you count it, adding deeper integation and usiing ChainRules types (rather than zexterning everythiing) would cut that to 1 I think,
and also give benifit of not unthunking everything.

Also, presumably many of Zygote's current adjoints are now redundant, given this PR. It'd be nice to remove those to show the overall impact on Zygote's codebase (hopefully a significant net improvement).

Easiest way to do would be to add convience functions to ChainRulesCore
for importing Zygote Rules's @adjoint defintions,
and then transferiing basically everything over.
Checking as one goes.
Then look at the method overwrite warnings.
JuliaDiff/ChainRulesCore.jl#44

I'ld rather not be doing that alone, and I'ld rather not have that blocking this PR.
I think it can be done in a follow up.
But if we need to use ChainRules types deeply then maybe not.

[MikeInnes]

Which bits do you think are adding levels? We don't recurse into _pullback anywhere in this PR.

If the user-defined call tree is foo -> bar -> baz, the transformed tree under this PR is _pullback(foo) -> _pullback_sct(foo) -> _pullback(bar) -> _pullback_sct(bar) -> _pullback(baz) -> _pullback_sct(baz) (since _pullback_sct recursively calls _pullback for each call inside the function it transforms).

Previously, it would have been _pullback_sct(foo) -> _pullback_sct(bar) -> _pullback_sct(baz). Because we have a type inference hack that lets us trick the compiler into seeing foo -> bar -> baz again here, this doesn't kill the compiler. Whereas now the callstack has three recursions into the same function _pullback, which can cause issues.

[oxinabox]

Because we have a type inference hack that lets us trick the compiler into seeing foo -> bar -> baz again here, this doesn't kill the compiler.

Where can I learn about this type inference hack?
And is the solution to just inline _pullback_sct back into _pullback by hand?

[MikeInnes]

I don't think it's documented, but you just set a flag on the codeinfo, as here.

I think manual inlining / having one uber-function for _pullback is probably the way to go (it can still be factored of course, as long as any function calls are compile-time).

[oxinabox]

Re inference
I am starting to feel like the easiest way
might be to more down to the Mike’s Little IR level
and also to just do a bunch of edge hooking myself

rather than using the intended if nothing==res=ChainRules.rrule(…_
thing I can just check the method it would hit,
and then check if it is the fallback,
and if not stick a call into it and and attach an edge
Or if it is hitting the fallback then attack an edge to the method table instead

2 problems:

1. this removes ability to bail out of a chainrule via returning nothing
2. this doesn’t work in 1.0

So I guess I have to put the whole thing including the decision of if to use a ChainRule or source2source
into a generated function that returns a CodeInfo,
so I can set the
method_for_inference_limit_heuristics flag on that.

But if I do that, don't I end up not having the compile know about what functions are called,
and so new ChainRules won't be picked up, without hand connecting edges,
and so it won't be possible to make it work in 1.0.

Is there a way to set method_for_inference_limit_heuristics flag
without having to use a generated function returning a CodeInfo?

Maybe I could overload https://github.com/JuliaLang/julia/blob/a63f2e9b26751ef2d1522fa2634ee9d56db8528d/base/compiler/utilities.jl#L138|
on _pullback as a plain function.

[oxinabox]

Maybe could go the otherway and make pullback_source2source(ctx, f, args....) set its method_for_inference_limit_heuristics to be pullback(ctx, f, args....)
rather than being set to f(args...) like it currently seems to be?

[oxinabox]

@MikeInnes @dhairyagandhi96 how are we going to resolve this?
Will my plan from above work?
#366 (comment)
Do we need to get compiler people in to a Google meet call to work though things?

[oxinabox]

This is much nicer now.
Because some of the things we were blacklisting are gone,
others are fixed,
and I realized we don't need to actually get rid of most of ChainRule's types since they define + in the way accum needs.

[oxinabox]

bors r+

[bors]

Build succeeded:

• ci/gitlab/gitlab.com