Eli Bendersky's website - LLVM & Clang

Adventures in JIT compilation: Part 3 - LLVM

2017-05-01T05:51:00-07:00

This is part 3 of my "Adventures in JIT compilation" series. Part 1 introduced the BF input language and presented a few interpreters in increasing degree of optimization. In part 2 we've seen how to roll a JIT compiler for BF from scratch, by emitting x64 machine code into executable memory and jumping to it for execution.

In this part, I'm going to present another JIT compiler for BF. This time, however, rather than rolling everything from scratch, I'll be using the LLVM framework.

LLVM as a programming language backend

I'll start by saying this post is not meant to be a full tutorial for LLVM. LLVM has a pretty good tutorial already (see also my Python port of it). I've also written more in-depth pieces about LLVM in the past; see for example Life of an instruction in LLVM. That said, I do hope that seeing how to apply LLVM to develop a complete JIT compiler for BF can be useful, and the complete code sample can serve as a starting point for programming language enthusiasts to develop their own backends with a fairly modern version of LLVM. If you have a bit more time, look for exercise and experiment suggestions in the footnotes.

Speaking of LLVM versions... LLVM's C++ API is notoriously volatile, and code working today has little chance of working in just a few months without any changes. However, LLVM does have numbered releases that can be used to maintain some sort of sanity since they are extensively tested. The code for this post was developed with the LLVM 4.0 release. Even if you're reading this in the year 2022, hopefully you should be able to download LLVM 4.0 and compile & link the sample code.

As opposed to the use of asmjit in part 2, LLVM offers several distinct advantages:

Since LLVM comes with many state-of-the-art optimizations on the IR level, we can generate fairly straightforward LLVM IR from our source language, without worrying too much about its efficiency. I'll demonstrate this shortly.
LLVM is multi-target. In this post I'm showing a JIT compiler that emits code for the machine it runs on (x64 in my case), but one can easily reuse the same code to compile BF to ARM, PowerPC, MIPS or a bunch of other backends supported by LLVM.
LLVM comes with a large set of tools useful to visualize and manipulate IR and other stages of compilation. I mention a couple uses of the opt tool throughout the post, and there are others.

Generating LLVM IR from BF

The core of the code generation code in this sample is fairly short - only ~130 lines of C++ in emit_jit_function. I'll walk through it, leaving some details out. Feel free to check the LLVM documentation or API headers for more information, if needed.

llvm::FunctionType* jit_func_type =
    llvm::FunctionType::get(void_type, {}, false);
llvm::Function* jit_func = llvm::Function::Create(
    jit_func_type, llvm::Function::ExternalLinkage, JIT_FUNC_NAME, module);

llvm::BasicBlock* entry_bb =
    llvm::BasicBlock::Create(context, "entry", jit_func);
llvm::IRBuilder<> builder(entry_bb);

We begin by creating a LLVM function to hold the emitted code. The function is named __llvmjit (which is what the constant JIT_FUNC_NAME contains) and its linkage is external so that we could call it from outside the LLVM module where it resides. We also create the entry basic block in the function where the initial code will go, as well as an IR builder that makes the job of emitting LLVM IR a bit easier than using raw APIs.

More setup follows:

llvm::AllocaInst* memory =
    builder.CreateAlloca(int8_type, builder.getInt32(MEMORY_SIZE), "memory");
builder.CreateMemSet(memory, builder.getInt8(0), MEMORY_SIZE, 1);
llvm::AllocaInst* dataptr_addr =
    builder.CreateAlloca(int32_type, nullptr, "dataptr_addr");
builder.CreateStore(builder.getInt32(0), dataptr_addr);

