Just-In-Time compilation of OCaml byte-code

This paper presents various improvements that were applied to OCamlJIT2, a Just-In-Time compiler for the OCaml byte-code virtual machine. OCamlJIT2 currently runs on various Unix-like systems with x86 or x86-64 processors. The improvements, including the new x86 port, are described in detail, and performance measures are given, including a direct comparison of OCamlJIT2 to OCamlJIT.


Introduction
The OCaml system [11,16] is the main implementation of the Caml language [6], featuring a powerful module system combined with a full-fledged object-oriented layer. It comes with an optimizing native code compiler ocamlopt, a byte-code compiler ocamlc with an associated runtime ocamlrun, and an interactive top-level ocaml. OCamlJit [18] and OCamlJit2 [15] provide Just-In-Time compilers for the byte-code used by ocamlrun and ocaml.
We describe a set of improvements that were applied to OCamlJit2, including a new port to the x86 architecture 1 and some interesting optimizations to further improve the execution speed of the JIT engine. OCamlJit2 is open-source 2 and is verified to work with Linux/amd64, Linux/i386 and Mac OS X (64-bit Intel Macs only), but should also run on other Unix-like systems with x86 or x86-64 processors.
The paper is organized as follows: Section 2 mentions existing systems and relevant sources of information. Section 3 describes the improvements applied to OCamlJit2, in particular the new x86 port and various floating-point optimizations. Detailed performance measures are given in section 4, including a direct comparison with OCamlJit. Section 5 concludes with possible directions for further work.

Existing systems
The implementation of OCamlJit2 described below is based on a previous prototype of OCamlJit2 [15], and earlier work on OCamlJit [18] and OCaml [10,11,12,16]. We assume that the reader is familiar with the internals of the OCaml byte-code and runtime, and the structure of the OCamlJit and OCamlJit2 Just-In-Time compilers. An overview of the relevant parts of the OCaml virtual machine is given in both [15] and [18], while [16] provides the necessary insights on the OCaml language.

Improvements
This section describes the improvements that were applied to the OCamlJit2 prototype described in [15]. This includes the new x86 port (section 3.1), a simple native code managment strategy (section 3.2), as well as improvements in the implementation of floating-point primitives (section 3.3), which reduced the execution time of floating-point benchmarks by up to 30%. Readers only interested in the performance results may skip straight to section 4.

32-bit support
OCamlJit2 now supports both x86-64 processors (operating in long mode) as well as x86 processors (operating in protected mode). This section provides a brief overview of the mapping of the OCaml virtual machine to the x86 architecture, especially the mapping of the virtual machine registers to the available physical registers. See [15] for a description of the implementation details of the x86-64 port.
The x86 architecture [9] provides 8 general purpose 32-bit registers %eax, %ebx, %ecx, %edx, %ebp, %esp, %edi and %esi, as well as 8 80-bit floating-point registers organized into a so-called FPU stack (and also used as 64-bit MMX registers). Recent incarnations also include a set of 8 SSE2 registers %xmm0, . . ., %xmm7. The System V ABI for the x86 architecture [17], implemented by almost all operating systems running on x86 processors, except Win32 which uses a different ABI, mandates that registers %ebp, %ebx, %edi, %esi and %esp belong to the caller and the callee is required to preserve their values across function calls. The remaining registers belong to the callee.
To share as much code as possible with the x86-64 port, we use a similar register assignment for the x86 architecture. This includes making good use of callee-save registers to avoid saving and restoring too many aspects of the machine state whenever a C function is invoked. Our register assignment therefore looks as follows: The virtual register accu is mapped to %eax. extra args goes into %ebx, whichjust like on x86-64 -contains the number of extra arguments as tagged integer. The environment pointer env goes into %ebp, and the stack pointer sp goes into %esi. %edi contains the cached value of the minor heap allocation pointer caml young ptr to speed up allocations of blocks in the minor heap; this is different from ocamlopt on x86, where caml young ptr is not cached in a register (unlike x86-64 where ocamlopt caches its value in %r15).
The setup and helper routines that form the runtime of the JIT engine are located in the file byterun/jit rt i386.S. These are mostly copies of their x86-64 counterparts in the file byterun/jit rt amd64.S, adapted to the x86 architecture. The adaption was mostly straight-forward, replacing the x86-64 registers with their x86 counterparts and the 64-bit opcodes with their 32-bit counterparts.

