Tapir's Internal Design - Novice Questions

[limarta]

Hi. I've been trying to read the Tapir source since it is a successful case study on using non-standard code generation in Julia. I'm hoping the code can help provide some instruction for a future non-AD project I have in mind. I'm a novice for anything related to Julia compiler plugins, so unfortunately the code is a bit opaque for my head. A couple of questions.

• Why is the Abstract Interpreter used here? What role does it play for the AD system? According to the comment, Enzyme and Diffractor seem to take a similar approach for "inlining" so I'm curious to see why it is needed.
• What IR are you constructing? The derivative program? What tools are you using to manipulate the IR? It seems everything is done ground up rather than using an existing tool (e.g. IRTools).
• A bit of a tangent, but where in the source do you make the "connection" between Tapir and the ChainRules ecosystem? In other words if I add a new chain rule, where in the source does Tapir go "aha! there is a custom rule here"?
• Do you happen to have a design doc explaining the internals? 😅

Amazing work and good luck with future progress!

[willtebbutt]

Glad to hear that someone has been looking through the internals! To answer your questions:

Why is the Abstract Interpreter used here? What role does it play for the AD system? According to the comment, Enzyme and Diffractor seem to take a similar approach for "inlining" so I'm curious to see why it is needed.

The TapirInterpreter is really only made use of in lookup_ir. The point is that its inlining policy prevents things that I don't want to be inlined from getting inlined (things for which I have hand-written rules).

Enzyme + Diffractor do indeed take a similar approach. The point is that all three packages have a slightly different set of primitives at any given point in time (i.e. a different set of hand-written rules), so they need to choose different things to not inline. We could probably abstract this out at some point.

What IR are you constructing? The derivative program? What tools are you using to manipulate the IR? It seems everything is done ground up rather than using an existing tool (e.g. IRTools).

This is a little bit complicated. Morally speaking, I'm working with Core.Compiler.IRCode. This is the SSA-form IR that the Julia compiler uses to perform a bunch of optimisations on (dead-code elimination etc).

However, in practice, I'm actually making use of a parallel data structure called BBCode, which is really just Core.Compiler.IRCode + names for every line and every basic block. This is needed in my use-case because I move things around a lot, and need to insert additional basic blocks all over the place. I would recommend working directly with Core.Compiler.IRCode if you can.

I'm not using something like IRTools.jl because I need to operate on optimised Julia IR. IRTools.jl is (roughly speaking) designed for working with unoptimised Julia IR.

The way that I produce code which can actually run is by converting this BBCode back into IRCode, and shoving this in an OpaqueClosure. See here.

At a high-level, we derive IRCode for two OpaqueClosures from the IRCode of the original function -- one which does the forwards-pass and one which does the reverse-pass.

A bit of a tangent, but where in the source do you make the "connection" between Tapir and the ChainRules ecosystem? In other words if I add a new chain rule, where in the source does Tapir go "aha! there is a custom rule here"?

There's only a very weak connection -- Tapir.rrule!! has quite different semantics in general, and if you write a ChainRules.rrule you definitely won't automatically get a rule in Tapir.jl. We only make the most simple ChainRules.rrules easy to turn into Tapir.rrule!!s right now -- see here and example usage here. Whether or not we extend this in future is still an open question -- we definitely won't do anything that we're not sure isn't 100% safe.

Do you happen to have a design doc explaining the internals? 😅

Alas, I do not yet. This is very high up my todo list!

Amazing work and good luck with future progress!

Thanks :) I hope you find it helpful!