How to make JIT compilation dependent on the values of function variables?

I am using LLVM to implement a simple language.

When a function that has not yet been JITTed yet and is being called from other already compiled function, I want to lazily JIT-compile it on demand. There is a neat tutorial on this here: 4. Building a JIT: Extreme Laziness - Using Compile Callbacks to JIT from ASTs — LLVM 6 documentation and generally it works.

Now, I want one additional feature - I want to make AST->IR compilation conditional on parameter’s value the first time the function is called.

That is, let us say we have two functions, f() and g(int x).

In AST I have two g implementations - one for positive and one for negative x:

g_pos(x) = x + 1; g_neg(x) = x - 1;

This can be more general, but the idea is that I want to lift the condition x > 0 out of my programme and into the JIT level.

So if now f() = g(7), after JITting I will have f always call g_pos(x).

Is there any way to get this kind of behaviour?

HI @Terminus6,

Good question, and not one that I’ve spent a lot of time thinking about yet.

I hope that one day ORC will provide utilities to help with dynamic specialization, but there’s nothing in-tree yet.

Now, I want one additional feature - I want to make AST->IR compilation conditional on parameter’s value the first time the function is called.

In the special case where you can fully-specialize the function based on the argument values in its first call I think that you could:

(1) Just IRGen g into the module alongside f and let them LLVM optimizers take care of it (set g’s linkage to available_externally if you want to avoid codegen’ing the actual definition in this module). This should be relatively straightforward, but maybe a less interesting answer than you’re looking for.

(2) Use a definition generator to encode information about the argument values into the callee name at the call-site using an encoding scheme. E.g.:

define void @f() {
  call i32 @g$x_eq_7(i32 7)
  ret void
}

declare i32 @g$x_eq_7(i32)

You would use a custom definition generator to handle lookups of symbols of the form <fn>$<call-properties>: when your generator sees a lookup of g$x_eq_7 it would look up the generic AST for g and specialize it based on the property x=7.

A more general runtime specialization scheme (one that could handle repeat calls to g with different arguments) might be to point the symbol g at a “dispatcher” that selects a specialized implementation by inspecting its arguments. E.g. if by analysis of the AST of g you determine that there is a really good optimization opportunities if x==42, and a pretty good one if x >= 100, then you could define your dispatcher as (in C for clarity):

int g(int x) {
  if (x == 42)
    return g$x_eq_42(x);
  else if (x >= 100)
    return g$x_ge_100(x);
  else
    return g$generic(x);
}

If you allow lazy compilation to apply to your specializations you won’t need to pay for compiling all of them up-front, and if you inline the dispatcher into calling modules (as an available_externally function) the optimizers might be able to prove which specialization will be chosen and inline the call directly to it.

Note that I haven’t actually tested these ideas in practice. If you try them out please let me know how it works out. :slight_smile:

– Lang.

Thanks Lang! I’ll try your suggestions, they sound very reasonable.