Address mapping and on-demand compilation
We use the existing scheme [15] to map byte-code to native machine code address, which replaces the instruction opcode word within the byte-code sequence with the offset of the generated native code relative to the base address caml jit code end. Whenever jumping to a byte-code address -i.e. during closure application or return -this offset is read from the instruction opcode, caml jit code end is added to it, and a jump to the resulting address is performed. The trampoline code for x86 -shown in Figure 1 -is therefore quite similar to the trampoline code for x86-64 (%ecx contains the address of the byte-code to execute and %edx is a temporary register). movl (%ecx), %edx addl caml_jit_code_end, %edx jmpl *%edx Adapting the on-demand compilation driver was also straight-forward, due to the similarities of the x86 and x86-64 architectures. It was merely a matter of adding a x86 byte-code compile trampoline -shown in Figure 2 -which is slightly longer than its x86-64 counterpart, because the C calling convention [17] requires all parameters to be passed on the C stack.
pushl %eax pushl %ecx call caml_jit_compile movl %eax, %edx popl %ecx popl %eax jmpl *%edx Whenever the byte-code compile trampoline is invoked, %eax contains the current accu value, %ecx contains the byte-code address and the remainder of the OCaml virtual machine state is located in global memory locations and callee-save registers. Therefore %eax has to be preserved on the C stack and %ecx must be passed as parameter to the caml jit compile function. Upon return %eax contains the address of the generated native code, which is saved to %edx. Afterwards the stack frame is removed, restoring %eax, and execution continues at the native code address.
The remaining porting effort was mostly related to generalizing the code emitter preprocessor macros to conditionally emit x86 or x86-64 code, and fiddling with some nasty details of the different addressing modes. The core of the code generator is almost free of #ifdef's, because the structure of the generated code for the two targets is pretty much equivalent. This is especially true for floating-point operations, which use SSE2 registers and instructions on both x86 and x86-64. OCamlJit2 therefore requires a processor with support for the SSE2 extension [1,9]; if OCamlJit2 detects a processor without SSE2, it disables the JIT engine and falls back to the byte-code interpreter 3 .

Native code management
With our earlier prototype, memory allocated to a chunk of generated native machine code was never freed by OCamlJit2, even if the byte-code segment, for which this native machine code was generated, was released (using the Meta.static release bytecode OCaml function). This may cause trouble for MetaOCaml [19] and long-lived interactive top-level session, where the life-time of most byte-code segments is limited. This was mostly due to the fact that all generated native machine code was stored incrementally in one large memory region, with no way to figure out which part of the memory region contains code for a certain byte-code segment. We have addressed this problem with a solution quite similar to that used in OCamlJit [18]. Instead of generating native machine code incrementally into one large region of memory, we divide the region into smaller parts, called native code chunks, and generate code to these chunks. Every byte-code segment has an associated list of chunks, allocated from the large memory region, now called chunk pool, which contain the generated code for the given byte-code segment, as shown in Figure 3.
Every segment starts out with an empty list of chunks. As soon as on-demand compilation triggers for the segment, a new chunk is allocated from the chunk pool and native machine code is generated into this chunk until it is filled up; when this happens, the code generator allocates another chunk for this segment, and emits a jmp instruction from the previous chunk to the new chunk if necessary. Once a byte-code segment is freed using Meta.static release bytecode, all associated native code chunks are also released. This way OCamlJit2 no longer leaks native code generated for previously released byte-code segments.
While this technique is both simple and effective, there are also several drawbacks. The speed of both the JIT compiler and the generated code is somewhat dependent on the size of the native code chunks. A small chunk size helps to reduce the size of the working set and the amount of memory wasted in long-lived interactive top-level sessions, but on the other hand decreases the throughput of the JIT compilation driver and leads to somewhat less efficient execution for common byte-code programs which use only a single byte-code segment. We have settled on a chunk size of 256 KiB for now, which seems to provide a good compromise.
A possible way to reduce the amount of wasted memory with the interactive top-level would be to store code generated for small byte-code segments to some special, shared code chunks, and manage the life-time of these shared chunks using a reference counting scheme.

