Dynamic recompilation with ORC JIT API

Project Infrastructure LLVM Dev List Archives
1

Dear all,
I am recently playing a little bit with the ORC JIT APIs with the goal
of creating a prototype of a JIT compiler that loads the IR module of
a C++ class, compiles it and instantiates the class.
During the execution, I would like to replace some class methods by
recompiling them with an optimized version.

I have started from the latest version of the KaleidoscopeJIT, using
the CompileOnDemandLayer to perform a lazy compilation of the class.
I have managed successfully to load the code and instantiate the
object, by getting the pointer of a "factory" method in the IR Module
that returns the object itself.

What I want to do now is to trigger the recompilation of a specific
class method (i.e., function in the IR Module), and perform some IR
transformations before the function is compiled. Then, I would like to
point the old function to the new version (through the
IndirectStubManager).
The questions are:
1. How can I trigger the re-optimization and re-compilation from the
existing ExecutionSession? Is there a method in the
CompileOnDemandLayer that allows it?
2. Do I need to "manually" update the trampoline through the
IndirectStubManager, or it is automatically done once the new function
is compiled?

Thanks a lot,
Sebastiano

2

+Lang for ORC JIT questions

3

Hi Sebastiano,

What I want to do now is to trigger the recompilation of a specific
class method (i.e., function in the IR Module), and perform some IR
transformations before the function is compiled. Then, I would like to
point the old function to the new version (through the
IndirectStubManager).
The questions are:

  1. How can I trigger the re-optimization and re-compilation from the
    existing ExecutionSession? Is there a method in the
    CompileOnDemandLayer that allows it?

There is no off-the-shelf support for this yet. The usual solution is to use an IR pass to insert a counter for each unoptimized function and a corresponding increment / check in the function entry block. I.e.:

define void @foo() {
entry:
body:
; …
ret void
}

becomes

@foo_counter = global i64 0

declare void @optimize_function(i8 *)

define void @foo() {

entry:
%counter = load i64, i64* @foo_counter
%counter.1 = add i64 counter, 1
store i64 %counter.1, i64* @foo_counter
%should_optimize_foo = icmp eq i64 %counter.1, 1000
br i1 %should_optimize_foo, label %optimize, label %old_entry
optimize:

call void @optimize_function(i8* bitcast (void ()* @foo to i8*))
br label %old_entry

old_entry:
body:
; …
ret void
}

Then you would implement optimize_function directly in your JIT process and make it available to JIT’d code (either by reflecting process symbols using a DynamicLibrarySearchGenerator, or by adding it using an absoluteSymbols call). You can use the address passed to optimize_function to identify the function to be optimized. You will also need to give the function a new name (e.g. foo_optimized) to avoid name clashes.

  1. Do I need to “manually” update the trampoline through the
    IndirectStubManager, or it is automatically done once the new function
    is compiled?

Yes: In optimize_function you should look up the address of foo_optimized, then use the IndirectStubManager to update the stub with:
ISM.updatePointer(“foo”, foo_optimized_address);

There are a couple of gotchas here:

(1) You’ll need to set up different compilers for your optimized and unoptimized functions. There are a few ways that you can do that, but one of the easiest would be to create your own IRCompileLayer (copying the existing code – it’s quite short) and then inspect the names of the functions being passed. If it ends in “_optimized” you turn the optimization level up.

(2) The updatePointer method uses linker mangled names. On Linux these are the same as the C / LLVM-IR function names, but on macOS you have to add an underscore prefix, and on Windows there are a couple of different prefix schemes.

This is just a rough overview, but I’m happy to help with any more specific questions as they come up. Also, in the future I’d love to see ORC grow some off-the-shelf support for this, so if you’re ever interested in working on that please let me know.

Regards,
Lang.

4

[snip]

In optimize_function you should look up the address of foo_optimized, then use the IndirectStubManager to update the stub with:
ISM.updatePointer("foo", foo_optimized_address);

There are a couple of gotchas here:

[snip]

You might also have to watch out for the symbols created and used for static initialisation and make sure they refer to the existing globals. For example if the code you're compiling has something like:

void my_function()
{
static some_type something = something_only_once();
}

When you recompile my_function you don't want a new copy of something with its own initialization guard variables, you'll want to re-use the something created the first time my_function ran. In fact if the system were really clever the optimised version of my_function wouldn't need any of the static initialization checking logic, at least not in this simple case. However the main thing is that something_only_once might have side-effects which you don't want happening twice because my_function got recompiled.

Regards,
Raoul.

5

Hi all,
thanks a lot for your answers and your help.

@Lang it is definitely more clear now! I guess this approach is also
very similar to the one applied to the SpeculativeJIT under the
example folder. Am I right?
The main difference in my use case is that I want to trigger the
recompilation outside of the running function (e.g., because other
external events happened that may impact on the function under
consideration).
I will do some tests in the next few weeks and I will definitely come
back in this thread with other questions.
Then, I will be happy to discuss with you and add some off-the-shelf
support in ORC for this.

@Raoul
This shouldn't happen in my case, but it is a great suggestion and
definitely would have caused me a lot of headaches if I did not know
beforehand.

Thanks again,
Sebastiano