Eli Bendersky's website - Code generation

Adventures in JIT compilation: Part 4 - in Python

2017-05-10T05:31:00-07:00

This is part 4 in a series of posts on writing JIT compilers for BF.

Even though, for performance considerations, most JIT compilers are written in down-to-metal languages like C and C++, in some scenarios it could be useful to write them in higher-level languages as well [1].

This post intends to do just that, by presenting yet another JIT for BF - this time written in Python. It presents both a low-level approach to JITing native machine code, and using a higher-level library to perform instruction encoding. The JIT implemented here is not optimized (equivalent to the "simple" JIT from part 2), so I'm not going to discuss any performance results. The idea is just to show how JITing in Python works.

How to JIT from Python

To be able to JIT-compile BF, we should first be capable of JITing any code at all. It's worth showing how to do this in plain Python without any third-party libraries. In part 2 of the series we've used the same approach as presented in my older How to JIT post for basic JITing by copying machine code into an executable memory segment and jumping to it.

That was done in C and C++, but it turns out it's not much harder in Python due to the magic of ctypes. I've written about doing runtime calls to C code via ctypes before, and if you're not familiar with this wonderful tool, I encourage you to read more about it. What follows is a complete code sample that implements the JITing presented in the "How to JIT" post; it's heavily commented, so it should be reasonably clear what's going on:

import ctypes
import mmap

# Equivalent to dlopen(NULL) to open the main process; this means the Python
# interpreter, which links the C library in. Alternatively, we could open
# libc.so.6 directly (for Linux).
libc = ctypes.cdll.LoadLibrary(None)

# Grab the mmap foreign function from libc, setting its signature to:
#
#     void *mmap(void *addr, size_t length, int prot, int flags,
#                int fd, off_t offset)
#
# Per the mmap(2) man page.
mmap_function = libc.mmap
mmap_function.restype = ctypes.c_void_p
mmap_function.argtypes = (ctypes.c_void_p, ctypes.c_size_t,
                          ctypes.c_int, ctypes.c_int,
                          ctypes.c_int, ctypes.c_size_t)

CODE_SIZE = 1024

# Allocate RWX memory with mmap. Using mmap from libc directly rather than
# Python's mmap module here because the latter returns a special "mmap object"
# and we just need the raw address of the mapping. However, we will use the
# PROT_* and MAP_* constants from the mmap module rather than duplicating them.
code_address = mmap_function(None, CODE_SIZE,
                             mmap.PROT_READ | mmap.PROT_WRITE | mmap.PROT_EXEC,
                             mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS,
                             -1, 0)

if code_address == -1:
    raise OSError('mmap failed to allocate memory')

# Emit code into the allocated address. This sequence of x86-64 machine code
# encodes:
#
#  mov %rdi, %rax
#  add $4, %rax
#  ret
code = b'\x48\x89\xf8\x48\x83\xc0\x04\xc3'
assert len(code) <= CODE_SIZE
ctypes.memmove(code_address, code, len(code))

# Declare the type for our JITed function: long (*JitFuncType)(long), and cast
# code_address (which is a void*) to this type so we could just call it.
JitFuncType = ctypes.CFUNCTYPE(ctypes.c_long, ctypes.c_long)
function = ctypes.cast(code_address, JitFuncType)
print(function(-100))

If you execute this code it will print -96, since the JITed function adds 4 to its argument.

From this point on, it should be easy to implement a full JIT-compiler for BF using hand-rolled instruction encoding, just like in part 2. Feel free to do it as an exercise :) You'll face the same issue discussed in that post - manually encoding instructions and resolving labels/branches is tedious, and it would be nice if a library could do it for us.

PeachPy

Enter PeachPy - an x86-64 assembler embedded in Python. It's a relatively new project mostly aimed at writing highly-optimized assembly kernels for numerical computations without leaving the domain of Python. It even has a cute logo:

PeachPy has some examples strewn around the web, and my code for this post can serve as another example. Let's start by replicating the simple JIT functionality shown above. PeachPy takes care of all the memory mapping behind the scenes [2], so the full code we need to write is just this:

import peachpy
import peachpy.x86_64

x = peachpy.Argument(peachpy.int64_t)

with peachpy.x86_64.Function("Add4", (x,), peachpy.int64_t) as asm_function:
    peachpy.x86_64.LOAD.ARGUMENT(peachpy.x86_64.rax, x)
    peachpy.x86_64.ADD(peachpy.x86_64.rax, 4)
    peachpy.x86_64.RETURN(peachpy.x86_64.rax)

abi = peachpy.x86_64.abi.detect()
encoded_function = asm_function.finalize(abi).encode()
python_function = encoded_function.load()

# The JIT call
print(python_function(-100))

PeachPy lets us write our code in assembly rather than directly in machine code, and it also handles the loading for us. As we'll see soon, it lets us use symbolic labels (just like asmjit) and takes care of jump target resolution automatically. In other words, it's just the tool we need to focus on our domain - compiling BF - without worrying too much about the low level details.

Note how PeachPy uses Python's context managers (a very common approach for DSLs embedded in Python). The with peachpy.x86_64.Function statement creates a context manager within which all assembly instructions (like ADD and RETURN) are appended to the function.

JIT-compiling BF via PeachPy

With the Add4 sample above, we actually already have most of the building blocks we need to JIT-compile BF. What remains is just a matter of digging in PeachPy's source and examples (sadly there's very little documentation) to figure out how to invoke and properly use its assembly primitives. The following is very similar to how simpleasmjit is implemented in part 2. The full code sample is available here; I'll just highlight the important snippets of code. First the function "declaration":

# Create a JITed function named "ppjit", with the C-style signature:
#   void ppjit(uint8_t* memptr)
#
memptr = peachpy.Argument(peachpy.ptr(peachpy.uint8_t))

with peachpy.x86_64.Function("ppjit",
                             [memptr],
                             result_type=None) as asm_function:
    # Use r13 as our data pointer; initially it points at the memory buffer
    # passed into the JITed function.
    dataptr = peachpy.x86_64.r13
    peachpy.x86_64.LOAD.ARGUMENT(dataptr, memptr)

In this code we'll be passing memory from the host side (from Python, in this case), so the function signature is void (*ppjit)(uint8_t*). We then give r13 the symbolic name dataptr. The usual instruction-by-instruction BF compilation loop follows:

for pc, instr in enumerate(parse_bf_program(bf_file), start=1):
    if instr == '>':
        peachpy.x86_64.ADD(dataptr, 1)
    elif instr == '<':
        peachpy.x86_64.SUB(dataptr, 1)
    elif instr == '+':
        peachpy.x86_64.ADD([dataptr], 1)
    elif instr == '-':
        peachpy.x86_64.SUB([dataptr], 1)

PeachPy dresses Python in assembly-like flavor. Registers placed in brackets like [dataptr] imply dereferencing. While dataptr refers to the value of the r13 register itself, [dataptr] refers to the value of the memory word whose address is held in dataptr - this is similar to the syntax of many assembly languages.

For emitting I/O operations, we resort to syscalls again [3]:

elif instr == '.':
    # Invoke the WRITE syscall (rax=1) with stdout (rdi=1).
    peachpy.x86_64.MOV(peachpy.x86_64.rax, 1)
    peachpy.x86_64.MOV(peachpy.x86_64.rdi, 1)
    peachpy.x86_64.MOV(peachpy.x86_64.rsi, dataptr)
    peachpy.x86_64.MOV(peachpy.x86_64.rdx, 1)
    peachpy.x86_64.SYSCALL()
elif instr == ',':
    # Invoke the READ syscall (rax=0) with stdin (rdi=0).
    peachpy.x86_64.MOV(peachpy.x86_64.rax, 0)
    peachpy.x86_64.MOV(peachpy.x86_64.rdi, 0)
    peachpy.x86_64.MOV(peachpy.x86_64.rsi, dataptr)
    peachpy.x86_64.MOV(peachpy.x86_64.rdx, 1)
    peachpy.x86_64.SYSCALL()

For emitting the loops, we use a stack of PeachPy Label objects with one label for the loop and another for the "after loop". Again, this is very similar to how it was done with asmjit in part 2.

elif instr == '[':
    # Create labels for the loop start and after-loop.
    loop_start_label = peachpy.x86_64.Label()
    loop_end_label = peachpy.x86_64.Label()
    # Jump to after the loop if the current cell is 0.
    peachpy.x86_64.CMP([dataptr], 0)
    peachpy.x86_64.JZ(loop_end_label)
    # Bind the "start loop" label here.
    peachpy.x86_64.LABEL(loop_start_label)
    open_bracket_stack.append(
        BracketLabels(loop_start_label, loop_end_label))