Floating-point optimizations
OCaml uses a boxed representation for floating-point values. It does this for various reasons, i.e. to simplify the interface to the garbage collector and the garbage collector itself. While this is an elegant and portable solution, it decreases the performance of OCaml programs using floating-point calculations, especially when executed with the byte-code interpreter.
The optimizing native code compiler ocamlopt applies various optimizations to avoid generating a boxed floating-point object in the heap for each and every floating-point value during the execution of the program. The byte-code interpreter ocamlrun however has to box every floating-point value, which is certainly slower than using the available floating-point registers, and also causes a non-negligible load on the garbage collector. Both OCamlJit and OCamlJit2 (as described in [15]) used the same strategy as the byte-code interpreter, namely allocating a heap object for every floating-point value during the execution, but both JIT engines applied various peephole optimizations to avoid the overhead of calling the floating-point C primitives.
We have implemented a new technique for OCamlJit2, which avoids the heap allocation for temporary floating-point values that appear as result of a byte-code instruction and are used as argument in the subsequent byte-code instruction. Figure 4 shows an example taken from the byte-code of the almabench.ml OCaml program. ccall caml_array_unsafe_get_float, 2 ccall caml_mul_float, 2 ccall caml_add_float, 2 ccall caml_add_float, 2 ccall caml_sqrt_float, 1 push Executing this piece of byte-code with the byte-code interpreter allocates five floatingpoint objects in the minor heap, one for each ccall. The first four objects will die immediately once the garbage collector is run, since they are only used as temporary results, while the result of the caml sqrt float call may indeed survive for a longer period of time. Furthermore, accessing the actual floating-point values of the temporary results requires at least four additional load instructions. If the temporary results would be kept in registers, there would be no need for the heap allocation and this would also eliminate the additional load instructions. While the heap allocations can seriously decrease performance, the additional load instructions are a minor issue, since the store-to-load forwarding techniques [3,14] implemented in modern processors will usually eliminate the load from memory.
We have implemented a clever optimization in OCamlJit2, which translates various floating-point primitives using SSE2 instructions and functions from the standard C math library, in a way that the result of each primitive is not stored in a heap-allocated object, but is left in the %xmm0 register. Subsequent floating-point primitives then take the value from the %xmm0 register instead of the memory location pointed to by %rax (or %eax). This process is repeated for all floating-point primitives in a row. The last floating-point instruction -the call to caml sqrt float in our example -then allocates a heap block for its result. Our optimization is particularly beneficial for the x86 port, where we were able to beat the optimizing native code compiler in the almabench.unsafe benchmark, but it also pays off for the x86-64 port, where we could reduce the execution time of floating-point benchmarks by up to 30% (compared to the earlier OCamlJit2 prototype).

