Eileen M. Uchitelle

Exit(ing) Through the YJIT

When optimizing code for the YJIT compiler it can be difficult to figure out what code is exiting and why. While working on tracing exits in a Ruby codebase, I found myself wishing we had a tool to reveal the exact line that was causing exits to occur. We set to work on building that functionality into Ruby and now we are able to see every side-exit and why. In this talk we’ll learn about side-exits and how we built a tracer for them. We’ll explore the original implementation, how we rewrote it in Rust, and lastly why it’s so important to always ask "can I make what I built even better?"

RubyConf 2022

00:00:00 Hello everyone! I hope you're having a good time at RubyConf in Houston. Thank you so much for coming to watch my talk on tracing exit locations in YJIT.
00:00:03 Before we get started, I want to thank the organizers and volunteers for their hard work in putting on both RubyConf Mini and Houston this year. I know that a lot of work goes into organizing conferences, and I'm grateful for everything that you do.
00:00:21 If we haven't met before, I’m Eileen. You can find me anywhere online at @EileenCodes, and I guess on this new Mastodon thing at Ruby Social. I’ve had an account since 2018, and then all of a sudden everyone else had one. But I never tooted or tweeted or whatever they call it.
00:00:58 I’ve been writing Ruby since 2010 and have been a contributor to the Rails framework since 2014. In 2017, I joined the Rails core team, which is responsible for the future of the framework from the technical side. We decide what new features go into the next release, fix bugs, and ensure the framework is secure.
00:01:19 Earlier this year, in February, I joined Shopify's Ruby and Rails infrastructure team as a senior staff production engineer. There are currently about 45 people on the Ruby and Rails infra team organization, and many of us are here at the conference, so if you see us, say hi!
00:01:39 Our mission is to make Ruby a hundred-year language by ensuring that Ruby can support our applications and, by extension, our merchants and customers for years to come. We've built teams at Shopify to invest in the growth of the language, focusing on the things that we all need to continue thriving on Ruby.
00:01:54 Today, we're going to talk about one small portion of the investment work that we’ve been doing on YJIT. As part of our work to optimize YJIT, Aaron Patterson and I examined the most common places that YJIT was exiting in Rails.
00:02:07 Initially, our work wasn't producing much because we could see the exits occurring but had trouble finding and understanding the code that was causing the exits. I lamented that I wished we could get more information. I wanted to know the exact stack trace of the code that was exiting.
00:02:30 If we could get the backtrace, it would give us more information about whether that was an area where YJIT could be improved to better support Rails applications, or if our approach should be to refactor Rails to work better with YJIT.
00:02:49 Aaron and I set out to write a feature to trace the exits and pass that information to StackProf so that we could see the Ruby code causing frequent YJIT exits in Rails. If all of that sounds like a whole lot of nonsense, don't worry. This talk includes an introduction to how to get started, so that you'll be able to not only run YJIT in production, but I'll also give you the tools to find optimizations as well.
00:03:03 Today, we're going to go back to the beginning and explore what a JIT is, how YJIT compares to MJIT, how to configure Ruby with YJIT enabled, and what exiting the JIT means. Once we're familiar with these topics, we'll look at what Aaron and I built, how to find improvements in YJIT, and how to optimize our Ruby code.
00:03:31 My hope is that by the end of this talk, you'll feel more confident working with YJIT and will even be convinced to contribute.
00:03:50 So, what is a JIT? JIT stands for just-in-time compiler. Traditional compilers compile code in a step separate from runtime, like when you compile C or Rust code. The compiler will make optimizations while running, but once the source is compiled, generally, no optimizations will be made at runtime.
00:04:05 The Ruby code you write in your applications or services is interpreted and therefore not compiled. Compiled code runs faster than interpreted code because compiled code is easier for the machine to understand.
00:04:24 This is where a JIT comes in. A JIT is kind of like a hybrid between interpreted code and compiled code, allowing us to get the ease of use of interpreted languages with the speed of compiled ones. A JIT compiler compiles code just in time during the execution of a program.
00:04:50 In general, JITs are slower to warm up, but the advantage is that your code gets faster over time. For many years, Ruby didn’t have an official JIT; there were a few attempts, one example being Rubinius, but they were never merged into core.
00:05:06 Currently, Ruby has two JITs in the standard library: MJIT and YJIT. MJIT is a method-based JIT, which means it compiles single methods, one at a time. When MJIT is enabled, a separate thread runs in parallel and takes instructions from a priority queue and translates them into machine code. The next time a compiled method is called in your program, it'll use the machine code that was generated in the parallel compilation thread.
00:05:25 While MJIT started out with promising benchmarks for Ruby code, it unfortunately made Rails applications slower. This is a big problem because while not all Ruby code is Rails, all Rails applications are Ruby. If Ruby is hurting the performance of Rails applications, they'll be unwilling or unable to upgrade, which ultimately harms the community and ecosystem.
00:05:42 In 2020, the Ruby team at Shopify started working on a different JIT, which would eventually be named YJIT. YJIT stands for 'yet another JIT'; the name refers to how Ruby already had a JIT, and this is yet another one. I’m not sure why we didn’t go with yajit, because I think that’s kind of funnier.
00:06:00 YJIT was designed by Maxime, who works at Shopify. The design is based on her paper on basic block versioning. She had been working at Shopify on Truffle Ruby, but due to her interest in and expertise in compiler design, she proposed that we build a new JIT for CRuby using the technique described in this paper.
00:06:25 While MJIT compiles entire methods, basic block versioning is able to compile parts of methods and will only compile the code that is actually run in your program. For example, if you have an if-else conditional in your code and only the if branch ever gets run, YJIT will only compile that part and ignore the rest unless it's run. By compiling only code that is run, YJIT avoids compiling instructions that aren't needed.
00:06:48 YJIT can also compile parts of methods based on what instructions are defined in YJIT. In this example, both the able-to-compile and not able-to-compile methods will be run by your program, but YJIT will recognize which parts it can't compile for the able-to-compile method.
00:07:02 It will generate machine code when the call method is run. For those that aren’t able to compile, it will fall back to the interpreter. Other JIT designs, including MJIT, compile entire methods regardless of whether they’re run. It’s more efficient to generate machine code only for the instructions that you need.
00:07:23 In addition, YJIT produces more straight-line code, which results in faster execution because CPUs are more efficient when branching is avoided. The goal ultimately is to avoid exiting at all, but Ruby is a 25-year-old language that wasn't designed to be efficient for a JIT, so there's a lot of work that needs to be done to ensure that YJIT works well with CRuby.
00:07:47 If you're interested in learning more about the in-depth design of YJIT, you can watch Maxime's Rubik's talk from last year or any of her talks that she’s given this year at Ruby Kaigi or read the paper. It’s really interesting.
00:08:05 YJIT was originally developed as a separate project and was written in C. It was merged into CRuby as of the 3.1 release last December. That version only runs on x86, so I would kind of skip it and go straight to 3.2.
00:08:28 Earlier this year, we rewrote YJIT in Rust because it has strong type safety and better tools to manage complexity. It also enabled the team to write a new backend for YJIT so that it can be run on ARM and x86 machines.
00:08:54 Both MJIT and YJIT are available JIT options in the Ruby standard library. You might be wondering why there are both and which one you should use.
00:09:05 If you want a JIT for production, you should use YJIT. It's stable, performant, and used in production at Shopify in some places, and there are more maintainers focused on optimizing it.
00:09:22 However, that doesn't mean that MJIT is going away. If you watch KoCoolBun's talk from Ruby Kaigi this year, he goes in depth into the future of MJIT. Mainly it’s an interface for bringing your own JIT like Aaron Patterson's Tender JIT or John Hawthorne's HotJIT.
00:09:36 It’s also a great place to experiment with ideas for JIT optimization that can later be applied to YJIT. So think of it as kind of a playground for JITs where you get to do more experimental stuff.
00:09:50 For the remainder of this talk, we're only going to focus on YJIT.
00:10:00 Now that we've dug into what a JIT is, the difference between MJIT and YJIT, I bet you're thinking, 'I’m so intrigued, I want to run it in production! How do I get started?'
00:10:16 If you want to use YJIT in production and you aren't ready to contribute or you don't want to build YJIT from source, I recently updated Ruby's nightly Docker images to be built with YJIT available.
00:10:38 Previously, your only option was to build from source or wait for a release. In most cases, though, if you are watching one of my talks, you’re probably not as interested in how to simply run YJIT in production.
00:10:52 My hope is that you also want to learn how to build it locally so that you could optimize and debug YJIT, and of course contribute back to CRuby.
00:11:09 To build Ruby with YJIT, you need the following prerequisites: you need Autoconf, Make, Bison, GCC or Clang, libyaml, OpenSSL 1.1, Rust, and Ruby. Yes, you need Ruby to build Ruby.
00:11:31 It can be any older version down to 1.8, but it has to be available for you to build from source. For OpenSSL, it’s really important you use 1.1, as OpenSSL 3 does not yet work with Ruby.
00:12:00 There's an open issue that's currently being worked on to fix that, but I'm not sure what the timeline for actually fixing it is, and you will most definitely have a bad day if you try to compile with OpenSSL 3. I've experienced it.
00:12:15 After you've cloned Ruby and gotten everything installed, you need to run autogen.sh, which is a common thing to do in C libraries. This will generate the required makefiles and other configuration for building CRuby.
00:12:35 You don't really need to know what it does; you just have to run it. Then we need to run our configure command; there are a lot of flags that you can pass to your Ruby build depending on whether you want debug mode or specific experimental features.
00:12:51 I’m not going to go over any of that today; we're only going to look at the ones that you need to build YJIT.
00:13:05 You used to have to enable YJIT explicitly, and up until yesterday, I had that in these slides until I saw that that was not true. We recently updated it so that if you have Rust installed, then YJIT will be enabled in your build.
00:13:19 It's not turned on by default, but it's available. We then need to set the prefix for where we want Ruby installed. The path should match what your version manager wants.
00:13:36 This example is for CRuby or Truby. I don’t know how people say it; I kind of like Truby. You can set the name to whatever you want as long as it doesn’t conflict with an existing Ruby version.
00:13:59 This is what you’ll call Ruby YJIT in your shell when you want to switch Ruby versions. Next, you need to set the path to OpenSSL 1.1, otherwise, Ruby will try to build with version 3, and it will compile.
00:14:20 And it’s also recommended to disable RDoc, as this will speed up the time it takes to run make. Once we have everything configured, we need to run make install, which will install Ruby into the directory that we set in the configure prefix.
00:14:38 I recommend taking these three things, autogen, configure, and make install, and storing them in a script so you don’t have to remember how to configure Ruby. Of course, once it’s configured, you can simply run make and make install.
00:15:01 But you'll need to run this the first time or anytime that you're switching how you want Ruby configured. While we built Ruby with YJIT enabled, it’s not turned on by default.
00:15:22 So you'll need to explicitly enable it. You can do that by passing the RUBY_YJIT_ENABLE environment variable set to one. You can also use command line arguments for JIT or YJIT; these all do the same thing—they enable YJIT.
00:15:39 We have multiple options because environment variables are better for applications or some scripts, and the JIT option predated YJIT because it used to enable MJIT in the past.
00:15:54 If you're running YJIT in production, you'll want to compile in release mode, which is the default. However, if you are working on YJIT, optimizing Ruby code, or want statistics about what Ruby code YJIT is compiling, then you're going to need to configure in Dev mode.
00:16:10 To do that, we need to add YJIT_ENABLE=dev to the configure command and reconfigure Ruby. The Dev flag will enable YJIT stats collection so that we can get more information about what the compiler is doing when you compile with stats enabled.
00:16:27 YJIT is going to run slower, so don't use this in production; this is only for debugging. Now that our YJIT version is configured to run in Dev mode, we can get a lot of information about what Ruby code YJIT is and is not able to compile.
00:16:46 This means we can use Dev mode to figure out when YJIT is exiting to the interpreter. An exit means that YJIT was not able to compile your code. This could be because your code is too complex for the JIT currently, because it’s invalid Ruby, or because it’s using an instruction that’s not yet defined in YJIT.
00:17:02 Because YJIT uses basic block versioning, it can exit or fall back to the interpreter when it can’t understand the code that we want it to compile. Dev mode will give us the option to collect statistics about our code, which will help us determine where exits are occurring and whether we should be optimizing YJIT or fixing our Ruby code.
00:17:27 Let’s take a look at some code and use YJIT stats to figure out where we have widget exits. This is an example of a simple Active Record script we often use in Rails to reproduce bug reports without having to create an entire application.
00:17:45 In the script, we establish a connection, create a post table, add a post model, and run a quick test for create. It’s pretty simple, but this code example represents the lifecycle of a simple Active Record query.
00:18:05 I made this instead of running an application test because it’s more pared down, and we’re going to see fewer dependencies in our exits, which will make it easier to understand what’s happening.
00:18:24 To run this script, we can add the YJIT stats argument, which will collect information about exits and compilation. We don't need to call YJIT here because the YJIT stats argument will automatically enable YJIT for us.
00:18:39 I recorded a demo because it's easier to see what's happening. Let’s run this script. When we run our script, we get a lot of information; we'll go over what each section means.
00:18:55 At the top, we have the reasons that the method calls exited. Usually, method exit reasons have to do with either code complexity or invalid Ruby. An example of complex code in YJIT that won’t compile is the keyword splat argument.
00:19:14 We also see an 'I:SE' error which is also causing an exit. Theoretically, we could find the code that's causing this I:SE error. We should fix that because we don’t actually want our I:SE errors; Ruby can still compile it fine, but YJIT can’t.
00:19:33 If we scroll down, we can see the exit reasons for some specific instructions. Here we can see that 'invoke super' exited because the method entry changed or because we have a block, while other instructions like 'expand array' show that all exits are zero.
00:19:55 This means that either we had no code that used the 'expand array' instruction, or that YJIT was able to compile 100 percent of our 'expand array' calls. The next section shows various stats that might be useful depending on the optimization work that we’re doing.
00:20:18 We can see the total side exit count and total exit count along with a lot of other information. Side exits differ in that side exit is exiting from an implemented instruction, whereas total exits count the exits that are not implemented in YJIT.
00:20:42 Lastly, at the bottom here, we have our top 20 most frequent exits which can be useful for figuring out where to put our efforts. Looking at this, we can see that 'op send without block' accounts for 43 percent of our exits, while 'send' accounts for 18 percent, and 'check match' for 0.6 percent.
00:21:06 These stats give us a lot of information about the code that YJIT can compile. However, there's some information that's missing from these stats: we know that YJIT is exiting and in some cases we have information about why, but we don’t actually know what the Ruby code that's exiting looks like.
00:21:27 There's nothing pointing to the source from these stats. We can't tell if the exit is in a gem, in Rails, or in our application code.
00:21:43 Earlier this year, I was working with Aaron on reducing YJIT exits in Rails. I complained about this, that the stats were useful for knowing what instruction exited, but not what the Ruby code looked like. With Rails, we have thousands of lines of code and no information on what file the code is in or what it looks like.
00:22:05 This means we didn’t have enough information to make routes faster with YJIT. The larger the application, the more important it is for YJIT to tell us the exact location of the code that's causing the exit; otherwise, we're spending a lot of time guessing what optimizations can be made.
00:22:22 If the instruction is implemented, like the 'invoke super' one, but it's exiting often, we want to find that code so we can decide how to handle it. In some cases, we might want to rewrite Rails to work better to avoid the complexity that makes YJIT exit.
00:22:47 For example, we could refactor our code to avoid I:SE errors so that we no longer exit there. In other cases, we might want to prioritize optimizing YJIT and work on improving 'op send without block' so that YJIT can compile in more scenarios.
00:23:08 While useful, the exit stats only give us some of the information that we need to decide whether to improve Rails or improve YJIT.
00:23:21 As always when I complain to Aaron about something, I wish we had, he says, 'Oh yeah, we can totally make that work.' So we paused work on finding exits in Rails to build a new tool for anyone working on optimizing YJIT.
00:23:41 We set out to build a tool to complement YJIT stats called 'YJIT Trace Exits' that would allow us to keep track of the backtrace of any exit from YJIT and use that to find the code that was causing the exits.
00:23:58 This feature will enable us to record the backtrace of every exit and marshal dump it to a file. We can then use StackProf to view the code backtrace for specific instructions. Instead of performance profile data, we see the percentage of exits from YJIT.
00:24:18 StackProf comes with a handy tool that I had actually never used before called walk, which will allow us to traverse the backtrace interactively. We’ll look at a demo in a minute.
00:24:38 To use the new functionality, we can pass 'YJIT Trace Exits' to our script. I’ve removed 'YJIT stats' because 'Trace Exits' will automatically turn on stats collection for us.
00:24:57 Let’s look at another demo. When we run our script, we see the same output as before except for now there is a new line that tells us where to find our dumped locations.
00:25:07 First, let's look at the file using StackProf's text argument. We can see the same exits in the top 20 frequent exits are present in our StackProf output. We have 'op send without block' with 43 percent.
00:25:27 Instead of performance data about time spent, we see the percentage of times exited, which, on its own, isn’t super useful; it’s still the same data we had from YJIT stats.
00:25:45 However, StackProf has a feature for interactively traversing this data that we’ve dumped. We can use the walk argument with the file name and the method argument with an instruction. Here I’ve chosen 'check match'.
00:26:01 If we choose one, we can view the source code of the exit. Now we know exactly what the code looks like, we know exactly what file it’s in, and we know that we know what 'check match' means.
00:26:18 I happen to know that 'check match' is not implemented in YJIT yet, so the optimization that we would do here is to implement the 'check match' instruction in YJIT.
00:26:36 Knowing what this code looks like will help us write a good test for the YJIT for the 'check match' instruction. It’s a lot easier to test that YJIT isn’t exiting if we know exactly what the Ruby code that uses the 'check match' instruction looks like.
00:26:53 Now let’s look at an instruction that we know is implemented, like 'invoke super'. The 'invoke super' is exiting for a few different reasons, and if we look at the code, we may be able to figure out which ones are exiting and why.
00:27:12 From the stats data, we know that 'invoke super' is either executing because we have a block or because the method entry changed. As you can see in our data here, we have more than just Rails, so we can improve the gems that are exiting frequently as well.
00:27:29 From this code, we can see that there’s a block around super, so this matches up with one of the exit reasons that we saw earlier. We can use this information to improve YJIT to avoid exiting, or we could refactor Concurrent Ruby to not need to use the 'invoke super' instruction with a block.
00:27:52 That’s probably less likely to be possible, but we have that information now for why it’s exiting. I want to show you an example of how we can use this exit information to actually refactor our Ruby code.
00:28:11 Rather than finding complex code in Rails, I’ve decided to do a demo that purposely implements code that I know will exit, and then we can all fix it together.
00:28:31 In this script, I've added a new method named 'call method' that takes one argument, 'options.' The options are then assigned to one, two, and three. Looking at the method, we’d expect the past arguments would be an array, but we've accidentally passed a single string 'hello'.
00:28:52 This is going to cause the 'expand array' instruction to exit because the argument passed is not an array. For this example, I’m going to add a new argument to our script that sets the call threshold to one.
00:29:10 By default, YJIT will only compile methods that it sees 10 times or more. Because our script is going to only call 'call method' once, we need to set the threshold to one; otherwise, we won’t see the exits from our script.
00:29:26 Let’s take a look at how this works. Here’s our code again. When we run it, we're going to see a new exit for 'expand array' at the bottom of our top 20 most frequent exits.
00:29:46 If we then go up to the exit reasons, we’ll see that this code exited because it’s not an array. By using 'Trace Exits' here, we can find the code that is incorrect and refactor it to avoid this exit.
00:30:02 This is a case where YJIT is not only helping make Ruby faster, but it’s also acting as a type checker in a way, notifying us when we have invalid Ruby. This Ruby code can run, but it doesn’t make it correct.
00:30:17 Since the Ruby code isn’t valid, YJIT can’t compile it. By using 'YJIT Trace Exits', we can easily find the offending code, avoiding an exit and improving our Ruby code at the same time.
00:30:45 Using Trace exit locations in StackProf, we could pinpoint the exact code that exited and even identify the caller if we continue up the stack. So here we can see our Hello that's incorrect; I mean—that's where our incorrect code is.
00:31:02 Let’s go back to our code and fix it to avoid exiting. Instead of passing 'hello', we can pass an array of 'hello' and 'RubyConf', and when we run our script again, we'll see that there are now no exits for the 'expand array' instruction.
00:31:21 Now, we just run this code, so we know we were expecting an exit; that’s why I wrote it that way. But I hope that you can see how that would be useful in an application where you could say, 'Oh, I know I did something wrong, and I can fix it now.' YJIT won't exit.
00:31:39 The exit locations tool gives us a lot more information about what YJIT can and can't compile so that we're able to make informed decisions about how to optimize our code for YJIT.
00:32:00 It’s really important to make sure that we have as much data as possible when working with YJIT so we can make Ruby faster by focusing on what will have the most impact rather than micro-optimizations.
00:32:20 With the information exit location gives us, we can make two choices: we can either optimize YJIT or we can optimize code for YJIT. Cases where we can focus on optimizing YJIT involve looking at whether the instructions that are exiting are implemented in YJIT.
00:32:43 If they aren't, we send a pull request to fix it. You can find whether an instruction is implemented in YJIT by looking at the codegen.rs file. Here we can see that 'invoke super' is implemented.
00:33:06 If we search this file for 'check match'—as of this talk—we won’t find that instruction implemented in this file, and that’s how we know YJIT doesn’t have that instruction.
00:33:24 We know that while 'invoke super' is implemented, it exits when the method entry changes or if we have a block. We can find where those exits are being recorded by looking for the counted exit and exit reason that we saw here.
00:33:40 We can see 'invoke super' blocked, so that means an 'invoke super' exited. We counted that exit when we saw a block.
00:34:02 Knowing that these are frequent exits in Rails means that we should prioritize implementing the complex code paths to fix the block or method entry change, because doing so will directly benefit Rails applications.
00:34:19 In most cases, we don’t want to have to rewrite our applications or framework code in order to get the benefits of a JIT. Ideally, we shouldn't be rewriting Rails, but that’s not always the case.
00:34:41 So in these cases, there are times where improving the JIT is too difficult. We can use YJIT exit locations to refactor our Ruby code to avoid writing code that’s hard for YJIT to compile, like our 'expand array' not-an-array exit.
00:35:04 We can see firsthand that sometimes in our script, we exit because we have invalid Ruby code with I:SE errors or not arrays. In other cases, we’re using functionality that we can’t get, like splatting keyword arguments.
00:35:24 Having the information about what code is exiting and why will help us refactor Rails to run faster on YJIT. We can look through the exits, use trace locations to understand the code, and update the code to go through a path that’s not going to exit.
00:35:47 In these cases, performance changes should focus on what real applications are exiting on and what impact refactoring will have. Additionally, we should focus on the exits that happen most frequently.
00:36:02 We want to avoid making micro optimizations that won’t have a real-world impact. But if a refactor can save a significant performance hit, it might be worth the trade-off of a rewrite.
00:36:17 Now that you know what a JIT is, how to build it, how to build YJIT, and how to investigate optimizations, I hope that you’ll consider running your application with YJIT enabled.
00:36:35 We need more companies willing to try out YJIT so we can continue to improve it and optimize it. If you have any issues or frequent exits, you now have the tools to investigate and report what you’ve found.
00:36:49 We build the best software when we have the most information about where it can be improved. I hope that in addition to running YJIT, you’ll also help us make YJIT better through contributions upstream.
00:37:06 I’ve talked about this a lot before in my prior talks, and it won’t ever not be true: when we fix things in Ruby upstream, we’re fixing it for everyone.
00:37:26 While we have a great team working on YJIT at Shopify, we still need your help to improve Ruby, optimize YJIT, implement frequently used instructions, and reduce exits.
00:37:49 Ideally, YJIT would never exit and every instruction would be implemented, but we can’t do that on our own. We need your help.
00:38:06 It’s so important for us to treat Ruby as an extension of our application. There is so much work to do, and there’s plenty of room for your ideas and your work.
00:38:27 I’m very excited for Ruby’s future. Our team at Shopify is brilliant, but so are all of you in this room!
00:38:41 All of us can run YJIT and contribute to making Ruby faster than anyone thought possible. Thank you so much for listening; I hope you enjoyed this talk and you're going to try out YJIT as soon as possible.
00:38:54 Again, you can find me online at @EileenCodes. Thank you.