Ruby JIT Hacking Guide

Ruby

Takashi Kokubun

Ruby JIT Hacking Guide

by Takashi Kokubun

In this talk titled "Ruby JIT Hacking Guide," Takashi Kokubun, a member of the Ruby Infrastructure Team at Shopify, discusses Just-In-Time (JIT) compilation in Ruby. The presentation aims to educate attendees about writing JIT compilers and introduces a new tutorial repository called Ruby Challenge. Key points covered include:

The evolution of JIT compilation in Ruby, starting with the introduction of MJIT in Ruby 2.6 and the more recent Arjit compiler in Ruby 3.3, which is designed for production workloads and is written in Ruby rather than C.
The performance improvements associated with Ruby 3.3, which include a 70-80% performance enhancement over Ruby 3.2.
How to use Arjit for experimental purposes, including building Ruby in a Rust-enabled environment and enabling JIT with the command line flag or an environment variable.
An overview of Ruby's internal workings, focusing on how the stack-based virtual machine compiles Ruby code to machine code, illustrated with examples demonstrating how to read assembly code generated by Ruby.
The process of writing a custom JIT compiler, highlighting the importance of understanding Intel's Software Developers Manual for instruction syntax.
Practical advice and optimization strategies for JIT compilers, including managing method definitions, handling local variables more efficiently, and implementing polymorphic methods for improved performance.

Kokubun encourages attendees to participate in writing their own JIT compilers, emphasizing collaborative community efforts to enhance Ruby's performance. The overall message is one of experimentation and contribution, inviting developers to explore the capabilities of JIT in Ruby, thus paving the way for future improvements in Ruby programming.

The key takeaway is the potential for creating faster Ruby environments through tailored JIT compilers, enhancing performance without competing against existing implementations. The session ends with a call to the community to engage in discussions about optimizing Ruby, ultimately benefiting the broader Ruby ecosystem.

00:00:08.120 Thank you.

00:00:11.300 Hello everyone. I'm going to talk about Just-In-Time compilation, which is just JIT for short.

00:00:15.420 Maybe some of you are feeling a bit sleepy now since it's day three and everyone is kind of tired, but I have a solution for you. It's this repository called Ruby Challenge, which is a tutorial for writing a Ruby JIT. The goal of this talk is to make you a Ruby JIT author, and I hope that writing JIT is a fun experience to keep you awake. So, let's please give it a try!

00:00:27.060 So who am I? I'm Takashi Kokubun, and I'm working on the Ruby Infrastructure Team at Shopify. I maintain Ruby and projects like Audit and template engines like ERB. All of these non-Watch projects are supported by GitHub Sponsors, so I want to thank you for your support. Over the last year, I've been quite active, making 999 commits, while another so-called ‘patch monster’ has made 996. So, you could say I’m a JIT monster now.

00:01:32.580 Subsequently, at Shopify and Maxine, we created another compatible JIT compiler, which was the second JIT compiler introduced in Ruby 3.1. This year, the day before yesterday, we released yet another JIT compiler called Arjit. Unlike MJIT, which was originally written in C and then rewritten in Rust, Arjit is purely written in Ruby and was released as a preview yesterday. Arjit replaces MJIT because MJIT catered to limited workloads, specifically for an NS emulator, which isn't a production workload. Instead, Arjit optimizes Rails benchmarks and web application workloads, as we aimed for a compiler that could be genuinely beneficial for production.

00:02:54.180 MJIT also relies on a C compiler, making the implementation more complicated. Thus, Arjit was designed to simplify the virtual machine. The first goal of Arjit was to remove MJIT to streamline the virtual machine, reusing the existing infrastructure to avoid complicating the implementation. Another goal was to introduce the ability to create custom JIT compilers, allowing developers to have fun experimenting with their own ideas that might not align with the design of MJIT.

00:05:03.900 The reason we're focusing on JIT is because compilers generally offer greater opportunities for optimizing performance compared to interpreters. If you listened to yesterday's keynote, we announced deploying the budget for a server to render the Shopify storefront application in production. On Ruby 3.2, we achieved a 10% performance increase on average. Now, for the remaining 1.1% of our cluster, we are rolling out Ruby 3.3, which delivers a performance enhancement of 70 to 80%. Thus, Ruby 3.2 can often be considered relatively slow. I encourage you to upgrade to Ruby 3.3 for improved performance.

00:05:40.740 How many of you are already using JIT? Please raise your hand if you are using JIT in production. Thank you! Now, let's discuss how to use Arjit. It’s primarily for experimental purposes; I don't expect anyone to deploy Arjit to production just yet. To utilize Arjit, first build Ruby in an environment that has Rust installed, as it's written in Rust. Once you install Ruby, it will automatically enable JIT for you.

