Feasible to perform pullback preparation at compile time?

[MasonProtter]

Are there plans / ambitions to have a compile-time preparation step for Mooncake.jl? Or is it heavily baked into the design that you need values to construct the tape?

[willtebbutt]

Hi @MasonProtter . There are two kinds of preparation which happen:

1. rule derivation / code generation, and
2. allocation of (co)tangent / shadow memory in which to store results.

The former is entirely a function of the types of the arguments, so can happen once those are available. The tangent memory is necessarily allocate at runtime (as in Enzyme) because you need e.g. size information about arrays in order to allocate it.

At the minute, in both DI and Mooncake's own interfaces, these two steps get bundled together.

Additionally, once a rule is constructed, it's beneficial to re-use it because there's a bunch of allocations which definitely occur the first time that you compute a gradient, but which are typically not required on subsequent runs of the function unless you hit different control flow paths / change the size of arguments etc.

Do you have a particular use case in mind?

[willtebbutt]

I'm going to close this @MasonProtter , but please feel free to re-open if there is more that you would like to discuss!

[MasonProtter]

The former is entirely a function of the types of the arguments, so can happen once those are available. The tangent memory is necessarily allocate at runtime (as in Enzyme) because you need e.g. size information about arrays in order to allocate it.

At the minute, in both DI and Mooncake's own interfaces, these two steps get bundled together.

I see, yeah I think this separation is the thing I was mainly wondering about. I guess the problem I have is that sometimes I want a very light-weight derivative, e.g. scalar derivatives or something, or a function that I take the derivative of in many difference places, but only ever once at a time. For these purposes, having to do all the type level stuff over and over again is quite annoying, whereas with e.g. Enzyme.jl it is just handled fine.

[willtebbutt]

Ah, yeah, I see your problem. I'd love to make it much easier to just access a given rule, but I unfortunately don't know how that would be achieved with the current tools available in Julia. Plainly Enzyme has clearly managed it somehow, but I'm not sure how it should be achieved at the Julia level.

[MasonProtter]

I think at least having an interface to do the type-level stuff as a separate step using only type inputs before the value-step could make it easier to play around with this. I have a bit of experience playing with abstract interpreters and compilation result caching, but can't promise anything.

[willtebbutt]

Ah, cool. So you can do this using an internal interface at the minute. Specifically,

Mooncake.jl/src/interpreter/s2s_reverse_mode_ad.jl

Line 1074 in 59565bc

function build_rrule(

(see also

Mooncake.jl/src/interpreter/s2s_reverse_mode_ad.jl

Line 1051 in 59565bc

build_rrule(sig::Type{<:Tuple}; kwargs...) = build_rrule(get_interpreter(), sig; kwargs...)

).

I'd be happy to make it part of the public interface in the next breaking release of Mooncake (I'm planning to swap out the kwargs for a single Mooncake.Config object) if it's something that would be helpful for you.

[MasonProtter]

Okay, I'll see if I can play around with this a bit some time. Can't promise anything of course. I don't think making them public would make much of a difference for me.

[willtebbutt]

Great. Let me know how you get on. Even if you don't manage to figure it out, I'd be very keen to see whatever approach you wind up taking, as I'm sure that I will learn something!