Ruby Compiler Fun: AOT and JIT Compilation

Who knew writing a compiler could be so much fun.

I managed to accomplish two things tonight. It's late and I have a flight home tomorrow, so I'll be brief.

jrubyc: JRuby's Ahead-Of-Time (AOT) Compiler

I have whipped together the very barest of command-line, ahead-of-time compilers, along with a simple script to invoke it.

~/NetBeansProjects/jruby $ jrubyc
Usage: jrubyc <filename> [<dest>]

It's mostly just a very thin wrapper around the existing compiler code, so it can only compile constructs it knows about. However, for really simple scripts without any unrecognized nodes, it works fine:

~/NetBeansProjects/jruby $ cat samples/fib.rb
# calculate Fibonacci(20)
# for benchmark
def fib(n)
if n<2
  n
else
  fib(n-2)+fib(n-1)
end
end
print(fib(20), "\n")
~/NetBeansProjects/jruby $ jrubyc samples/fib.rb tmp
~/NetBeansProjects/jruby $ ls tmp/samples
fib$MultiStub0.class    fib.class

At the moment, two classes are generated; one is a class to hold the script entry points and the other is a stub class for all the actual blocks of code contained within the script (toplevel code, method code, etc). This will soon be a single class file, so pay the MultiStub no mind.

We can then execute the script like you'd expect, specifying the JRuby and ASM jar files on the classpath:

~/NetBeansProjects/jruby $ export CLASSPATH=lib/jruby.jar:lib/asm-2.2.2.jar:tmp        
~/NetBeansProjects/jruby $ java samples/fib
6765

Huzzah! Compilation!

Now of course, as I mentioned, this only compiles scripts containing constructs it knows about. If you try to compile a script it can't handle, you'll get an error:

~/NetBeansProjects/jruby $ jrubyc lib/ruby/1.8/singleton.rb
Error -- Not compileable: Can't compile node: ModuleNode[]

The compiler currently supports only literal fixnums, strings, and arrays, simple method definitions, while loops, if/else, and calls that don't involve blocks or splatted arguments. More will come as time progresses. The benefit of building the compiler piecemeal like this becomes more apparent in the next section...

JIT Compilation

The current compiler only understands enough of Ruby to handle my experimentation and research. The compiler also does not output one-to-one Ruby-to-Java classes or even a single large method: it outputs a class containing a method for every semantically separate block of code in a given script. In Ruby's case, that means toplevel code, code found within the body of a class, and code found within the body of a method definition. By combining these two traits, we have everything necessary for a simple JIT.

A JIT, or Just-In-Time compiler, performs its compilation at runtime, usually based on some gathered information about the executing code. HotSpot, for example, has an extensive array of optimizations it can perform on running code just by watching how it executes and eliminating unnecessary overhead. My vastly simpler JIT uses a much more basic metric: the number of times a method has been invoked.

The actual compiler code is the same as that used for the AOT compiler, with one major difference. Instead of the generated code being dumped to a file for later execution, it's immediately loaded, instantiated, and snuggled away in the same location where interpreted code used to live. The logic goes like this:

1. A method is called. We'll name it "foo"
2. foo's code is written in Ruby, so it's just a sequence of AST nodes to be interpreted
3. we interpret foo's nodes, but each time we increment a counter. When the counter reaches some number (currently 50), the compiler kicks in
4. if the code can't be compiled, we continue to interpretation, but we set a flag and never try to compile again
5. if the code can be compiled, we save the generated code and use it for all future invocations

Because the compiler can generate these small pieces of code, we're able to JIT Ruby code that was not compiled before execution began, gaining the benefits of a compiled platform without losing the flexibility of an agile script-based development model. It also means we can start benefiting from bytecode compilation even before the compiler is complete.

So how well does it perform? Very well, provided you don't go outside the narrow range of AST nodes the script supports:

~/NetBeansProjects/jruby $ cat test/bench/bench_fib_recursive.rb
require 'benchmark'

def fib_ruby(n)
if n < 2
  n
else
  fib_ruby(n - 2) + fib_ruby(n - 1)
end
end

puts Benchmark.measure { fib_ruby(30) }
puts Benchmark.measure { fib_ruby(30) }

Here we have a fib benchmark script with a few nodes the compiler can't handle. For example, the blocks at the bottom of the script won't compile correctly at present. So it's a good candidate for the JIT.

Once the JRuby JIT's been wired up, we can simply run the code as normal:

~/NetBeansProjects/jruby $ jruby test/bench/bench_fib_recursive.rb
compiled: Object.fib_ruby
2.877000   0.000000   2.877000 (  2.876000)
2.955000   0.000000   2.955000 (  2.955000)

You will notice the "compiled" logging output I currently have in the JIT. The only method hit hard enough to be compiled during this run was the fib_ruby method defined on the toplevel Object instance. Now this performance is drastically increased over the current trunk, largely due to compilation but also due to a faster dynamic method invocation algorithm we're experimenting with. And there's still a lot of optimization left to be done at both the compiler and runtime levels. But it's already a vast improvement over JRuby from even a month ago. Things are moving very quickly now.

We also look better running under the Java 6 server VM. The "server" VM performs more aggressive optimizations of Java code than does the default "client" VM. Generally this is because the optimizations involved cause the server VM to start up a bit more slowly, since it waits longer and gathers more information before JITing. However in this case, the results are very impressive when we compare the JRuby JIT running under the Java 6 server VM against Ruby 1.8.5:

~/NetBeansProjects/jruby $ jruby SERVER test/bench/bench_fib_recursive.rb
compiled: Object.fib_ruby
1.645000   0.000000   1.645000 (  1.645000)
1.452000   0.000000   1.452000 (  1.453000)
~/NetBeansProjects/jruby $ ruby test/bench/bench_fib_recursive.rb
1.670000   0.000000   1.670000 (  1.677901)
1.660000   0.000000   1.660000 (  1.671957)

The future's looking pretty bright.

None of this code is in trunk at the moment, but it should land fairly soon. The AOT compiler may come before the JIT, since it's minimally invasive and won't affect normal interpreted mode execution. Look for both to be available in JRuby proper within a week or two, and watch for the compiler itself move toward completion over the coming weeks.