00:06:08.340 To enable it, simply pass the -J flag to the Ruby command or set an environment variable called RUBY_JIT_ENABLE. This brings us to the main part of this talk: Ruby JIT hacking. The idea is to introduce a guide that covers the Ruby virtual machine and internals with a focus on JIT compilation. The previous Ruby hacking guide does not cover JIT, as it is quite dated. Today, we're introducing the Ruby JIT hacking guide, which explains the Ruby JIT internals.

00:06:54.180 Let’s discuss how Ruby works internally. When you have methods that simply calculate and return a value, the Ruby interpreter compiles the Ruby code into a parse tree. Then, the interpreter translates that tree into a sequence of operations specific to the Ruby virtual machine. As mentioned in yesterday's talk, Ruby is a stack-based machine, where instructions push objects onto the stack. If you execute the first two instructions of a method, you’ll have calculations on the stack, and when you reach the addition instruction, Ruby will compute the result and push it back onto the stack before returning it.

00:08:30.720 In this talk, I’ll help you understand how to read assembly. This is an actual generated code example showing how you can observe your Ruby code in the machine code format using the -J --dump option. This command can help you understand the instructions Ruby runs by commenting. Even if you aren't familiar with assembly code, you can read the comments to understand what is happening in the code.

00:09:01.560 For instance, let’s look at the assembly generated for adding two numbers. The first line puts the object '2' onto the stack and the second line plus instruction calculates the result. After computing, the method will return `3` as the integer representation of the addition. Additionally, there are instructions for overflow checks to prevent issues with large numbers.

00:10:35.520 Now, I want to take you through the process of writing a JIT compiler from scratch. Resources like the official Intel Software Developers Manual provide comprehensive documentation on all the instructions that you will find handy while working on a JIT. However, often it is easier to reference web resources that summarize these instructions rather than browsing through long PDFs. For example, the first instruction for moving a value into a register involves several bytes that need to be encoded into the binary structures of code.

00:11:34.560 Having understood how to translate Ruby bytecode into machine encoding, you can write Ruby code that translates simple operations into those instructions. Therefore, one would create an array in Ruby that represents the bytecode, allowing you to push integers as necessary. Once you have your byte array, the next step is to convert that into an executable format stored in the memory using standard libraries, thus creating a functional JIT compiler.

00:12:55.440 In conclusion, creating a custom JIT is not only feasible with Ruby version 3.2 but can also be very rewarding after incorporating Arjit. The language includes features that allow the creation of multiple methods for each individual method without worrying about compatibility. As demonstrated, people such as Aaron Patterson and John Horton have already begun writing custom JIT compilers for Ruby.

00:13:54.680 Before wrapping it up, I'd like you to try using Arjit because, through practical development, you can develop methods to compute intense Fibonacci sequences. On this graph shown, if the higher the value is, it indicates better performance. You will notice that Ruby 3.3 with no JIT is longer than Arjit, and the idea is that if you follow through with the tutorial, you can create a JIT compiler faster than what has been previously established.

00:15:57.180 Of course, the most critical takeaway from this experience is learning how to improve performance through JIT development, and as I explained, I want more people to participate in these discussions and ideas to understand how they might be able to enhance Ruby's performance in the future.

00:17:12.900 Now, let me quickly summarize the transitions we can explore while writing a JIT compiler. First up is site exit, which is crucial, as static exits can help optimize various workflows by reducing overhead and improving efficiency during runtime. You want to make sure that your JIT understands how integer calculations typically do not overflow, allowing it to optimize easier procedures without invoking unnecessary operations.

00:18:11.940 Next is message redefinition, which is crucial in Ruby's dynamic environment. When redefining methods, one can reduce the redundancy in jumps to the original definition since redundancy can bog down performance. With a hook in place to handle this for method definition invalidation in Arjit, one can ensure that old code paths don’t slow down new executions of redefined methods.

00:19:45.420 In addition, managing local variables offers an opportunity for improved performance as well. By rethinking how local variables are handled in an optimized approach, we can shift from a standard memory-centric model to a more register-oriented approach, significantly improving performance through faster operations.

00:20:59.520 Finally, I want to emphasize the idea of polymorphic methods and inline caching optimization techniques. This method allows various array operations to be handled without incurring performance penalties from repeated calls. By using these advanced optimizations, we can conclusively create a faster Ruby environment, demonstrating the capabilities and benefits of writing a tailored JIT compiler, paving the way for more exciting developments in Ruby.

00:23:25.180 To wrap up, I aim to inspire and encourage everyone to explore creating their own JIT compiler through the Ruby community. Remember that this isn't to compete against existing implementations but to contribute back the insights derived through experimentation, potentially making Ruby more efficient overall. Thank you for your attention!

RubyKaigi 2023