In this version of the BF JIT, I decided to keep the data memory on the stack of the JITed function. This makes it easier for LLVM to optimize accesses to the memory (as we'll see), because the memory is private - it can't be aliased from outside the function and can't have side effects, which frees the optimizer to be more aggressive [1]. The data pointer itself is kept on the stack too (created with an alloca instruction) - more on this very shortly. The data pointer is initialized to 0, per BF semantics [2].

Finally, as usual for one-pass code emission from BF, we have to keep a stack of open brackets (in the asmjit version the type was called BracketLabels):

std::stack<BracketBlocks> open_bracket_stack;

Now we're ready for the compilation loop that takes the next BF instruction and emits the LLVM IR for it. First, pointer movement:

for (size_t pc = 0; pc < program.instructions.size(); ++pc) {
  char instruction = program.instructions[pc];
  switch (instruction) {
  case '>': {
    llvm::Value* dataptr = builder.CreateLoad(dataptr_addr, "dataptr");
    llvm::Value* inc_dataptr =
        builder.CreateAdd(dataptr, builder.getInt32(1), "inc_dataptr");
    builder.CreateStore(inc_dataptr, dataptr_addr);
    break;
  }
  case '<': {
    llvm::Value* dataptr = builder.CreateLoad(dataptr_addr, "dataptr");
    llvm::Value* dec_dataptr =
        builder.CreateSub(dataptr, builder.getInt32(1), "dec_dataptr");
    builder.CreateStore(dec_dataptr, dataptr_addr);
    break;
  }

To move the data pointer, we load its value from its storage on the stack, update it (by either incrementing or decrementing) and store it back. If this seems inefficient, read on! Later on, the post describes why this style of LLVM IR emission is not only acceptable, but desirable.

For data memory updates, the code is not much more complicated:

case '+': {
  llvm::Value* dataptr = builder.CreateLoad(dataptr_addr, "dataptr");
  llvm::Value* element_addr =
      builder.CreateInBoundsGEP(memory, {dataptr}, "element_addr");
  llvm::Value* element = builder.CreateLoad(element_addr, "element");
  llvm::Value* inc_element =
      builder.CreateAdd(element, builder.getInt8(1), "inc_element");
  builder.CreateStore(inc_element, element_addr);
  break;
}
case '-': {
  llvm::Value* dataptr = builder.CreateLoad(dataptr_addr, "dataptr");
  llvm::Value* element_addr =
      builder.CreateInBoundsGEP(memory, {dataptr}, "element_addr");
  llvm::Value* element = builder.CreateLoad(element_addr, "element");
  llvm::Value* dec_element =
      builder.CreateSub(element, builder.getInt8(1), "sub_element");
  builder.CreateStore(dec_element, element_addr);
  break;
}

We similarly load the data pointer, this time using it to compute an offset into the memory (using LLVM's getelementptr instruction). We then load the element from memory, update it and store it back. Note how we use the inbounds variant of getelementptr; BF doesn't define the behavior of stepping outside the bounds of memory - we leverage this fact to let LLVM produce more optimized code.

For I/O, we use a technique similar to the one employed with asmjit in part 2: call the getchar/putchar functions from the host. If you look at the signature of emit_jit_function, it takes a llvm::Function* for each of getchar and putchar; these are declared in the caller by adding their declaration to the module. Later, in the section dealing with the JIT we'll see how these get resolved at run-time.

case '.': {
  llvm::Value* dataptr = builder.CreateLoad(dataptr_addr, "dataptr");
  llvm::Value* element_addr =
      builder.CreateInBoundsGEP(memory, {dataptr}, "element_addr");
  llvm::Value* element = builder.CreateLoad(element_addr, "element");
  llvm::Value* element_i32 =
      builder.CreateIntCast(element, int32_type, false, "element_i32_");
  builder.CreateCall(putchar_func, element_i32);
  break;
}
case ',': {
  llvm::Value* user_input =
      builder.CreateCall(getchar_func, {}, "user_input");
  llvm::Value* user_input_i8 =
      builder.CreateIntCast(user_input, int8_type, false, "user_input_i8_");
  llvm::Value* dataptr = builder.CreateLoad(dataptr_addr, "dataptr");
  llvm::Value* element_addr =
      builder.CreateInBoundsGEP(memory, {dataptr}, "element_addr");
  builder.CreateStore(user_input_i8, element_addr);
  break;
}

As usual, the trickiest part of generating code for BF is handling the loops, which could be nested. LLVM makes this fairly easy by having basic blocks as first-class citizens. Every loop body gets its basic block, and the original block gets split to two - the first part jumping into the loop, the last part happening after the loop (since we can skip directly to it). Here's how the control-flow graph within a function with just a single loop looks [3]:

Note how the jumps between basic blocks happen on a condition and have T (true) or F (false) clauses. These just encode the usual BF semantics:

For a [, we compare the current memory cell with 0; if it's 0, we skip the loop (the T clause); if it's not 0, we enter the loop (the F clause).
For a ], we compare the current memory cell with 0; if it's 0, we jump back to the loop start; otherwise we end the loop.

Here's the code generating LLVM IR from [:

case '[': {
  llvm::Value* dataptr = builder.CreateLoad(dataptr_addr, "dataptr");
  llvm::Value* element_addr =
      builder.CreateInBoundsGEP(memory, {dataptr}, "element_addr");
  llvm::Value* element = builder.CreateLoad(element_addr, "element");
  llvm::Value* cmp =
      builder.CreateICmpEQ(element, builder.getInt8(0), "compare_zero");

  llvm::BasicBlock* loop_body_block =
      llvm::BasicBlock::Create(context, "loop_body", jit_func);
  llvm::BasicBlock* post_loop_block =
      llvm::BasicBlock::Create(context, "post_loop", jit_func);
  builder.CreateCondBr(cmp, post_loop_block, loop_body_block);
  open_bracket_stack.push(BracketBlocks(loop_body_block, post_loop_block));
  builder.SetInsertPoint(loop_body_block);
  break;
}

With the overview and control-flow graph above, I hope it's clear. The most interesting part is using the basic blocks to represent parts of the loop. When [ is encountered, we create both the loop and "post loop" blocks, since the branch instruction has to refer to them. We also save these blocks on the open bracket stack to refer to them when the matching ] is encountered. Finally, we set our IR builder to insert all subsequent instructions into the loop block.

This is how ] is handled:

case ']': {
  if (open_bracket_stack.empty()) {
    DIE << "unmatched closing ']' at pc=" << pc;
  }
  BracketBlocks blocks = open_bracket_stack.top();
  open_bracket_stack.pop();

  llvm::Value* dataptr = builder.CreateLoad(dataptr_addr, "dataptr");
  llvm::Value* element_addr =
      builder.CreateInBoundsGEP(memory, {dataptr}, "element_addr");
  llvm::Value* element = builder.CreateLoad(element_addr, "element");
  llvm::Value* cmp =
      builder.CreateICmpNE(element, builder.getInt8(0), "compare_zero");
  builder.CreateCondBr(cmp, blocks.loop_body_block, blocks.post_loop_block);
  builder.SetInsertPoint(blocks.post_loop_block);
  break;
}

Pretty much as we'd expect: the relevant blocks are popped from the stack and used as targets for another conditional branch [4].

IR sample for a simple BF program

Let's see what LLVM IR our code emits for the simple count1to5.bf program:

++++++++ ++++++++ ++++++++ ++++++++ ++++++++ ++++++++

>+++++
[<+.>-]

The full LLVM-based JIT is available in llvmjit.cpp. When invoked in verbose mode, it will dump LLVM IR for the translated program before and after LLVM optimization. Let's start by looking at the pre-optimization IR. I've added comments on lines starting with ###:

define void @__llvmjit() {
entry:

  ### The entry BB starts by declaring the memory and data pointer, and
  ### initializing the data pointer to 0.

  %memory = alloca i8, i32 30000
  call void @llvm.memset.p0i8.i64(i8* %memory, i8 0, i64 30000, i32 1, i1 false)
  %dataptr_addr = alloca i32
  store i32 0, i32* %dataptr_addr

  ### The following 5 instructions increment memory[dataptr] (with dataptr
  ### remaining at 0, as initialized), and they are repeated 48 times...

  %dataptr = load i32, i32* %dataptr_addr
  %element_addr = getelementptr inbounds i8, i8* %memory, i32 %dataptr
  %element = load i8, i8* %element_addr
  %inc_element = add i8 %element, 1
  store i8 %inc_element, i8* %element_addr

  ### ... 47 more times skipped ...

  store i8 %inc_element184, i8* %element_addr182
  %dataptr185 = load i32, i32* %dataptr_addr
  %element_addr186 = getelementptr inbounds i8, i8* %memory, i32 %dataptr185
  %element187 = load i8, i8* %element_addr186
  %inc_element188 = add i8 %element187, 1
  store i8 %inc_element188, i8* %element_addr186

  ### Now incrementing dataptr and 5 more data increments.

  %dataptr189 = load i32, i32* %dataptr_addr
  %inc_dataptr = add i32 %dataptr189, 1
  store i32 %inc_dataptr, i32* %dataptr_addr

  %dataptr190 = load i32, i32* %dataptr_addr
  %element_addr191 = getelementptr inbounds i8, i8* %memory, i32 %dataptr190
  %element192 = load i8, i8* %element_addr191
  %inc_element193 = add i8 %element192, 1
  store i8 %inc_element193, i8* %element_addr191

  ### ... 4 more increments skipped

  ### Load memory[dataptr] and compare it to 0; on true, jump to %post_loop;
  ### on false jump to %loop_body.

  %dataptr210 = load i32, i32* %dataptr_addr
  %element_addr211 = getelementptr inbounds i8, i8* %memory, i32 %dataptr210
  %element212 = load i8, i8* %element_addr211
  %compare_zero = icmp eq i8 %element212, 0
  br i1 %compare_zero, label %post_loop, label %loop_body

loop_body:

  ### Decrement dataptr

  %dataptr213 = load i32, i32* %dataptr_addr
  %dec_dataptr = sub i32 %dataptr213, 1
  store i32 %dec_dataptr, i32* %dataptr_addr

  ### Increment memory[dataptr]

  %dataptr214 = load i32, i32* %dataptr_addr
  %element_addr215 = getelementptr inbounds i8, i8* %memory, i32 %dataptr214
  %element216 = load i8, i8* %element_addr215
  %inc_element217 = add i8 %element216, 1
  store i8 %inc_element217, i8* %element_addr215

  ### Invoke putchar on memory[dataptr]

  %dataptr218 = load i32, i32* %dataptr_addr
  %element_addr219 = getelementptr inbounds i8, i8* %memory, i32 %dataptr218
  %element220 = load i8, i8* %element_addr219
  %element_i32_ = zext i8 %element220 to i32
  %0 = call i32 @putchar(i32 %element_i32_)

  ### Increment dataptr

  %dataptr221 = load i32, i32* %dataptr_addr
  %inc_dataptr222 = add i32 %dataptr221, 1
  store i32 %inc_dataptr222, i32* %dataptr_addr

  ### Decrement memory[dataptr]

  %dataptr223 = load i32, i32* %dataptr_addr
  %element_addr224 = getelementptr inbounds i8, i8* %memory, i32 %dataptr223
  %element225 = load i8, i8* %element_addr224
  %sub_element = sub i8 %element225, 1
  store i8 %sub_element, i8* %element_addr224

  ### Load memory[dataptr] and compare it to 0; on true, jump back to
  ### %loop_body; on false jump to %post_loop.

  %dataptr226 = load i32, i32* %dataptr_addr
  %element_addr227 = getelementptr inbounds i8, i8* %memory, i32 %dataptr226
  %element228 = load i8, i8* %element_addr227
  %compare_zero229 = icmp ne i8 %element228, 0
  br i1 %compare_zero229, label %loop_body, label %post_loop

post_loop
  call void @dump_memory(i8* %memory)
  ret void
}

As discussed before, this code is doing a huge amount of repetetive and mostly unnecessary work. It keeps storing and reloading values it should already have. But this is precisely what the LLVM optimizer is designed to fix. Let's see the post-optimization code LLVM produces:

define void @__llvmjit() local_unnamed_addr {
loop_body.preheader:
  %memory290 = alloca [30000 x i8], align 1
  %memory290.sub = getelementptr inbounds [30000 x i8], [30000 x i8]* %memory290, i64 0, i64 0
  %0 = getelementptr inbounds [30000 x i8], [30000 x i8]* %memory290, i64 0, i64 2
  call void @llvm.memset.p0i8.i64(i8* nonnull %0, i8 0, i64 29998, i32 1, i1 false)
  store i8 48, i8* %memory290.sub, align 1
  %element_addr191 = getelementptr inbounds [30000 x i8], [30000 x i8]* %memory290, i64 0, i64 1
  store i8 5, i8* %element_addr191, align 1
  br label %loop_body

loop_body:
  %1 = tail call i32 @putchar(i32 49)
  %2 = tail call i32 @putchar(i32 50)
  %3 = tail call i32 @putchar(i32 51)
  %4 = tail call i32 @putchar(i32 52)
  %5 = tail call i32 @putchar(i32 53)
  store i8 53, i8* %memory290.sub, align 1
  store i8 0, i8* %element_addr191, align 1
  call void @dump_memory(i8* nonnull %memory290.sub)
  ret void
}

The LLVM optimizer is extremely aggressive! Not only all the repetition is gone, but there isn't even a loop any more because LLVM statically computed it will just run 5 times and unrolled it completely. putchar is invoked 5 times with the values the loop would have produced, and that's all. The reason LLVM kept memory around was only so the call to dump_memory would have data. If we remove this debugging call, the function turns into:

define void @__llvmjit() local_unnamed_addr #0 {
entry:
  %0 = tail call i32 @putchar(i32 49)
  %1 = tail call i32 @putchar(i32 50)
  %2 = tail call i32 @putchar(i32 51)
  %3 = tail call i32 @putchar(i32 52)
  %4 = tail call i32 @putchar(i32 53)
  ret void
}

Observing how LLVM handles our loops

So the count1to5 sample was a bit too simple for the LLVM optimizer. To see our loops actually being emitted, we'll have to resort to more tricks - by placing "input" instructions (, in BF) in stratetic locations so that LLVM can't just infer their value statically. Take for example the (slightly-nonsensical) BF program:

>,

[<+.>,]

It places user input into cell 1, and then repeatedly:

Increments cell 0, printing its value out
Places new user input in cell 1

It will terminate when the user input is 0. Not very useful, but it does the job. Here's the optimized IR LLVM emits for it:

define void @__llvmjit() local_unnamed_addr {
entry:
  %memory29 = alloca [30000 x i8], align 1
  %memory29.sub = getelementptr inbounds [30000 x i8], [30000 x i8]* %memory29, i64 0, i64 0
  call void @llvm.memset.p0i8.i64(i8* nonnull %memory29.sub, i8 0, i64 30000, i32 1, i1 false)
  %user_input = tail call i32 @getchar()
  %user_input_i8_ = trunc i32 %user_input to i8
  %element_addr = getelementptr inbounds [30000 x i8], [30000 x i8]* %memory29, i64 0, i64 1
  store i8 %user_input_i8_, i8* %element_addr, align 1
  %compare_zero = icmp eq i8 %user_input_i8_, 0
  br i1 %compare_zero, label %post_loop, label %loop_body.preheader

loop_body.preheader:
  br label %loop_body

loop_body:
  %element730 = phi i8 [ %inc_element, %loop_body ], [ 0, %loop_body.preheader ]
  %inc_element = add i8 %element730, 1
  %element_i32_ = zext i8 %inc_element to i32
  %0 = tail call i32 @putchar(i32 %element_i32_)
  %user_input13 = tail call i32 @getchar()
  %user_input_i8_14 = trunc i32 %user_input13 to i8
  %compare_zero20 = icmp eq i8 %user_input_i8_14, 0
  br i1 %compare_zero20, label %post_loop.loopexit, label %loop_body

post_loop.loopexit:
  store i8 %inc_element, i8* %memory29.sub, align 1
  store i8 0, i8* %element_addr, align 1
  br label %post_loop

post_loop:
  call void @dump_memory(i8* nonnull %memory29.sub)
  ret void
}

The two most important points to note are:

There's no extra stack movement; like, at all. LLVM statically computed what happens at address 0 and address 1 and just juggles virtual registers with these values, actually storing to memory only outside the loop!
To be able to do that, it generated a phi instruction at the loop body start.

The latter is especially important: this is LLVM converting the IR to SSA form, where every value is assigned only once and special phi nodes are required to merge multiple possible values. SSA is well outside the scope of this humble post, but I suggest reading about it. LLVM's mem2reg pass converts our naive stack-using code to SSA with virtual registers, and SSA form makes it much easier for the compiler to analyze the IR and optimize it aggressively.

On the other hand, emitters of LLVM IR don't have to worry about efficient usage of virtual registers and can just assign stack slots for all "variables", leaving the optimization to LLVM. For our simple use case of BF this may not matter too much, but think about a classical programming language where a function may have dozens of variables to track [5].

JITing LLVM IR to executable machine code

Being able to optimize LLVM IR is just one of the strengths of LLVM llvmjit is using. The other is the ability to efficiently execute this IR at run-time by JITing it into an excutable in-memory chunk of machine code.

In llvmjit, the part doing this is:

SimpleOrcJIT jit(/*verbose=*/verbose);
module->setDataLayout(jit.get_target_machine().createDataLayout());
jit.add_module(std::move(module));

llvm::JITSymbol jit_func_sym = jit.find_symbol(JIT_FUNC_NAME);
if (!jit_func_sym) {
  DIE << "Unable to find symbol " << JIT_FUNC_NAME << " in module";
}

using JitFuncType = void (*)(void);
JitFuncType jit_func_ptr =
    reinterpret_cast<JitFuncType>(jit_func_sym.getAddress());
jit_func_ptr();

Which, of course, leaves the question - what is SimpleOrcJit? It's a simplified instantiation of LLVM's newest JIT engine - ORC. For the full code see this header and its accompanying source file. My implementation of the JIT is very similar to the one in the LLVM tutorial, with the addition of dumping the JITed machine code to a binary file prior to returning an executable pointer to it.

LLVM's ORC JIT looks formidable, but there's no magic to it. In its essence, it's just doing the thing my introduction to JITing describes, with with many (many!) more layers on top. One interesting thing I would like to mention, though, is how the JIT finds host functions to call, since it's designed differently from the previous JITs I showed.

For llvmjit, all we do is declare (without defining) the host functions in the module, so that we can insert a call instruction to these functions. For example:

declare i32 @putchar(i32) local_unnamed_addr #0

Unlike the approach with asmjit-based JITs, we don't encode the address of the host function in the JITed code. Rather we let the JIT resolve it automatically. This is done by adding the following resolver to the JIT:

auto resolver = orc::createLambdaResolver(
    [this](const std::string& name) {
      if (auto sym = find_mangled_symbol(name)) {
        return sym;
      }
      return JITSymbol(nullptr);
    },
    [](const std::string& name) {
      if (auto sym_addr =
              RTDyldMemoryManager::getSymbolAddressInProcess(name)) {
        return JITSymbol(sym_addr, JITSymbolFlags::Exported);
      }
      return JITSymbol(nullptr);
    });

The second lambda function passed to the resolver is invoked when the LLVM IR being compiled contains a call to a symbol that wasn't defined in the module. In this case, what we do amounts to a call to dlsym with the symbol name. The constructor of SimpleOrcJIT calls LoadLibraryPermanently(nullptr) which translates to dlopen(nullptr, ...), meaning that all symbols contained in the host executable are visible to the JIT. To accomplish this, llvmjit is compiled with the -rdynamic linker flag. This way the JITed code can call any function found in the host code, including standard C library functions like putchar.

How fast does it run?

Let's see how llvmjit compares to the JITs developed in part 2. To make this measurement fair, I'm instrumenting the program to report execution time separately from compilation time, since LLVM may take a while optimizing a large IR program.

With this instrumentation, llvmjit takes 0.92 seconds to run mandelbrot, which is virtually the same time optasmjit took; it takes 0.14 seconds to run factor, half of the time of optasmjit. That said, for these programs LLVM optimizations took roughly the same amount of time as execution, so if you include everything in - llvmjit is equal on factor and loses by ~2x on mandelbrot [6]. So, as often is the case with JITs, it all depends on the use case - if we care about run-time of long programs, it makes sense to spend more time optimizing; it we mostly care about short programs, it doesn't. This is why many modern JITs (think JavaScript) are multi-stage - starting with a fast interpreter or baseline (very light on optimization) JIT, and switching to a more optimizing JIT if the program turns out to be longer-running than initially expected.

Optimization-time aside, it's quite amazing that LLVM is managing to reach the speed of optasmjit, which emits tight machine code and uses domain-specific optimizations to convert whole loops from the BF code to simple O(1) sequences. This is due to the power of LLVM's optimizer; it is able, on its own, to infer that some of these loops are really accomplishing trivial computations. I suggest taking this as an exercise: write simple BF programs and run them through llvmjit --verbose to observe the post-optimization IR. You'll see that LLVM is able to remove loops like [-] and [<->+], replacing them by simple data initialization & movement. Bonus points for uncovering other optimizations that LLVM did but our optasmjit didn't - after all, LLVM does run the factor benchmark much faster, so there must be something interesting there.

Conclusion

After playing with manual code-generation in part 2, this post demonstrates using an industrial-strength compiler backend to create a JIT for BF. LLVM is a large and complex library with a steep learning curve, but once you're past the initial ramp-up stage it's an extremely powerful tool. LLVM lets you emit very straightforward code without worrying about things like registers and recomputing the same values over and over again. Its state-of-the art optimizer removes all these inefficiencies, producing extremely tight IR. Moreover, LLVM lets us target multiple architectures very easily from the same code - just choose which targets you want when configuring & compiling LLVM itself, and you have a multi-target compiler.

However, LLVM has downsides as well - it's large, and its compile time is considerable. This sometimes makes it less desirable in situations where a small footprint and fast compilation are required. For an interesting account of how the WebKit developers replaced LLVM with a custom backend for their "last-tier" optimized JIT, read Introducing the B3 JIT compiler.

Links to all posts in this series:

[1]

As an exercise, move this memory outside the function - by allocating it on the host and passing a pointer to the JITed code, similarly to how it was done in part 2. Observe how this affects LLVM's optimization, if at all. Read on LLVM's concepts of volatile and aliasing and think about how they may affect optimizations.

[2]	A fun experiment is to remove this initialization and observe what LLVM does with the code. I had this bug initially, and it's a brilliant demonstration of LLVM's capability of optimizing code that has undefined behavior (in this case using uninitialized values).

[3] This CFG is produced automatically by LLVM and is a good example of the tools provided by default with this framework. To generate it, I took the LLVM IR file dumped by llvmjit in verbose mode, and ran LLVM's opt tool on it with the -view-cfg-only flag. This makes opt dump a Graphviz .dot file which can then be converted to an image.

[4]	Exercise: why do we need to save `loop_body_block` at all? When a `]` is encountered, shouldn't the matching `loop_body_block` be the same one as the currently populated block?

[5]	A cool experiment to try is compile some C function to LLVM IR with the Clang front-end and observe stack usage in unoptimized code, followed by running `opt -O3` to see what the optimizer turned it into.

[6]	For `mandelbrot`, the textual representation of LLVM IR generated for the BF program runs to ~42,000 lines, which takes LLVM 0.8 seconds to fully optimize.

Updates for building my LLVM & Clang samples for release 3.7

2015-09-06T06:22:00-07:00

The Makefile in my LLVM & Clang samples project is using g++ by default as the C++ compiler. Even though the officially recommended compiler to build Clang and LLVM is Clang, g++ has worked well because not everyone has Clang installed on their system.

This is going to change very soon, when LLVM and Clang 3.7 are released (any day now, as the release is in RC2 stage). When you build the samples vs. the binary of release 3.7, you'll likely run into compile errors about -Wcovered-switch-default, and maybe other problems. The reason for this has to do with the following:

When LLVM & Clang binary releases are created, they are built with Clang.
Starting with 3.7, LLVM & Clang are built with the -Wcovered-switch-default flag by default.
This also means that the llvm-config utility will emit this flag when queried for compile flags. And my samples' Makefile uses llvm-config.
GCC doesn't yet support -Wcovered-switch-default.

There are two possible solutions for this issue. The best one is to just use Clang to build these samples (and any other code linking with LLVM or Clang themselves). You don't even have to have Clang installed for this; simply download the latest binary release, and point CXX to bin/clang++ in the untarred directory. Everything should work, since Clang will find GCC's include paths, libstd++ and everything else.

An alternative approach, if you just must stick to GCC for some reason, is to sanitize the flags emitted by llvm-config, removing those that aren't supported before passing them to the compiler. This means you won't be able to use my Makefile as is, though.

In the long run, I suggest to just use Clang to build LLVM & Clang. This is what the LLVM core developers do to test things and build releases, so this is the way least susceptible to unexpected problems.

Calling back into Python from llvmlite-JITed code

2015-04-18T05:50:00-07:00

Update (2025-01-25): these days llvmlite has binary wheels on PyPI. For self-contained examples (including this post's), check out this GitHub repository.

This post is about a somewhat more interesting and complex use of llvmlite than the basic example presented in my previous article on the subject.

I see compilation as a meta-tool. It lets us build new levels of abstraction and expressiveness within our code. We can use it to build additional languages on top of our host language (common for C, C++ and Java-based systems, less common for Python), to accelerate some parts of our host language (more common in Python), or anything in between.

To fully harness the power of runtime compilation (JITing), however, it's very useful to know how to bridge the gap between the host language and the JITed language; preferably in both directions. As the previous article shows, calling from the host into the JITed language is trivial. In fact, this is what JITing is mostly about. But what about the other direction? This is somewhat more challenging, but leads to interesting uses and additional capabilities.

While the post uses llvmlite for the JITing, I believe it presents general concepts that are relevant for any programming environment.

Callback from JITed code to Python

Let's start with a simple example: we want to be able to invoke some Python function from within JITed code.

from ctypes import c_int64, c_void_p, CFUNCTYPE
import sys

import llvmlite.ir as ir
import llvmlite.binding as llvm

def create_caller(m):
    # define i64 @caller(i64 (i64, i64)* nocapture %f, i64 %i) #0 {
    # entry:
    #   %mul = shl nsw i64 %i, 1
    #   %call = tail call i64 %f(i64 %i, i64 %mul) #1
    #   ret i64 %call
    # }
    i64_ty = ir.IntType(64)

    # The callback function 'caller' accepts is a pointer to FunctionType with
    # the appropriate signature.
    cb_func_ptr_ty = ir.FunctionType(i64_ty, [i64_ty, i64_ty]).as_pointer()
    caller_func_ty = ir.FunctionType(i64_ty, [cb_func_ptr_ty, i64_ty])

    caller_func = ir.Function(m, caller_func_ty, name='caller')
    caller_func.args[0].name = 'f'
    caller_func.args[1].name = 'i'
    irbuilder = ir.IRBuilder(caller_func.append_basic_block('entry'))
    mul = irbuilder.mul(caller_func.args[1],
                        irbuilder.constant(i64_ty, 2),
                        name='mul')
    call = irbuilder.call(caller_func.args[0], [caller_func.args[1], mul])
    irbuilder.ret(call)

create_caller creates a new LLVM IR function called caller and injects it into the given module m.

If you're not an expert at reading LLVM IR, caller is equivalent to this C function:

int64_t caller(int64_t (*f)(int64_t, int64_t), int64_t i) {
  return f(i, i * 2);
}

It takes f - a pointer to a function accepting two integers and returning an integer (all integers in this post are 64-bit), and i - an integer. It calls f with i*2 and i as the arguments. That's it - pretty simple, but sufficient for our demonstration's purposes.

Now let's define a Python function:

def myfunc(a, b):
    print('I was called with {0} and {1}'.format(a, b))
    return a + b

Finally, let's see how we can pass myfunc as the callback caller will invoke. This is fairly straightforward, thanks to the support for callback functions in ctypes. In fact, it's exactly similar to the way you'd pass Python callbacks to C code via ctypes without any JITing involved:

def main():
    module = ir.Module()
    create_caller(module)

    llvm.initialize()
    llvm.initialize_native_target()
    llvm.initialize_native_asmprinter()

    llvm_module = llvm.parse_assembly(str(module))
    tm = llvm.Target.from_default_triple().create_target_machine()

    # Compile the module to machine code using MCJIT.
    with llvm.create_mcjit_compiler(llvm_module, tm) as ee:
        ee.finalize_object()

        # Obtain a pointer to the compiled 'caller' - it's the address of its
        # JITed code in memory.
        CBFUNCTY = CFUNCTYPE(c_int64, c_int64, c_int64)
        cfptr = ee.get_pointer_to_function(llvm_module.get_function('caller'))
        callerfunc = CFUNCTYPE(c_int64, CBFUNCTY, c_int64)(cfptr)

        # Wrap myfunc in CBFUNCTY and pass it as a callback to caller.
        cb_myfunc = CBFUNCTY(myfunc)
        print('Calling "caller"')
        res = callerfunc(cb_myfunc, 42)
        print('  The result is', res)

If we run this code, we get the expected result:

Calling "caller"
I was called with 42 and 84
  The result is 126

Registering host functions in JITed code

When developing a JIT, one need that comes up very often is to delegate some of the functionality in the JITed code to the host language. For example, if you're developing a JIT to implement a fast DSL, you may not want to reimplement a whole I/O stack in your language. So you'd prefer to delegate all I/O to the host language. Taking C as a sample host language, you just want to call printf from your DSL and somehow have it routed to the host call.

How do we accomplish this feat? The solution here, naturally, depends on both the host language and the DSL you're JITing. Let's take the LLVM tutorial as an example. The Kaleidoscope language does computations on floating point numbers, but it has no I/O facilities of its own. Therefore, the Kaleidoscope compiler exposes a putchard function from the host (C++) to be callable in Kaleidoscope. For Kaleidoscope this is fairly simple, because the host is C++ and is compiled into machine code in the same process with the JITed code. All the JITed code needs to know is the symbol name of the host function to call and the call can happen (as long as the calling conventions match, of course).

Alas, for Python as a host language, things are not so straightforward. This is why, in my reimplementation of Kaleidoscope with llvmlite, I resorted to implementing the builtins in LLVM IR, emitting them into the module along with compiled Kaleidoscope code. These builtins just call the underlying C functions (which still reside in the same process, since Python itself is written in C) and don't call into Python.

But say we wanted to actually call Python. How would we go about that?

Well, we've seen a way to call Python from JITed code in this post. Can this approach be used? Yes, though it's quite cumbersome. The problem is that the only place where we have an actual interface between Python and the JITed code is when we invoke a JITed function. Somehow we should use this interface to convey to the JIT side what Python functions are available to it and how to call them. Essentially, we'll have to implement something akin to the following schematic symbol table interface in the JITed code:

typedef int64_t (*CallbackType)(int64_t, int64_t);
std::unordered_map<std::string, CallbackType> symtab;

void register_callback(std::string name, CallbackType callback) {
  symtab[name] = callback;
}

CallbackType get_callback(std::string name) {
  auto iter = symtab.find(name);
  if (iter != symtab.end()) {
    return iter->second;
  } else {
    return nullptr;
  }
}

To register Python callbacks with the JIT, we'll call register_callback from Python, passing it a name and the callback (CFUNCTYPE as shown in the code sample at the top). The JIT side will remember this mapping in a symbol table. When it needs to invoke a Python function it will use get_callback to get the pointer by name.

In addition to being cumbersome to implement [1], this is also inefficient. It seems wasteful to go through a symbol table lookup for every call to a Python builtin. It's not like these mappings ever change in a typical use case! We are emitting code at runtime here and have so much flexibility at our command - so this lookup feels like a crutch.

Moreover, this is a simplified example - every callback takes two integer arguments. In real scenarios, the signatures of callback functions can be arbitrary, so we'd have to implement a full blown FFI-dispatching on the calls.

Breaching the compile/run-time barrier

We can do better. For every Python function we intend to call from the JITed code, we can emit a JITed wrapper. This wrapper will hard-code a call to the Python function, thus removing this dispatching (the symbol table shown above) from run-time; this totally makes sense because we know at compile time which Python functions are needed and where to find them.

Let's write the code to do this with llvmlite:

import ctypes
from ctypes import c_int64, c_void_p
import sys

import llvmlite.ir as ir
import llvmlite.binding as llvm

cb_func_ty = ir.FunctionType(ir.IntType(64),
                             [ir.IntType(64), ir.IntType(64)])
cb_func_ptr_ty = cb_func_ty.as_pointer()
i64_ty = ir.IntType(64)

def create_addrcaller(m, addr):
    # define i64 @addrcaller(i64 %a, i64 %b) #0 {
    # entry:
    #   %f = inttoptr i64% ADDR to i64 (i64, i64)*
    #   %call = tail call i64 %f(i64 %a, i64 %b)
    #   ret i64 %call
    # }
    addrcaller_func_ty = ir.FunctionType(i64_ty, [i64_ty, i64_ty])
    addrcaller_func = ir.Function(m, addrcaller_func_ty, name='addrcaller')
    a = addrcaller_func.args[0]; a.name = 'a'
    b = addrcaller_func.args[1]; b.name = 'b'
    irbuilder = ir.IRBuilder(addrcaller_func.append_basic_block('entry'))
    f = irbuilder.inttoptr(irbuilder.constant(i64_ty, addr),
                           cb_func_ptr_ty, name='f')
    call = irbuilder.call(f, [a, b])
    irbuilder.ret(call)

The IR function created by create_addrcaller is somewhat similar to the one we've seen above with create_caller, but there's a subtle difference. addcaller does not take a function pointer at runtime. It has knowledge of this function pointer encoded into it when it's generated. The addr argument passed into create_addrcaller is the runtime address of the function to call. addrcaller converts it to a function pointer (using the inttoptr instruction, which is somewhat similar to a reinterpret_cast in C++) and calls it [2].

Here's how to use it:

def main():
    CBFUNCTY = ctypes.CFUNCTYPE(c_int64, c_int64, c_int64)
    def myfunc(a, b):
        print('I was called with {0} and {1}'.format(a, b))
        return a + b
    cb_myfunc = CBFUNCTY(myfunc)
    cb_addr = ctypes.cast(cb_myfunc, c_void_p).value
    print('Callback address is 0x{0:x}'.format(cb_addr))

    module = ir.Module()
    create_addrcaller(module, cb_addr)
    print(module)

    llvm.initialize()
    llvm.initialize_native_target()
    llvm.initialize_native_asmprinter()

    llvm_module = llvm.parse_assembly(str(module))

    tm = llvm.Target.from_default_triple().create_target_machine()

    # Compile the module to machine code using MCJIT
    with llvm.create_mcjit_compiler(llvm_module, tm) as ee:
        ee.finalize_object()
        # Now call addrcaller
        print('Calling "addrcaller"')
        addrcaller = ctypes.CFUNCTYPE(c_int64, c_int64, c_int64)(
            ee.get_pointer_to_function(llvm_module.get_function('addrcaller')))
        res = addrcaller(105, 23)
        print('  The result is', res)

The key trick here is the call to ctypes.cast. It takes a Python function wrapped in a ctypes.CFUNCTYPE and casts it to a void*; in other words, it obtains its address [3]. This is the address we pass into create_addrcaller. The code ends up having exactly the same effect as the previous sample, but with an important difference: whereas previously the dispatch to myfunc happened at run-time, here it happens at compile-time.

This is a synthetic example, but it should be clear how to extend it to the full thing mentioned earlier: for each built-in needed by the JITed code from the host code, we emit a JITed wrapper to call it. No symbol table dispatching at runtime. Even better, since these builtins can have arbitrary signatures, the JITed wrapper can handle all of that efficiently. PyPy uses this technique to make calls into C (via the cffi library) much more efficient than they are with ctypes. ctypes uses libffi, which has to pack all the arguments to a function at runtime, according to a type signature it was provided. However, since this type signature almost never changes during the runtime of one program, this packing can be done much more efficiently with JITing.

Conclusion

Hopefully it's clear that while this article focuses on a very specific technology (using llvmlite to JIT native code from Python), its principles are universal. The overarching idea here is that the difference between what happens when the program is compiled and what happens when it runs is artificial. We can breach and overlay it in many ways, and use it to build increasingly complex abstractions. Some languages, like the Lisp family, list this mixture of compile-time and run-time as one of their unique strengths, and have been preaching it for decades. I fondly recall my own first real-world use of this technique many years ago - reading a configuration file and generating code at runtime that unpacks data based on that configuration. That task, emitting Perl code from Perl according to a XML config may appear worlds away from the topic of this post - emitting LLVM IR from Python according to a function signature, but if you really think about it, it's exactly the same thing.

I suspect this is one of the most obtuse articles I've written lately; if you read this far, I sure hope you found it interesting and helpful. Let me know in the comments if anything isn't clear or if you have relevant ideas - I love discussing this topic!

[1]	We'll have to compile the equivalent of a hash table implementation into our JITed code. While not impossible, this may be an overkill if you really just want a quick-and-simple DSL.

[2]	This delightful mixture of compile-time and run-time is by far the most important part of this article; if you remember just one thing from here, this should be it. Let me know in the comments if it's not clear.

[3] The concept of "address" for a Python function may raise an eyebrow. Keep in mind that this isn't a pure Python function we're talking about here. It's wrapped in a ctypes.CFUNCTYPE, which is a dispatcher created by ctypes ("thunk" in the nomenclature of libffi, the underlying mechanism behind ctypes) to perform argument conversion and make the actual call.

Python version of the LLVM tutorial

2015-02-08T15:28:00-08:00

The LLVM tutorial is a venerable and important part of the project's documentation. It's been there for as long as I've been using LLVM (and according to the logs, a few years before that), almost always the first resource newcomers to the project are pointed to. It strikes just the right balance between simplicity and interesting content to provide an enticing introduction to LLVM. For a motivated reader, it shouldn't take more than a work day or two to go through it from start to finish, building a full compiler for a simple but "real" programming language in the process; how cool is that?

Anyway, it occurred to me that since the "official" version of the tutorial is in C++ and the only alternative checked into the tree is in OCaml, it may be interesting to re-implement the tutorial in Python. While I wouldn't write an industrial strength compiler in Python, it's a great prototyping platform, and when thinking about compilers and languages in general, prototyping is very imporatnt. You want to try all kinds of possibilities and combinations of features quickly, to get a feel for writing code in the language before it's fully done - and Python (with LLVM) is great for that.

Enter Pykaleidoscope, a project I put on GitHub that follows the steps of the official LLVM tutorial, but implementing the Kaleidoscope compiler in Python, using llvmlite as the binding to LLVM.

Installing llvmlite is fairly easy - see this post if you have any issues.

While working on Pykaleidoscope, I was impressed with llvmlite's maturity and compatibility with the C++ LLVM IR APIs. I didn't run into any significant problems, except maybe lack of documentation. But documentation isn't a strong side of LLVM either, which is one of the problems the tutorial helps with. So I hope this Python version will help folks understand how to use llvmlite to build non-trivial LLVM IR in Python.

Building and using llvmlite - a basic example

2015-01-26T05:03:00-08:00

Update (2025-01-25): these days llvmlite has binary wheels on PyPI. For self-contained examples (including this post's), check out this GitHub repository.

A while ago I wrote a short post about employing llvmpy to invoke LLVM from within Python. Today I want to demonstrate an alternative technique, using a new library called llvmlite. llvmlite was created last year by the developers of Numba (a JIT compiler for scientific Python), and just recently replaced llvmpy as their bridge to LLVM. Since the Numba devs were also the most active maintainers of llvmpy in the past couple of years, I think llvmlite is definitily worth paying attention to.

One of the reasons the Numba devs decided to ditch llvmpy in favor of a new approach is the biggest issue heavy users of LLVM as a library have - its incredible rate of API-breaking change. The LLVM C++ API is notoriously unstable and will remain this way for the foreseeable future. This leaves library users and all kinds of language bindings (like llvmpy) in a constant chase after the latest LLVM release, if they want to benefit from improved optimizations, new targets and so on. The Numba developers felt this while maintaining llvmpy and decided on an alternative approach that will be easier to keep stable going forward. In addition, llvmpy's architecture made it slow for Numba's users - llvmlite fixes this as well.

The main idea is - use the LLVM C API as much as possible. Unlike the core C++ API, the C API is meant for facing external users, and is kept relatively stable. This is what llvmlite does, but with one twist. Since building the IR using repeated FFI calls to LLVM proved to be slow and error-prone in llvmpy, llvmlite re-implemented the IR builder in pure Python. Once the IR is built, its textual representation is passed into the LLVM IR parser. This also reduces the "API surface" llvmlite uses, since the textual representation of LLVM IR is one of the more stable things in LLVM.

I found llvmlite pretty easy to build and use on Linux (though it's portable to OS X and Windows as well). Since there's not much documentation yet, I thought this post may be useful for others who wish to get started.

After cloning the llvmlite repo, I downloaded the binary release of LLVM 3.5 - pre-built binaries for Ubuntu 14.04 mean there's no need to compile LLVM itself. Note that I didn't install LLVM, just downloaded and untarred it.

Next, I had to install the libedit-dev package with apt-get, since it's required while building llvmlite. Depending on what you have lying around on your machine, you may need to install some additional -dev packages.

Now, time to build llvmlite. I chose to use Python 3.4, but any modern version should work (for versions below 3.4 llvmlite currently requires the enum34 package). LLVM has a great tool named llvm-config in its binary image, and the Makefile in llvmlite uses it, which means building llvmlite with any version of LLVM I want is just a simple matter of running:

$ LLVM_CONFIG=<path/to/llvm-config> python3.4 setup.py build

This compiles the C/C++ parts of llvmlite and links them statically to LLVM. Now, you're ready to use llvmlite. Again, I prefer not to install things unless I really have to, so the following script can be run with:

$ PYTHONPATH=$PYTHONPATH:<path/to/llvmlite> python3.4 basic_sum.py

Replace the path with your own, or just install llvmlite into some virtualenv.

And the sample code does the same as the previous post - creates a function that adds two numbers, and JITs it:

from ctypes import CFUNCTYPE, c_int
import sys

import llvmlite.ir as ll
import llvmlite.binding as llvm

llvm.initialize()
llvm.initialize_native_target()
llvm.initialize_native_asmprinter()

# Create a new module with a function implementing this:
#
# int sum(int a, int b) {
#   return a + b;
# }
module = ll.Module()
func_ty = ll.FunctionType(ll.IntType(32), [ll.IntType(32), ll.IntType(32)])
func = ll.Function(module, func_ty, name='sum')

func.args[0].name = 'a'
func.args[1].name = 'b'

bb_entry = func.append_basic_block('entry')
irbuilder = ll.IRBuilder(bb_entry)
tmp = irbuilder.add(func.args[0], func.args[1])
ret = irbuilder.ret(tmp)

print('=== LLVM IR')
print(module)

# Convert textual LLVM IR into in-memory representation.
llvm_module = llvm.parse_assembly(str(module))

tm = llvm.Target.from_default_triple().create_target_machine()

# Compile the module to machine code using MCJIT
with llvm.create_mcjit_compiler(llvm_module, tm) as ee:
    ee.finalize_object()
    print('=== Assembly')
    print(tm.emit_assembly(llvm_module))

    # Obtain a pointer to the compiled 'sum' - it's the address of its JITed
    # code in memory.
    cfptr = ee.get_pointer_to_function(llvm_module.get_function('sum'))

    # To convert an address to an actual callable thing we have to use
    # CFUNCTYPE, and specify the arguments & return type.
    cfunc = CFUNCTYPE(c_int, c_int, c_int)(cfptr)

    # Now 'cfunc' is an actual callable we can invoke
    res = cfunc(17, 42)
    print('The result is', res)

This should print the LLVM IR for the function we built, its assembly as produced by LLVM's JIT compiler, and the result 59.

Compared to llvmpy, llvmlite now seems like the future, mostly due to the maintenance situation. llvmpy is only known to work with LLVM up to 3.3, which is already a year and half old by now. Having just been kicked out of Numba, there's a good chance it will fall further behind. llvmlite, on the other hand, is very actively developed and keeps track with the latest stable LLVM release. Also, it's architectured in a way that should make it significantly easier to keep up with LLVM in the future. Unfortunately, as far as uses outside of Numba go, llvmlite is still rough around the edges, especially w.r.t. documentation and examples. But the llvmlite developers appear keen on making it useful in a more general setting and not just for Numba, so that's a good sign.