elif instr == ']':
    if not len(open_bracket_stack):
        die('unmatched closing "]" at pc={}'.format(pc))
    labels = open_bracket_stack.pop()
    # Jump back to loop if the current cell is not 0.
    peachpy.x86_64.CMP([dataptr], 0)
    peachpy.x86_64.JNZ(labels.open_label)
    # Bind the "after-loop" label here.
    peachpy.x86_64.LABEL(labels.close_label)

Finally, this is how the JITed function is invoked:

# Allocate memory as a ctypes array and initialize it to 0s. Then perform
# the JIT call.
memsize = 30000
MemoryArrayType = ctypes.c_uint8 * memsize
memory = MemoryArrayType(*([0] * memsize))

There's just a little bit of trickiness here in our usage of ctypes. Since we have to pass arguments to C functions, ctypes lets us declare C-like types like MemoryArrayType which is a uint8_t[30000]. The funky syntax on the following line is nothing more than Python's *args destructuring of a list of 30000 zeros. So memory is now an object we can safely pass to our JITed function which expects a uint8_t* argument:

python_function(memory)

This calls the JITed function, the result of which is I/O and modified memory cells in memory.

PeachPy - more features and some limitations

My usage of PeachPy in this post has been fairly limited, and the library has many more features. For example:

Since PeachPy was mainly developed to write fast mathematical kernels, it has fairly good support for the newest vector extensions like AVX512.
There's some support for automatic register allocation. In the BF JIT, the r13 register is used directly, but we don't really care which register it is. We could ask PeachPy for a "general purpose 64-bit register" and it would allocate one for us. When writing a much more complicated piece of assembly code, this can be very useful.
There's also some support for generating ARM code, though I don't know how mature it is.

However, PeachPy is also a fairly new library which means there are still limitations and rough edges. When developing the JIT I ran into a bug and sent a PR - which was promptly accepted. PeachPy also doesn't do well with a large number of assembly instructions. It has some recursive analysis logic that blows up the Python stack when run on large code. In the BF JIT, I'm setting the stack to a higher limit artificially; with that, PeachPy doesn't crash but takes quite a while to assemble large programs like Mandelbrot.

According to PeachPy's maintainer, this is due to the design of PeachPy being aimed towards writing assembly code manually rather than generating it from some other language. Fair enough - but definitely something to keep in mind. All in all, the maintainer is responsive and the library seems to be improving quickly - so these limitations may go away in the future.

Python JIT alternative - llvmlite

A somewhat higher-level alternative to JITing from Python is using llvmlite, the new Python binding to LLVM. I wrote about llvmlite before and also ported the LLVM official tutorial to Python using it.

llvmlite would definitely do the job here (and most likely not have the limitations of PeachPy), but I wanted to go with something different in this part. After all, a BF JIT with LLVM was already covered in part 3, and llvmlite is a binding to the same library, just in Python. PeachPy offers an altenative approach for machine code generation from Python - with more direct control of the emitted instructions, though none of the optimizations LLVM provides automatically.

[1]	For example Numba - a JIT compiler for numeric code written in Python; it takes Python code that uses Numpy and JIT-compiles it to optimized native code.

[2]	No magic is involved here - in fact PeachPy uses precisely the same approach as I've shown to invoke `mmap` from Python via `ctypes`.

[3]	As an exercise, modify this code to invoke `putchar` and `getchar`.

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.

Adventures in JIT compilation: Part 2 - an x64 JIT

2017-03-22T06:32:00-07:00

In the first part of the series I've briefly introduced the BF source language and went on to present four interpreters with increasing degree of optimization. That post should serve as a good backgroud before diving into actual JIT-ing.

Another important part of the background puzzle is my How to JIT - an introduction post from 2013; there, I discuss some of the basic tools needed to emit executable x64 machine code at run-time and actually run it on Linux. Please go through it quickly if these things are new to you.

The two phases of JIT

As I wrote previously, the JIT technique is easier to understand when divided into two distinct phases:

Create machine code at program run-time.
Execute that machine code, also at program run-time.

Phase 2 for our BF JIT is exactly identical to the method described in that introductory post. Take a look at the JitProgram class in jit_utils for details. We'll be more focused on phase 1, which will be translating BF to x64 machine code; per the definition quoted in part 1 of the series, we're going to develop an actual BF compiler (compiling from BF source to x64 machine code).

Compilers, assemblers and instruction encoding

Traditionally, compilation was divided into several stages. The actual source language compiler would translate some higher-level language to target-specific assembly; then, an assembler would translate assembly to actual machine code [1]. There's a number of important benefits assembly language provides over raw machine code. Salient examples include:

Instruction encoding: it's certainly nicer to write inc %r13 to increment the contents of register r13 than to write 0x49, 0xFF, 0xC5. Instruction encoding for the popular architectures is notoriously complicated.
Naming labels and procedures for jumps/calls: it's easier to write jl loop than to figure out the encoding for the instruction, along with the relative position of the loop label and encoding the delta to it (not to mention this delta changes every time we add instructions in between and needs to be recomputed). Similarly for functions, call foo instead of doing it by address.

One of my guiding principles through the field of programming is that before diving into the possible solutions for a problem (for example, some library for doing X) it's worth working through the problem manually first (doing X by hand, without libraries). Grinding your teeth over issues for a while is the best way to appreciate what the shrinkwrapped solution/library does for you.

In this spirit, our first JIT is going to be completely hand-written.

Simple JIT - hand-rolling x64 instruction encoding

Out first JIT for this post is simplejit.cpp. Similarly to the interpreters of part 1, all the action happens in a single function (here called simplejit) invoked from main. simplejit goes through the BF source and emits x64 machine code into a memory buffer; in the end, it jumps to this machine code to run the BF program.

Here's its beginning:

std::vector<uint8_t> memory(MEMORY_SIZE, 0);

// Registers used in the program:
//
// r13: the data pointer -- contains the address of memory.data()
//
// rax, rdi, rsi, rdx: used for making system calls, per the ABI.

CodeEmitter emitter;

// Throughout the translation loop, this stack contains offsets (in the
// emitter code vector) of locations for fixup.
std::stack<size_t> open_bracket_stack;

// movabs <address of memory.data>, %r13
emitter.EmitBytes({0x49, 0xBD});
emitter.EmitUint64((uint64_t)memory.data());

As usual, we have our BF memory buffer in a std::vector. The comments reveal some of the conventions used througout the emitted program: our "data pointer" will be in r13.

CodeEmitter is a very simple utility to append bytes and words to a vector of bytes. Its full code is here. It's platform independent except the assumption of little-endian (for EmitUint64 it will write the lowest byte of the 64-bit word first, then the second lowest byte, etc.)

Our first bit of actual machine code emission follows:

// movabs <address of memory.data>, %r13
emitter.EmitBytes({0x49, 0xBD});
emitter.EmitUint64((uint64_t)memory.data());

And it's a cool one, mixing elements from the host (the C++ program doing the emission) and the JITed code. First note the usage of movabs, a x64 instruction useful for placing 64-bit immediates in a register. This is exactly what we're doing here - placing the address of the data buffer of memory in r13. The call to EmitBytes with a cryptic sequence of hex values is preceded by a snippet of assembly in a comment - the assembly conveys the meaning for human readers, the hex values are the actual encoding the machine will understand.

Then comes the BF compilation loop, which looks at the next BF instruction and emits the appropriate machine code for it. Our compiler works in a single pass; this means that there's a bit of trickiness in handling the jumps, as we will soon see.

