[RFC] JIT Design Details

👋 Hey,

Summary

Determine the Goals and Design of our JIT.

Detailed Explanation

This started out in #4262 where we aren't really sure how we want to implement function calls, and how inlined / outlined we want to go with stuff in the JIT. So I went and figured the broad details of how a bunch of other popular JIT's work.

This is a bit more focused on how the JIT's interact with the VM's / Interpreters, but in general it looks like we also don't have a set direction on where to go with the JIT, so It would be great if folks could chime in with their opinions on this!

How other JIT's work

Numba

Numba JIT's entire functions at a time. Functions that are to be JIT'ed are marked with a @jit decorator. That marks a function for compilation on its first call. The user can also provide a signature for numba to specialize ahead of time.

See the docs on its internals, they are really good!

Numba has two compilation modes. object mode makes all values python objects and uses the Python C API to perform all operations. And as expected they mention that Its not very fast.

However Numba can also specialize some functions (or single loops!) if they know the types of the objects, and if they do so, they are called nopython mode functions. These are closer to what we have today in our JIT, they require type inference (or type annotations), but since all operations are natively compiled they are way faster.

Other links:

• List of all Compiler stages.
• Environment Object
• Memory Management

Ruby's YJIT

Ruby has recently gained a JIT. It compiles a method at a time, and instead of doing an all or nothing approach they try to codegen as much as possible natively, and if they hit an unsupported operation they return control to the interpreter.

• YJIT - Entering and Exiting the Interpreter

LuaJIT / PyPy

LuaJIT is a Tracing JIT. PyPy fits into a similar category of Meta-Tracing JIT's

They trace functions, loops, and other places. Figure out what operations are executed in that function/loop/etc.. and compile just that trace into native code. This means that a trace is a very specific subset of a function.

This is significantly different from how our JIT works today, so It's not really worth discussing unless we want to significantly reorganize our JIT.

• Mike Pall's summary of how LuaJIT's compiler works

What should our JIT look like?

Currently our JIT looks a lot like Numba's nopython mode! We either compile the function correctly and execute the whole thing as native code or we fail compilation and interpret the function.

This requires us to know the precise types of all items at every point in the function, which in turn is probably going to require some advanced type inference.

Personally I like YJIT's approach of optimizing what it can, and falling back to the Interpreter otherwise. However that requires some expensive rebuilding when it has to return control to the interpreter. I think they try to re-synchronize the stack from the JITed code back into the interpreter at the point that it failed.

It would be nice if we could avoid falling back to the interpreter so hard, but we could try to reuse some operations individually, if we can't natively support them. So, sort of mixing numba's object mode when we don't precisely know the types of stuff, but for types that we DO know implementing those natively like numba's nopython mode.

But I really don't have a strong preference for any of those designs!

Drawbacks, Rationale, and Alternatives

Having a JIT design pinned down is helpful to resolve a bunch of other design questions (such as #4262). For example if we are going with a sort of nopython mode, then that is definitely the correct answer. But if we are going with a strict object mode, then we need to call back to the interpreter instead.

As mentioned above I think there could be some benefits to building a mixed nopython/object mode. However that probably requires us to do some type inference, which may not be particularly simple, we can also specialize functions per types, but even then that probably still requires some type inference.

It is also slower than a real nopython mode, however we do gain the advantage that it is significantly easier to compile methods, even if we don't have the full type information.

By the way I've never worked in a JIT so, take all of this with a huge grain of salt.

Unresolved Questions

Everything else!

GC and Memory Management are probably the big ones. Another one that recently came up in #2293 and #4269 is how do we handle exceptions. Do we want to specialize functions based on signatures? Do we want to optimize just loops? And probably a bunch of other ones that I haven't thought of yet!

These all have some impact on the design of the JIT and would be really useful to have pinned down.

首先阅读相关 issue #4262、#2293 和 #4269，然后比较为 Numba、YJIT、LuaJIT 和 PyPy 链接的 JIT 方案。项目就 JIT 方向达成一致，并解决关于编译范围、解释器回退、类型特化、垃圾回收、内存管理和异常的所列问题后，即视为完成。