Performance
With the x86 port in place we were finally able to compare the performance of OCamlJit [18] and OCamlJit2 running on the same machine. We measured the performance on three different systems, one x86 box for comparison with OCamlJit, and two x86-64 machines with different processors and operating systems to test-drive the recent improvements with our primary 64-bit targets: • A MacBook with an Intel Core 2 Duo "Penryn" 2.4 GHz CPU ( The benchmark programs used to measure the performance are the following test programs from the testsuite/tests folder of the OCaml 3.12.0 distribution: • almabench is a number-crunching benchmark designed for cross-language comparisons. almabench.unsafe is the same program compiled with the -unsafe compiler switch.
• bdd is an implementation of binary decision diagrams, and therefore a good test for the symbolic computation performance.
• boyer is a term manipulation benchmark.
• fft is an implementation of the Fast Fourier Transformation.
• nucleic is another floating-point benchmark.
• quicksort is an implementation of the well-known QuickSort algorithm on arrays and serves as a good test for loops.
• soli is a simple solitaire solver, well suited for testing the performance of non-trivial, short running programs.
• sorts is a test bench for various sorting algorithms.
For our tests we measured the total execution time of the benchmarks, given as combined system and user CPU time, in seconds.     values are better. The times were collected by executing each benchmark 5 times with every engine, and using the timings of the fastest run. Figure 5 highlights the speedup of the JIT engines relative to the byte-code interpreter on the three test systems. The x86-64 port received another nice speedup in the floating point benchmarks compared to the earlier prototype due to the improvements described in section 3.3. On x86 OCamlJit2 provides a performance boost of 2.2 to 9.6 compared to the byte-code interpreter, and is 1.1 to 2.0 times faster than OCamlJit. This is especially noticable with short running programs like soli.unsafe, where OCamlJit2 benefits from the reduced compilation overhead, and in floating-point programs, where OCamlJit wins because of clever floating point optimizations and its use of the SSE2 extension. Figure 6 shows the performance of the byte-code interpreter and the JIT engines relative to the optimizing native code compiler ocamlopt. One rather surprising result was the bad performance of the generated x86 native code for the almabench.unsafe benchmark, where OCamlJit2 was able to beat the native code compiler by a factor of 1.2. This may be related to the use of SSE2 instead of x87 instructions, which are generally faster, but it also seems that we have spotted a problem within the x86 port of ocamlopt here (unfortunately, we were unable to track down the issue). On a related note, we have also spotted some issues with the x86-64 floating-point code generated by ocamlopt and already submitted an appropriate patch 4 , which will be available in OCaml 3.12.1, yielding performance improvements of 6 − 13% with floating-point programs.
In general, floating-point benchmarks benefit a lot when used with OCamlJit2, which was somewhat expected, in particular with the optimizations implemented in the latest prototype. OCamlJit does a respectable job, but is unable to compete with OCamlJit2 performance-wise. This is probably caused by the better compilation scheme used by OCamlJit2 and also related to the use of Gnu Lightning within OCamlJit. Nevertheless, it is this use of Gnu Lightning which makes OCamlJit slightly more portable than OCamlJit2 (three supported platforms in case of OCamlJit, compared to only two with OCamlJit2).

Conclusions and further work
Our results show that Just-In-Time compilation of OCaml byte-code can lead to some significant speedup (at least twice as fast the byte-code interpreter in all benchmarks). Starting out with a simple compilation scheme and applying some clever optimizations led to impressive performance gains. But we have also noticed that our approach is somewhat limited, which is mostly related to the nature of the OCaml byte-code and the design choices made for the interpreter runtime. The OCaml byte-code and runtime were certainly designed with "fast interpretation" in mind [10] and perform quite well in this area, but this same fact also limits the possibilities for effective Just-In-Time compilation if one wants to avoid touching too many aspects of the runtime. Using a register machine as used by Dalvik [5] or Parrot [7] instead of the stack machine for the byte-code virtual machine would make it easier to apply JIT techniques -it would in fact make Just-In-Time compilation the natural implementation choice for the virtual machine.
Implementations of other runtimes, like the various JVMs [13] or the Common Language Runtime [8], show that it is indeed possible to perform efficient JIT compilation with stack virtual machines, but these runtimes use expensive compilation techniques for instruction selection, register allocation and instruction scheduling, whose applicability is questionable within the scope of the OCaml byte-code runtime. For example, in order to effectively reduce the overhead of closure application and return in OCaml byte-code execution, one would most likely need to perform interprocedural register allocation, which mandates the availability of global control and data flow information, both of which are difficult to collect efficiently. It may indeed be possible to design a Just-In-Time compiler for the OCaml byte-code, which generates code as fast as the code generated by ocamlopt, using standard compilation techniques [2], but such a JIT engine comes at a high cost with respect to maintainability and execution speed of the JIT compiler, and it is questionable whether it is worth this cost, especially since there is already an optimizing native code compiler, which limits the possible use cases for the Just-In-Time compiler.
The main application of OCamlJit2 is the interactive top-level and other dynamic code generation environments like MetaOCaml [19]. The OCaml repository already contains an experimental "native top-level" ocamlnat for this purpose, which uses the functionality of ocamlopt to generate efficient native code at runtime and execute it via the native code runtime. Improving ocamlnat may provide a better way to gain an efficient top-level, and we are currently evaluating what would need to be done in this area.