for (size_t pc = 0; pc < p.instructions.size(); ++pc) {
  char instruction = p.instructions[pc];
  switch (instruction) {
  case '>':
    // inc %r13
    emitter.EmitBytes({0x49, 0xFF, 0xC5});
    break;
  case '<':
    // dec %r13
    emitter.EmitBytes({0x49, 0xFF, 0xCD});
    break;
  case '+':
    // Our memory is byte-addressable, so using addb/subb for modifying it.
    // addb $1, 0(%r13)
    emitter.EmitBytes({0x41, 0x80, 0x45, 0x00, 0x01});
    break;
  case '-':
    // subb $1, 0(%r13)
    emitter.EmitBytes({0x41, 0x80, 0x6D, 0x00, 0x01});
    break;

These are pretty straightforward; since r13 is the data pointer, > and < increment and decrement it, while + and - increment and decrement what it's pointing to. One slightly subtle aspect is that I chose a byte-value memory for our BF implementations; this means we have to be careful when reading or writing to memory and do byte-addressing (the b suffixes on add and sub above) rather than the default 64-bit-addressing.

The code emitted for . and , is a bit more exciting; in the effort of avoiding any external dependencies, we're going to invoke Linux system calls directly. WRITE for .; READ for ,. We're using the x64 ABI here with the syscall identifier in rax:

  // To emit one byte to stdout, call the write syscall with fd=1 (for
  // stdout), buf=address of byte, count=1.
  //
  // mov $1, %rax
  // mov $1, %rdi
  // mov %r13, %rsi
  // mov $1, %rdx
  // syscall
  emitter.EmitBytes({0x48, 0xC7, 0xC0, 0x01, 0x00, 0x00, 0x00});
  emitter.EmitBytes({0x48, 0xC7, 0xC7, 0x01, 0x00, 0x00, 0x00});
  emitter.EmitBytes({0x4C, 0x89, 0xEE});
  emitter.EmitBytes({0x48, 0xC7, 0xC2, 0x01, 0x00, 0x00, 0x00});
  emitter.EmitBytes({0x0F, 0x05});
  break;
case ',':
  // To read one byte from stdin, call the read syscall with fd=0 (for
  // stdin),
  // buf=address of byte, count=1.
  emitter.EmitBytes({0x48, 0xC7, 0xC0, 0x00, 0x00, 0x00, 0x00});
  emitter.EmitBytes({0x48, 0xC7, 0xC7, 0x00, 0x00, 0x00, 0x00});
  emitter.EmitBytes({0x4C, 0x89, 0xEE});
  emitter.EmitBytes({0x48, 0xC7, 0xC2, 0x01, 0x00, 0x00, 0x00});
  emitter.EmitBytes({0x0F, 0x05});
  break;

The comments certainly help, don't they? I hope these snippets are a great motivation for using assembly language rather than encoding instructions manually :-)

The jump instructions are always the most interesting in BF. For [ we do:

case '[':
  // For the jumps we always emit the instruciton for 32-bit pc-relative
  // jump, without worrying about potentially short jumps and relaxation.

  // cmpb $0, 0(%r13)
  emitter.EmitBytes({0x41, 0x80, 0x7d, 0x00, 0x00});

  // Save the location in the stack, and emit JZ (with 32-bit relative
  // offset) with 4 placeholder zeroes that will be fixed up later.
  open_bracket_stack.push(emitter.size());
  emitter.EmitBytes({0x0F, 0x84});
  emitter.EmitUint32(0);
  break;

Note that we don't know where this jump leads at this point - it will go to the matching ], which we haven't encountered yet! Therefore, to keep our compilation in a single pass [2] we use the time-honored technique of backpatching by emitting a placeholder value for the jump and fixing it up once we encounter the matching label. Another thing to note is always using a 32-bit pc-relative jump, for simplicity; we could save a couple of bytes with a short jump in most cases (see my article on assembler relaxation for the full scoop), but I don't think it's worth the effort here.

Compiling the matching ] is a bit trickier; I hope the comments do a good job explaining what's going on, and the code itself is optimized for readability rather than cleverness:

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

  // cmpb $0, 0(%r13)
  emitter.EmitBytes({0x41, 0x80, 0x7d, 0x00, 0x00});

  // open_bracket_offset points to the JZ that jumps to this closing
  // bracket. We'll need to fix up the offset for that JZ, as well as emit a
  // JNZ with a correct offset back. Note that both [ and ] jump to the
  // instruction *after* the matching bracket if their condition is
  // fulfilled.

  // Compute the offset for this jump. The jump start is computed from after
  // the jump instruction, and the target is the instruction after the one
  // saved on the stack.
  size_t jump_back_from = emitter.size() + 6;
  size_t jump_back_to = open_bracket_offset + 6;
  uint32_t pcrel_offset_back =
      compute_relative_32bit_offset(jump_back_from, jump_back_to);

  // jnz <open_bracket_location>
  emitter.EmitBytes({0x0F, 0x85});
  emitter.EmitUint32(pcrel_offset_back);

  // Also fix up the forward jump at the matching [. Note that here we don't
  // need to add the size of this jmp to the "jump to" offset, since the jmp
  // was already emitted and the emitter size was bumped forward.
  size_t jump_forward_from = open_bracket_offset + 6;
  size_t jump_forward_to = emitter.size();
  uint32_t pcrel_offset_forward =
      compute_relative_32bit_offset(jump_forward_from, jump_forward_to);
  emitter.ReplaceUint32AtOffset(open_bracket_offset + 2,
                                pcrel_offset_forward);
  break;
}

This concludes the compiler loop; we end up with a bunch of potentially executable machine code in vector. This code refers to the host program (the address of memory.data()), but that's OK since the host program's lifetime wraps the lifetime of the JITed code. What's remaining is to actually invoke this machine code:

// ... after the compilation loop

// The emitted code will be called as a function from C++; therefore it has to
// use the proper calling convention. Emit a 'ret' for orderly return to the
// caller.
emitter.EmitByte(0xC3);

// Load the emitted code to executable memory and run it.
std::vector<uint8_t> emitted_code = emitter.code();
JitProgram jit_program(emitted_code);

// JittedFunc is the C++ type for the JIT function emitted here. The emitted
// function is callable from C++ and follows the x64 System V ABI.
using JittedFunc = void (*)(void);

JittedFunc func = (JittedFunc)jit_program.program_memory();
func();

The call should be familiar from reading the How to JIT post. Note that here we opted for the simplest function possible - no arguments, no return value; in future sections we'll spice it up a bit.

Taking our JIT for a spin

In part 1, I presented a trivial BF program that prints the numbers 1 to 5 to the screen:

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

Let's see what our compiler translates it to. Even though the code vector inside simplejit is ephemeral (lives only temporarily in memory), we can serialize it to a binary file which we can then disassemble (with objdump -D -b binary -mi386:x86-64). The following is the disassembly listing with comments I embedded to explain what's going on:

 # The runtime address of memory.data() goes into r13; note that this will
 # likely be a different value in every invocation of the JIT.

  0:   49 bd f0 54 e3 00 00    movabs $0xe354f0,%r13
  7:   00 00 00

 # A sequence of 48 instructions that all do the same, for the initial sequence
 # of +s; this makes me miss our optimizing interpreter, by worry not - we'll
 # make this go away later in the post.

  a:   41 80 45 00 01          addb   $0x1,0x0(%r13)
  f:   41 80 45 00 01          addb   $0x1,0x0(%r13)

 # [...] 46 more 'addb'

 # >+++++

 fa:   49 ff c5                inc    %r13
 fd:   41 80 45 00 01          addb   $0x1,0x0(%r13)
102:   41 80 45 00 01          addb   $0x1,0x0(%r13)
107:   41 80 45 00 01          addb   $0x1,0x0(%r13)
10c:   41 80 45 00 01          addb   $0x1,0x0(%r13)
111:   41 80 45 00 01          addb   $0x1,0x0(%r13)

 # Here comes the loop! Note that the relative jump offset is already inserted
 # into the 'je' instruction by the backpatching process.

116:   41 80 7d 00 00          cmpb   $0x0,0x0(%r13)
11b:   0f 84 35 00 00 00       je     0x156
121:   49 ff cd                dec    %r13
124:   41 80 45 00 01          addb   $0x1,0x0(%r13)

 # The '.' is translated into a syscall to WRITE

129:   48 c7 c0 01 00 00 00    mov    $0x1,%rax
130:   48 c7 c7 01 00 00 00    mov    $0x1,%rdi
137:   4c 89 ee                mov    %r13,%rsi
13a:   48 c7 c2 01 00 00 00    mov    $0x1,%rdx
141:   0f 05                   syscall
143:   49 ff c5                inc    %r13
146:   41 80 6d 00 01          subb   $0x1,0x0(%r13)
14b:   41 80 7d 00 00          cmpb   $0x0,0x0(%r13)

 # Jump back to beginning of loop

150:   0f 85 cb ff ff ff       jne    0x121

 # We're done

156:   c3                      retq

How does it perform?

It's time to measure the performance of our JIT against the interpreters from part 1. optinterp3 was about 10x faster than the naive interpreter - how will this JIT measure up? Note that it has no optimizations (except not having to recompute the jump destination for every loop iteration as the naive interpreter did). Can you guess? The results may surprise you...

The simple JIT runs mandelbrot in 2.89 seconds, and factor in 0.94 seconds - much faster still than opt3interp; here's the comparison plot (omitting the slower interpreters since they skew the scale):

Why is this so? opt3interp is heavily optimized - it folds entire loops into a single operation; simplejit does none of this - we've just seen the embarrassing sequence of addbs it emits for a long sequence of +s.

The reason is that the baseline performance of the JIT is vastly better. I've mentioned this briefly in part 1 - imagine what's needed to interpret a single instruction in the fastest interpreter.

Advance pc and compare it to program size.
Grab the instruction at pc.
Switch on the value of the instruction to the right case.
Execute the case.

This requires a whole sequence of machine instructions, with at least two branches (one for the loop, one for the switch). On the other hand, the JIT just emits a single instruction - no branches. I would say that - depending on what the compiler did while compiling the interpreter - the JIT is between 4 and 8 times faster at running any given BF operation. It has to run many more BF operations because it doesn't optimize, but this difference is insufficient to close the huge baseline gap. Later in this post we're going to see an optimized JIT which performs even better.

But first, let's talk about this painful instruction encoding business.

Manually encoding instructions

As promised, simplejit is completely self-contained. It doesn't use any external libraries, and encodes all the instructions by hand. It's not hard to see how painful that process is, and the code is absolutely unreadable unless accompanied by detailed comments; moreover, changing the code is a pain, and changes happen in unexpected ways. For example, if we want to use some other register in an instruction, the change to emitted code won't be intuitive. add %r8, %r9 is encoded as 0x4C, 0x01, 0xC8, but add %r8, %r10 is 0x4C, 0x01, 0xD0; since registers are specified in sub-byte nibbles, one needs very good memory and tons of experience to predict what goes where.

Would you expect related instructions to look somewhat similar? They don't. inc %r13 is encoded as 0x49, 0xFF, 0xC0, for example. To put it bluntly - unless you're Mel, you're going to have a hard time. Now imagine that you have to support emitting code for multiple architectures!

This is why all compilers, VMs and related projects have their own layers to help with this encoding task, along with related tasks like labels and jump computations. Most are not exposed for easy usage outside their project; others, like DynASM (developed as part of the LuaJIT project) are packaged for separate usage. DynASM is an example of a low-level framework - providing instruction encoding and not much else; some frameworks are higher-level, doing more compiler-y things like register allocation. One example is libjit; another is LLVM.

asmjit

While looking for a library to help me encode instructions, I initially tried DynASM. It's an interesting approach - and you can see Josh Haberman's post about using it for a simple BF JIT, but I found it to be a bit too abandonware-ish for my taste. Besides, I don't like the funky preprocessor approach with a dependency on Lua.

So I found another project that seemed to fit the bill - asmjit - a pure C++ library without any preprocessing. asmjit began about 3 years ago to ease its author's development of fast kernels for graphics code. Its documentation isn't much better than dynasm's, but being just a C++ library I found it easier to dive into the source when questions arose the docs couldn't answer. Besides, the author is very active and quick in answering questions on GitHub and adding missing featuers. Therefore, the rest of this post shows BF JITs that use asmjit - these can also serve as a non-trivial tutorial for the library.

simpleasmjit - JIT with sane instruction encoding

Enter simpleasmjit.cpp - the same simple JIT (no optimizations) as simplejit, but using asmjit for the instruction encoding, labels and so on. Just for fun, we'll mix things up a bit. First, we'll change the JITed function signature from void (*)(void) to void (*)(uint64_t); the address of the BF memory buffer will be passed as argument into the JITed function rather than hard-coded into it.

Second, we'll use actual C functions to emit / input characters, rather than system calls. Moreover, since putchar and getchar may be macros on some systems, taking their address can be unsafe. So we'll wrap them in actual C++ functions, whose address it is safe to take in emitted code:

void myputchar(uint8_t c) {
  putchar(c);
}

uint8_t mygetchar() {
  return getchar();
}

simpleasmjit starts by initializing an asmjit runtime, code holder and assembler [3]:

asmjit::JitRuntime jit_runtime;
asmjit::CodeHolder code;
code.init(jit_runtime.getCodeInfo());
asmjit::X86Assembler assm(&code);

Next, we'll give a mnemonic name to our data pointer, and emit a copy of the address of the memory buffer into it (it's in rdi initially, as the first function argument in the x64 ABI):

