How do I implement branch prediction for JIT?

Question

I have a language with a working interpreter, and I am trying to add just-in-time compilation (JIT) to it. For ease of implementation, only individual branches are compiled, while control flow structures are still handled by the interpreter. Unlike traditional JIT, I can't compile on a function level, since this is an esolang that doesn't have functions.

For performance reasons, I plan to have the JIT compile some branches on a background thread before they’re actually needed. Consider some code like the following:

a();
if (someCondition) {
  b();
  if (anotherCondition) {
    c();
  }
} else {
  d();
}
e();

When a() is hit, either branch b() or d() is next, so they should be prioritized over c() and e(). My initial thought for how to implement this was by storing each branch yet to be compiled in a priority queue, and use the following algorithm:

If the interpreter decides to execute a branch, and that branch is already compiled, run the compiled version of that branch. Otherwise, increase the priority of that branch by 1 (so that hot paths reach the front more quickly) and increase the priority of its following branches by 2 (so that branches that still have a chance to be compiled are prioritized over ones we already missed), and continue interpreting.

However, there’s a flaw in this algorithm. What if this happens?

1. a() is hit and not yet compiled, so its priority is increased by 1, and b() and d() have their priorities increased by 2.
2. b() is compiled just as a() finishes execution.
3. The condition ends up being true, and b() is executed.

Now, d() has the highest priority while c() and e() have the lowest priority, which is the opposite of what we want. If I change the priority of branches following compiled ones, 100% compiled code will spend time searching the priority queue to adjust the priority of things that aren’t even there.

How do I choose what branches to compile next for maximum performance?

Answer 1

I thought "branch prediction for JIT" means "branch prediction in JIT", but apparently you want branch prediction outside JIT...

Unfortunately, I don't think it will work in terms of performance for several reasons.

First, pretty much all modern JIT compilers work on a function-by-function granularity level. One of the reasons, on a boundary between the interpreter and JITed code you have to perform some adaptation code — save/restore certain registers, fill the right registers or variables where JITed code will grab the data and alike. Both on the enter and the exit. If you're going to do it before/after every branch, this marshalling code will dominate the execution time.

Another big source of increased performance is the set of certain optimizations, in particular constant propagation, common subexpression elimination, loop-invariant code motion and register allocation. On a sub-function scope there's not that much code to gain significant speedup from these optimisations, while JIT still has to spend time on them.

Lastly, your scheme with updating weights of nested branches is going to introduce pretty random memory accesses which would thrash the cashes and degrade performance even further.

Thus I would advise implementing per-function JIT, or even using off-the-self JIT implementation (Cranelift for instance).

Update

As long as you don't have functions at all, I guess you don't expect users to write too big programs, so (JIT) compiling the whole program at once sounds feasible. That will avoid a lot of JIT grief (calling convention adaptation, on-stack replacement and so on).

"Good old" tracing JIT might present a viable approach for you. Laurence Tratt has a pretty detailed explanation of it with application to implementing new languages.