// We pass the data pointer as an argument to the JITed function, so it's
// expected to be in rdi. Move it to r13.
asmjit::X86Gp dataptr = asmjit::x86::r13;
assm.mov(dataptr, asmjit::x86::rdi);

Then we get to the usual BF processing loop that emits code for every BF op:

for (size_t pc = 0; pc < p.instructions.size(); ++pc) {
  char instruction = p.instructions[pc];
  switch (instruction) {
  case '>':
    // inc %r13
    assm.inc(dataptr);
    break;
  case '<':
    // dec %r13
    assm.dec(dataptr);
    break;
  case '+':
    // addb $1, 0(%r13)
    assm.add(asmjit::x86::byte_ptr(dataptr), 1);
    break;
  case '-':
    // subb $1, 0(%r13)
    assm.sub(asmjit::x86::byte_ptr(dataptr), 1);
    break;

Notice the difference! No more obscure hex codes - assm.inc(dataptr) is so much nicer than 0x49, 0xFF, 0xC5, isn't it?

For input and output we emit calls to our wrapper functions:

case '.':
  // call myputchar [dataptr]
  assm.movzx(asmjit::x86::rdi, asmjit::x86::byte_ptr(dataptr));
  assm.call(asmjit::imm_ptr(myputchar));
  break;
case ',':
  // [dataptr] = call mygetchar
  // Store only the low byte to memory to avoid overwriting unrelated data.
  assm.call(asmjit::imm_ptr(mygetchar));
  assm.mov(asmjit::x86::byte_ptr(dataptr), asmjit::x86::al);
  break;

The magic is in the imm_ptr modifier, which places the address of the function in the emitted code.

Finally, the code handling [ and ] is also much simpler due to asmjit's labels, which can be used before they're actually emitted:

case '[': {
  assm.cmp(asmjit::x86::byte_ptr(dataptr), 0);
  asmjit::Label open_label = assm.newLabel();
  asmjit::Label close_label = assm.newLabel();

  // Jump past the closing ']' if [dataptr] = 0; close_label wasn't bound
  // yet (it will be bound when we handle the matching ']'), but asmjit lets
  // us emit the jump now and will handle the back-patching later.
  assm.jz(close_label);

  // open_label is bound past the jump; all in all, we're emitting:
  //
  //    cmpb 0(%r13), 0
  //    jz close_label
  // open_label:
  //    ...
  assm.bind(open_label);

  // Save both labels on the stack.
  open_bracket_stack.push(BracketLabels(open_label, close_label));
  break;
}
case ']': {
  if (open_bracket_stack.empty()) {
    DIE << "unmatched closing ']' at pc=" << pc;
  }
  BracketLabels labels = open_bracket_stack.top();
  open_bracket_stack.pop();

  //    cmpb 0(%r13), 0
  //    jnz open_label
  // close_label:
  //    ...
  assm.cmp(asmjit::x86::byte_ptr(dataptr), 0);
  assm.jnz(labels.open_label);
  assm.bind(labels.close_label);
  break;
}

We just have to remember which label we used for the jump and emit the exact same Label object - asmjit handles the backpatching on its own! Moreover, all the jump offset computations are performed automatically.

Finally, after emitting the code we can call it:

using JittedFunc = void (*)(uint64_t);

JittedFunc func;
asmjit::Error err = jit_runtime.add(&func, &code);
// [...]
// Call it, passing the address of memory as a parameter.
func((uint64_t)memory.data());

That's it. This JIT emits virtually the same exact code as simplejit, and thus we don't expect it to perform any differently. The main point of this exercise is to show how much simpler and more pleasant emitting code is with a library like asmjit. It hides all the icky encoding and offset computations, letting us focus on what's actually unique for our program - the sequence of instructions emitted.

optasmjit - combining BF optimizations with a JIT

Finally, it's time to combine the clever optimizations we've developed in part 1 with the JIT. Here, I'm essentially taking optinterp3 from part 1 and bolting a JIT backend onto it. The result is optasmjit.cpp.

Recall that instead of the 8 BF ops, we have an extended set, with integer arguments, that conveys higher-level ops in some cases:

enum class BfOpKind {
  INVALID_OP = 0,
  INC_PTR,
  DEC_PTR,
  INC_DATA,
  DEC_DATA,
  READ_STDIN,
  WRITE_STDOUT,
  LOOP_SET_TO_ZERO,
  LOOP_MOVE_PTR,
  LOOP_MOVE_DATA,
  JUMP_IF_DATA_ZERO,
  JUMP_IF_DATA_NOT_ZERO
};

The translation phase from BF ops to a sequence of BfOpKind is exactly the same as it was in optinterp3. Let's take a look at how a couple of the new ops are implemented now:

case BfOpKind::INC_PTR:
  assm.add(dataptr, op.argument);
  break;

As before with the interpreters, an increment of 1 is replaced by the addition of an argument. We use a different instruction for this - add instead of inc [4]. How about something more interesting:

case BfOpKind::LOOP_MOVE_DATA: {
  // Only move if the current data isn't 0:
  //
  //   cmpb 0(%r13), 0
  //   jz skip_move
  //   <...> move data
  // skip_move:
  asmjit::Label skip_move = assm.newLabel();
  assm.cmp(asmjit::x86::byte_ptr(dataptr), 0);
  assm.jz(skip_move);

  assm.mov(asmjit::x86::r14, dataptr);
  if (op.argument < 0) {
    assm.sub(asmjit::x86::r14, -op.argument);
  } else {
    assm.add(asmjit::x86::r14, op.argument);
  }
  // Use rax as a temporary holding the value of at the original pointer;
  // then use al to add it to the new location, so that only the target
  // location is affected: addb %al, 0(%r13)
  assm.mov(asmjit::x86::rax, asmjit::x86::byte_ptr(dataptr));
  assm.add(asmjit::x86::byte_ptr(asmjit::x86::r14), asmjit::x86::al);
  assm.mov(asmjit::x86::byte_ptr(dataptr), 0);
  assm.bind(skip_move);
  break;
}

I'll just note again how much simpler this code is to write with asmjit than without it. Also note the careful handling of the byte-granulated data when touching memory - I ran into a number of nasty bugs when developing this. In fact, using the native machine word size (64 bits in this case) for BF memory cells would've made everything much simpler; 8-bit cells are closer to the common semantics of the language and provide an extra challenge.

Performance

Let's see how optasmjit fares against the fastest interpreter and the unoptimized JIT - 0.93 seconds for mandelbrot, 0.3 seconds for factor - another factor of 3 in performance:

Notably, the performance delta with the optimized interpreter is huge: the JIT is more than 4x faster. If we compare it all the way to the initial simple interpreter, optasmjit is about 40x faster - making it hard to even compare on the same chart :-)

JITs are fun!

I find writing JITs lots of fun. It's really nice to be able to hand-craft every instruction emitted by the compiler. While this is quite painful to do without any encoding help, libraries like asmjit make the process much more pleasant.

We've done quite a bit in this part of the series. optasmjit is a genuine optimizing JIT for BF! It:

Parses BF source
Translates it to a sequence of higher-level ops
Optimizes these ops
Compiles the ops to tight x64 assembly in memory and runs them

Let's connect these steps to some real compiler jargon. BfOpKind ops can be seen as the compiler IR. Translation of human-readable source code to IR is often the first step in compilation (though it in itself is sometimes divided into multiple steps for realistic languages). The translation/compilation of ops to assembly is often called "lowering"; in some compilers this involves multiple steps and intermediate IRs.

I left a lot of code out of the blog post - otherwise it would be huge! I encourage you to go back through the full source files discussed here and understand what's going on - every JIT is a single standalone C++ file.

Links to all posts in this series:

[1]

I said traditionally because many modern compilers no longer work this way. For example, LLVM compiles IR to another, much lower-level IR that represents machine-code level instructions; assembly can be emitted from this IR, but also machine code directly - so the assembler is integrated into the compiler.

[2]	Some compilers would do two passes; this is similar to our first interpreter optimization in part 1: the first pass collects information (such as location of all matching `]`s), so the second pass already knows what offsets to emit.

[3]	Please refer to asmjit's documentation for the full scoop. I'll also mention that asmjit has a "compiler" layer which does more sophisticated things like register allocation; in this post I'm only using the base assembly layer.

[4] Wondering whether we could have just used add 1 instead of inc in the first place? Certainly! In fact, while there probably used to be a good reason for a separate inc instruction, in these days of complex multi-port pipelined x64 CPUs, it's not clear which one is faster. I just wanted to show both for diversity.

Adventures in JIT compilation: Part 1 - an interpreter

2017-03-20T06:25:00-07:00

This is the first post in a series about JIT compilers. The plan is to take a simple input language and develop some interpreters and JITs for it, in roughtly increasing degree of complexity. It's my hope that by the end of the series readers will have a good understanding of what it takes to develop a JIT compiler and what are some of the tools available to assist with the task.

The input language will be Brainfuck, or BF as I'll be referring to it from now and throughout the series. I think it's a good language for the purpose since it really boils programmability down to the essentials. Even though it's very tedious to program in, BF is fairly "mainstream" as far as programming languages go, with some concepts like memory pointers and loops mapping directly to familiar C constructs.

As the implementation language I'll be using C++. This is, perhaps, not the most commonly used "starter" language; that said, most compilers I know are written in C++ (or C), and hence many of the most popular low-level code-generation libraries in existence are in these languages. In later parts of this series we'll be using some C++ libraries, and this is by far easiest to do from C++ itself. Besides, I try to keep my code straightforward throughout the series - there is very little use of advanced C++ features here.

The BF language

The BF language is simple to describe, but I don't do this here. Please take a look at the spec, read the Wikipedia page, or one of the other existing resources. An in-browser interpreter such as this one can be very useful.

I'll just give an example to develop a taste for the language. The following BF program prints the numbers 1 to 5 to the screen:

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

Here's what it does:

Line 1 initializes memory cell 0 to the value 48, which happens to be the ASCII code for 0.
Line 2 initializes memory cell 1 to 5, which is our loop counter.
Line 3 is a loop that, at each iteration, increments cell 0 and prints its value out, then decrements cell 1 and checks if it has reached the value 0.

A simple interpreter

To get an initial feel for the language and to have a reliable reference implementation, we'll start with a simple interpreter that processes one BF character at a time and does what's necessary to "execute" it.

One of the reasons for my choosing BF as the source language is its simplicity. You'll find a lot of tutorials online that purport to develop interpreters or compilers but end up focusing 90% of their time on writing the parser. I think the later stages of compilation are much more interesting, so my "parser" for BF looks like this:

struct Program {
  std::string instructions;
};

Program parse_from_stream(std::istream& stream) {
  Program program;

  for (std::string line; std::getline(stream, line);) {
    for (auto c : line) {
      if (c == '>' || c == '<' || c == '+' || c == '-' || c == '.' ||
          c == ',' || c == '[' || c == ']') {
        program.instructions.push_back(c);
      }
    }
  }
  return program;
}

Note that this is a valid implementation according to the BF spec: all characters except the 8 supported ones are to be treated as comments and ignored [1]. This parser is going to serve us throughout the series.

With that out of the way, here's the actual interpreter:

constexpr int MEMORY_SIZE = 30000;

void simpleinterp(const Program& p, bool verbose) {
  // Initialize state.
  std::vector<uint8_t> memory(MEMORY_SIZE, 0);
  size_t pc = 0;
  size_t dataptr = 0;

  while (pc < p.instructions.size()) {
    char instruction = p.instructions[pc];
    switch (instruction) {
    case '>':
      dataptr++;
      break;
    case '<':
      dataptr--;
      break;
    case '+':
      memory[dataptr]++;
      break;
    case '-':
      memory[dataptr]--;
      break;
    case '.':
      std::cout.put(memory[dataptr]);
      break;
    case ',':
      memory[dataptr] = std::cin.get();
      break;

    // [...]

All these cases are rather trivial. The more interesting ones are the control flow ops - [ and ]. We'll start with [ - jump forward if the current data location is zero. This op makes it possible to skip a loop or implement a simple if-like condition.

case '[':
  if (memory[dataptr] == 0) {
    int bracket_nesting = 1;
    size_t saved_pc = pc;

    while (bracket_nesting && ++pc < p.instructions.size()) {
      if (p.instructions[pc] == ']') {
        bracket_nesting--;
      } else if (p.instructions[pc] == '[') {
        bracket_nesting++;
      }
    }

    if (!bracket_nesting) {
      break;
    } else {
      DIE << "unmatched '[' at pc=" << saved_pc;
    }
  }
  break;

The most important thing to note here is that the [ and ] brackets in BF can be nested; therefore, when figuring out where to jump, we have to find the matching bracket. If this seems like something wasteful to do at run-time, you're right - keep reading!

For ] we do something very similar. In BF, ] is jumping to an earlier [ if the current data location is not zero. This is how loops advance to the next iteration (or stop).

case ']':
  if (memory[dataptr] != 0) {
    int bracket_nesting = 1;
    size_t saved_pc = pc;

    while (bracket_nesting && pc > 0) {
      pc--;
      if (p.instructions[pc] == '[') {
        bracket_nesting--;
      } else if (p.instructions[pc] == ']') {
        bracket_nesting++;
      }
    }

    if (!bracket_nesting) {
      break;
    } else {
      DIE << "unmatched ']' at pc=" << saved_pc;
    }
  }
  break;

And that's it! The interpreter loop concludes with:

  default: { DIE << "bad char '" << instruction << "' at pc=" << pc; }
  }

  pc++;
}

The full code for this simple interpreter can be found in simpleinterp.cpp.

Measuring the performance of BF programs

Whenever we develop something like an interpreter or compiler, execution speed is a paramount concern. Therefore, it's common for compiler writers to have benchmark suites they refer to for measurements. For BF, I'll be using a couple of programs throughout the series to measure how fast our implementation is. One is a Mandelbrot generator; another is a factorization program invoked on the large-ish prime 179424691 [2].

The simple interpreter shown above takes 38.6 seconds on mandelbrot and 16.5 seconds on factor [3]. Now let's see how we can greatly improve these numbers.

Optimized interpreter - take 1

The most obvious optimization opportunity for the simple interpreter is to avoid laboriously looking for the matching bracket every time a [ or ] is encountered. Imagine a realistic program with a hot inner loop (by "hot" here I mean it runs many, many - possibly billions - of times throughtout the execution of the program). Is it really necessary to scan the source to find the matching bracket every single time? Of course not. We can just precompute these jump destinations ahead of time, since the BF program doesn't change throughout its execution.

This is the idea behind optinterp.cpp - our first optimized interpreter. Much of the code is the same as for the simple interpreter, so I'll just highlight the differences. A crucial addition is this function, which is run before the actual interpretation happens:

std::vector<size_t> compute_jumptable(const Program& p) {
  size_t pc = 0;
  size_t program_size = p.instructions.size();
  std::vector<size_t> jumptable(program_size, 0);

  while (pc < program_size) {
    char instruction = p.instructions[pc];
    if (instruction == '[') {
      int bracket_nesting = 1;
      size_t seek = pc;

      while (bracket_nesting && ++seek < program_size) {
        if (p.instructions[seek] == ']') {
          bracket_nesting--;
        } else if (p.instructions[seek] == '[') {
          bracket_nesting++;
        }
      }

      if (!bracket_nesting) {
        jumptable[pc] = seek;
        jumptable[seek] = pc;
      } else {
        DIE << "unmatched '[' at pc=" << pc;
      }
    }
    pc++;
  }

  return jumptable;
}

It computes the jump destinations for all the [ and ] ops in the program. Its operation is essentially identical to the scanning forward / backward for a matching bracket in the main loop of the simple interpreter. The result is the vector jumptable, where for every [ and ] at offset i in the program, jumptable[i] holds the offset of the matching bracket. For any other op at offset i, jumptable[i] is simply 0.

The actual main loop of optinterp is the same as in simpleinterp, except for the clauses for brackets, which become simply:

case '[':
  if (memory[dataptr] == 0) {
    pc = jumptable[pc];
  }
  break;
case ']':
  if (memory[dataptr] != 0) {
    pc = jumptable[pc];
  }
  break;

As you'd expect, optinterp is quite a bit faster; it takes only 18.4 seconds to run mandelbrot and 6.7 seconds to run factor - more than a factor of 2 improvement!

Optimized interpreter - take 2

The optimization applied in the previous section was very beneficial, but it's also fairly trivial - we avoid completely unnecessary work at run-time, if we can just precompute it at compile time. To make our interpreter even faster, we'll have to get more creative.

The first step in optimizing anything is measuring and profiling the current code. Some past experience helps avoid needless steps in this process. For example, it's fairly clear that almost 100% of the run-time of the interpreter will be spent in the single function that interprets the program; therefore, function/call profiling won't be of much help.

The main loop is fairly small, however, and there doesn't appear to be much to optimize at first glance (disregarding micro-optimizaitons which I won't worry about here). Well, except that this loop runs for every BF instruction encountered in the source program, so it can run tons of times. So what we'll do is get a breakdown of the ops that execute during a typical program run. The code for optinterp already has this tracing included - it's protected with the BFTRACE preprocessor macro because it's costly and we want to avoid doing it in "real" runs.

Here's the execution profile we get from a typical run of the factor benchmark on the prime 179424691. On the left is the operation, and on the right the number of times it was executed by the interpreter for the program at hand:

.  -->  21
,  -->  10
+  -->  212,428,900
]  -->  242,695,606
<  -->  1,220,387,704
-  -->  212,328,376
>  -->  1,220,387,724
[  -->  118,341,127
.. Total: 3,226,569,468

A couple of immediate observations:

The total number of operations is huge: over 3 billion times around the main interpreter loop. It's a good thing we're using C++ for the interpreter - running 3 billion iterations of anything in a higher level language would be painful!
The ratio of pointer movement instructions to loops is suspiciously high. There's something like 242 million loop iterations executed (the count for ]) but a total of 2.4 billion pointer moves: < and >.

We'd expect and hope for the hot loops to be short and tight - why is every loop doing so much?

A cursory glance at the source of factor.bf provides a clue. Here's a representative snippet:

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

Note the fairly long sequences of <<<<<<<, >>>>>>>. Just a bit of thought about the semantics of BF makes this clear - these are necessary to get anything done because we want to be able to get from cell to cell to update data.

Now let's think what it means to execute a sequence such as <<<<<<< in our interpreter. Our main loop executes 7 times, each time doing:

Advance pc and compare it to program size.
Grab the instruction at pc.
Switch on the value of the instruction to the right case.
Execute the case.

That's quite expensive. What if we could compress all the long sequences of <<<<<<<? After all, what we do for a single < is:

case '<':
  dataptr--;
  break;

So for seven <s we could do:

case ... something representing 7 <s ...:
  dataptr -= 7;
  break;

This is easy to generalize. We can detect any consecutive sequence in the BF source and encode it as a pair: the operation, and the repetition count. Then at execution time we simply repeat the op the required number of times.

The full code for this interpreter is optinterp2.cpp. Previously, we kept a separate jump table correlated to the [ and ] instructions in the input program. Now we need extra information for every BF instruction, so we'll just translate the Program into a sequences of ops of the type:

enum class BfOpKind {
  INVALID_OP = 0,
  INC_PTR,
  DEC_PTR,
  INC_DATA,
  DEC_DATA,
  READ_STDIN,
  WRITE_STDOUT,
  JUMP_IF_DATA_ZERO,
  JUMP_IF_DATA_NOT_ZERO
};

// Every op has a single numeric argument. For JUMP_* ops it's the offset to
// which a jump should be made; for all other ops, it's the number of times the
// op is to be repeated.
struct BfOp {
  BfOp(BfOpKind kind_param, size_t argument_param)
      : kind(kind_param), argument(argument_param) {}

  BfOpKind kind = BfOpKind::INVALID_OP;
  size_t argument = 0;
};

The interpretation happens in two steps. First we run translate_program to read the program and generate a std::vector<BfOp>. This translation is pretty straight-forward: it detects repetitions in ops like < and encodes them in the argument field. A slightly tricky aspect here is handling the jumps, since the offsets of all ops in the program change (a run of seven consecutive <s turns into a single DEC_PTR, for example). Take a look at the code for the full details.

As planned, the main interpreter loop becomes:

switch (kind) {
case BfOpKind::INC_PTR:
  dataptr += op.argument;
  break;
case BfOpKind::DEC_PTR:
  dataptr -= op.argument;
  break;
case BfOpKind::INC_DATA:
  memory[dataptr] += op.argument;
  break;
case BfOpKind::DEC_DATA:
  memory[dataptr] -= op.argument;
  break;
// [...] etc.

How fast is it? The mandelbrot benchmark now takes 11.9 seconds, and factor takes 3.7 seconds; another 40% reduction in run-time.

Optimized interpreter - take 3

Our optimized interpreter now runs the mandelbrot benchmark more than 3x faster than the original, naive interpreter. Can we do even better?

First, let's take a look at instruction tracing for optinterp2, repeating the previous experiment:

.  -->  21
]  -->  242,695,606
,  -->  10
+  -->  191,440,613
<  -->  214,595,790
-  -->  205,040,514
>  -->  270,123,690
[  -->  118,341,127
.. Total: 1,242,237,371

The total instruction count went down almost 3x. Also, now the number of BF loop executions is more comparable to the number of other instructions, meaning that we don't do too much work in every iteration. This was our goal with the optimization of repetitions, after all.

In fact, this execution profile is annoyingly flat. Performance gurus don't like flat profiles because there's nothing in particular sticking out that one could optimize. This usually means we should measure / trace something else as well.

An interesting question worth answering is - what is every BF loop doing. In other words, what are the hottest loops we are running, and can we spend some more specialized effort to optimize them? This would require more sophisticated tracing machinery, which I've already included in the code of optinterp2. This machinery traces loops and records the instruction sequence executed by each loop iteration in the program. It then sorts them by the number of appearances and shows the most common (hottest) loops. Here is the result for the factor benchmark:

-1<10+1>10      --> 32,276,219
-1              --> 28,538,377
-1<4+1>4        --> 15,701,515
-1>3+1>1+1<4    --> 12,581,941
-1>3+1>2+1<5    --> 9,579,970
-1<3+1>3        --> 9,004,028
>3              --> 8,911,600
-1<1-1>1        --> 6,093,976
-1>3+1<3        --> 6,085,735
-1<1+1<3+1>4    --> 5,853,530
-1>3+2<3        --> 5,586,229
>2              --> 5,416,630
-1>1+1<1        --> 5,104,333

What do these traces mean? The first, most common one says:

Decrement current memory cell
Move 10 cells to the left
Increment current memory cell
Move 10 cells to the right

The loop doing this was executed 32 million times! Similarly, a loop doing the simple "decrement the current cell" was executed 28 million times. If you look in the source of factor.bf, these loops are easy to spot. The first one is [-<<<<<<<<<<+>>>>>>>>>>]; the second one is just [-].

What if we could optimize these loops entirely away? After all, they are doing something that is much easier to express in a higher-level language. [-] merely sets the current memory cell to 0. [-<<<<<<<<<<+>>>>>>>>>>] is more involved, but not much more: all it does is add the value of the current memory cell 10 cells to the left. The trace shown above features many loops of this kind, along with another; loops like [>>>] move to the right in jumps of 3 until encountering a non-zero cell.

In optinterp2 we've added higher-level ops to the interpreter. We can add some even higher-level ops to optimize away these loops. optinterp3.cpp does just that. It adds a few more operation kinds for encoding common loops:

enum class BfOpKind {
  INVALID_OP = 0,
  INC_PTR,
  DEC_PTR,
  INC_DATA,
  DEC_DATA,
  READ_STDIN,
  WRITE_STDOUT,
  LOOP_SET_TO_ZERO,
  LOOP_MOVE_PTR,
  LOOP_MOVE_DATA,
  JUMP_IF_DATA_ZERO,
  JUMP_IF_DATA_NOT_ZERO
};

The new ops are LOOP_SET_TO_ZERO which replaces [-], LOOP_MOVE_PTR for loops like [>>>] and LOOP_MOVE_DATA for loops like [-<<<+>>>]. We'll now need a slightly more sophisticated translation step that detects these loops in the input program and emits the proper LOOP_* ops. For an example of how it's done, here's the translation for [-]:

std::vector<BfOp> optimize_loop(const std::vector<BfOp>& ops,
                                size_t loop_start) {
  std::vector<BfOp> new_ops;

  if (ops.size() - loop_start == 2) {
    BfOp repeated_op = ops[loop_start + 1];
    if (repeated_op.kind == BfOpKind::INC_DATA ||
        repeated_op.kind == BfOpKind::DEC_DATA) {
      new_ops.push_back(BfOp(BfOpKind::LOOP_SET_TO_ZERO, 0));

  // [...]

This function is called when translating the BF program to a sequence of ops. loop_start is the index in ops where the most recent loop starts. The code shown above detects the case where the only contents of the loop is a single - (or + since in BF memory cells hold unsigned values with wrap-around). In such cases, a LOOP_SET_TO_ZERO op is emitted. When the interpreter itself runs into a LOOP_SET_TO_ZERO, it does just what you'd expect:

case BfOpKind::LOOP_SET_TO_ZERO:
  memory[dataptr] = 0;
  break;

The other loop optimizations are a tiny bit more involved, but all follow the same basic idea.

We expect this optimization to be very significant - we've just taken some of the hottest loops the program runs and folded them into a single, efficient instruction (or a sequence of efficient instructions for pointer-movement loops). And indeed, optinterp3 is very fast: 3.9 seconds on mandelbrot and 1.97 seconds on factor.

The overall speedup is dramatic. optinterp3 is almost 10x faster than simpleinterp on our benchmarks. While we could certainly make it even faster, I think these optimizations are sufficient for our needs; let's talk about what we can learn from them instead.

On compilers, bytecode and tracing JITs

It turns out there's a surprising amount of insight to be gained from the exercise this post went through.

First, let's start with the distinction between compilers and interpreters. According to Wikipedia, a compiler is:

a computer program (or a set of programs) that transforms source code written in a programming language (the source language) into another computer language (the target language), with the latter often having a binary form known as object code.

gcc would be the canonical example of this: it transforms source code written in C (or C++) into assembly language for, say, Intel CPUs. But there are many other kinds of compilers: Chicken compiles Scheme into C; Rhino compiles Javascript to JVM bytecode; the Clang frontend compiles C++ to LLVM IR. CPython, the canonical Python implementation compiles Python source code to bytecode, and so on. In general, the term Bytecode refers to any intermediate representation / virtual instruction set designed for efficient interpretation.

Based on this definition, while simpleinterp is indeed just a BF interpreter, the optimized interpreters described here are more like compilers + bytecode interpreters. Consider optinterp3 for example. The source language is BF; the target language is bytecode with the following instruction set:

enum class BfOpKind {
  INVALID_OP = 0,
  INC_PTR,
  DEC_PTR,
  INC_DATA,
  DEC_DATA,
  READ_STDIN,
  WRITE_STDOUT,
  LOOP_SET_TO_ZERO,
  LOOP_MOVE_PTR,
  LOOP_MOVE_DATA,
  JUMP_IF_DATA_ZERO,
  JUMP_IF_DATA_NOT_ZERO
};

... where each instruction has a single argument. optinterp3 first compiles BF to this bytecode, and only then executes the bytecode. So if we squint a bit, there's a JIT compiler here already - with the caveat that the compilation target is not executable machine code but rather this specialized bytecode. Worry not - we'll get to a real JIT in the next part of the series.

Finally, I'd like to point out that the loop optimizations performed in optinterp3 are the static version of a tracing JIT. We used tracing to observe which loops occur most commonly in our benchmarks, and optimized these loops. While the loops we optimized were very generic and surely appear in most large BF programs, we could take it further. We could optimize even more of the hottest loops, but with time we'd get into more specialized paths in our particular benchmarks.

To be fully generic, we'd have to defer this optimization to run-time, which is what a tracing JIT does. A tracing JIT interprets code in the source language and keeps track of the hottest loops (and for dynamic languages, of the actual run-time types of values flowing through the loops). When the same loop goes over some threshold (say, invoked more than a million times) the loop is optimized and compiled into efficient machine code.

Parting words for part 1

In this post, we've seen an interpreter for BF being gradually refined from a naive approach to a compile-to-optimized-bytecode approach, speeding it up 10x in the process. Hopefully this provides a good feel for the source language, as well as some of the tradeoffs involved in optimizing for it.

In the next part of the series I'll present an actual JIT for BF - compiling BF into tight x64 machine code and invoking it, all at run-time. I'll show how to construct such a JIT compiler entirely from scratch (using nothing but the standard system libraries) and also how to use assembly encoding libraries for easier development. Stay tuned - it's going to be fun!

Links to all posts in this series:

[1]

An observant reader will note I could have used a switch statement or a lookup table here for more efficiency. I'm an avid adherent to the "until you've measured it, it ain't slow" philosophy. The parsing stage of BF takes negligible time for any realistic program, so it's hardly important to optimize this part. On my machine, the largest BF program I could find (Mandelbrot) is ~11,000 instructions and takes 360 microseconds to parse, most of which is almost certainly dominated by the file read time.

[2]

I'm using two different programs to prevent overfitting for one particular benchmark. Naturally it would be more professional to use a whole suite of benchmarks, but this is just a hobby blog post so let's not overdo it B-)

In fact, even real-world benchmark suites for large projects tend to be poor approximations for the real world.

[3]

All the performance numbers for this series were collected on my Haswell Linux (Ubuntu 14.04) box; the C++ code was compiled with the default gcc 4.8.4. I got slightly different results with Clang in some cases, but this series is not about comparing host C++ compilers, so I'll just use gcc throughout.

It's also worth noting these execution times are end-to-end for the whole binary and include loading the interpreter binary, reading the BF file, doing any pre-processing / optimization / code emission and actually running the benchmark. For any run time beyond ~100 ms, all secondary factors are negligible compared to the actual benchmark run-time (the time it took to execute the BF program at hand).

Some thoughts on LLVM vs. libjit

2014-01-15T05:49:56-08:00

Having recently completed a series of articles on using libjit and understanding how it works, I couldn't stop comparing it to LLVM inside my head. This is hardly surprising, since LLVM has been a significant part of my professional life for the past 3.5 years, and will remain in this position in the foreseeable future.

So this post is some unfiltered thoughts on libjit and LLVM - what's similar, what's different, and which one is more suitable for new projects.

First, a tiny bit of background. LLVM hardly needs introduction for the readers of my blog, but just for completeness - it was first released in late 2003. However, according to Chris Lattner's (one of the creators of LLVM and its main developer for many years) recent presentation, the first ideas for it came together in 2000 and the first prototype was hacked together in early 2001.

Today, LLVM is one of the most significant infrastructural backbones of modern computing, no less. It's the default compiler for the OS X and iOS platforms (FreeBSD is deprecating gcc in favor of LLVM & Clang too), and is heavily used in production by almost all major software companies. New things pop up all the time, like the (Windows-based, mind you) toolchain of the new Sony PS4 being completely built on top of LLVM.

libjit's first release was in April 2004, as part of the now-defunct DotGNU project. This was a much less complete release than LLVM 1.0, though (version 0.0.0f - doesn't sound very reassuring for running your nuclear reactor, does it?). Its goals, from the beginning, were quite similar to LLVM's, if somewhat more modest. It also aimed to provide the backend of a compiler from a target-independent IR.

Some internet archaeology brought up this fascinating mailing list thread from May 2004, in which Chris Lattner asks Rhys Weatherley (the creator of libjit) whether the two projects can "join forces", since their goals are similar. If you have any interest in open-source, you can't miss this discussion - go read it now. After a few weak justification attempts, Chris and Rhys got to the real issue. In open-source, people often start new projects just because it's fun, as well as to retain complete control. There may be an existing project that does something similar. The effort to use this existing project is, in all likeness, smaller than the effort to roll your own. But what would you prefer - spending time adapting a large body of existing code you didn't write (and whose brace style you hate!) to your needs and quibbling over methodology on mailing lists, or gloriously hacking into the night - your fingers burning holes in the keyboard producing new code? Especially when this isn't a "day job", but something you do for fun... Yeah, I thought so. Joel's In Defense of the NIH syndrome is a classic now (it's 13 years old, oh my) but it rings as true, and possibly truer, for open-source as for the corporate setting Joel first aimed it at.

But I digress... The short version of the paragraph above is that libjit and LLVM really do have similar goals, and if the mood of their overlords had been different on a few pleasant spring nights of 2004, we could've had just a single project today; that would, in all likeness, be LLVM. Could this also help the DotGNU project on the whole, which lost more and more steam over the years until its official abandonment in 2012? Who knows... If you have a time machine, let me know please.

But lo and behold - libjit is still alive today. It seems to live in a state of life support because there's not too much active development going on, but it's definitely not dead. Its current maintainer is fairly responsive and seems keen to fix things, including documentation.

As I mentioned, libjit's initial goals were much less ambitious than LLVM's, and it remained so. LLVM's philosophy of "everything is a pluggable library" kept driving its development, and if you look at LLVM today - everything is indeed quite pluggable, including newer parts like MC. LLVM supports pluggability in its very core. The so called mid-level IR infrastructure (arguably LLVM's best-designed part) is pluggability executed to perfection. You write generic passes that analyze and transform IR, and hook them together in any way you see fit (along with dozens of industrial-strength compilation passes already provided by LLVM). But the backend parts too: you can choose from more than one instruction selector (or write your own), multiple register allocators, schedulers, code emitters, assemblers, and so on.

libjit is not like that. While it was designed to be able to emit code for multiple architectures (it has both flavors of x86, some experimental ARM support as well as old "attic" code for Alpha), its internals are not modular; sure, you can add more optimizations (or a different register allocator) to libjit, but you'll have to hack it into libjit itself - the requisite APIs and data structures are not really exposed on the library level, the way they are with LLVM.

The libjit IR is also somewhat more limited, having been designed with a single front-end in mind (DotGNU). True, the designer aimed to make it fairly abstract, but some things are definitely missing compared to LLVM, which has been shape by years of use in multiple different front-ends. For example, the type system of libjit is far less flexible, supporting only a subset of integer and floating-point types (like int32, not the arbitrary LLVM jungle of int73 if you want it).

But generality comes at a cost: LLVM is both relatively slow and its code is relatively difficult to grok. It's not slow for a regular (AOT - Ahead of Time) compiler, but it is slow for a JIT compiler, a common source of pain for some LLVM users.

Now, a disclaimer: I did not benchmark libjit's compilation speed vs. LLVM; I was simply too lazy to generate equivalent and large-enough inputs for both. I'm also not saying that libjit is faster than LLVM, perhaps it isn't. After all, the amount of engineering power expended on LLVM has been orders of magnitude larger than for libjit, which may very well have led to much more optimized code. All I want to imply here is that libjit could be faster than LLVM. Generality is almost always at odds with performance in software, a sad fact we all wish wasn't true but, oh, it is. Is this very important? Not for most uses of LLVM today, but since libjit's main goal is JIT, then perhaps. But without real measurements and profiling, I don't have anything intelligent to contribute here.

libjit is also much simpler to understand. I've spent enough time with LLVM to appreciate its complexity, as well as see how confused many programmers are about it. libjit was dramatically easier to grok for me. This may have two reasons: one is definitely its significantly smaller size and simpler internal structure. There is simply much less code to read, and much less abstraction levels to keep in the head. The other reason is my (sometimes controversial) opinion about the choice of programming language. libjit is written in C, which IMHO is easier to understand than C++. But that's a whole other issue, unrelated to this post (you know where to find me for your hatemail).

There's another significant difference between LLVM and libjit. LLVM, despite the historical meaning of its acronym and the numerous attempts that have been made over the years to make it work for dynamic languages, is still in its core a backend for compiling static languages like C, C++ and Objective C. By this, I mean languages that don't really have a significant runtime, like Java or C#. libjit, on the other hand, was designed to serve as a backend for a .NET implementation. That said, like the rest of libjit, these parts are not really battle-proven, so it's hard to attest to their completeness and stability.

So, which one should you use or learn? As for use, there's no question; LLVM is an industrial strength, heavily tested and verified system. It's being used daily on millions (maybe billions) of devices; it's being hacked on by hundreds of programmers from dozens of companies. If it fits your needs, LLVM is the way to go. And these days it fits most needs. Not all of them, mind you. For example, as I've already mentioned above, LLVM is not great for fast JIT-ing. And to be honest, I don't think it's possible to create a really fast JIT within the framework of LLVM, because of its modularity. The faster the JIT, the more you'll have to deviate from the framework of LLVM. This is a serious problem for Portable Native Client, for instance.

libjit, on the other hand, is much more limited, aimed at dynamic runtimes and can (potentially) JIT-compile faster. It's also not being used too much, and lacks credibility w.r.t. features and stability. libjit itself only comes with toy examples - not very reassuring if you want to bet your future project on it.

In terms of educational value, libjit is great. If you want to learn about compiler backends, I would definitely start with libjit and only move to LLVM later. I hope my series of articles will be useful here.