Maple Ong

Implementing `concatarray` in YJIT

blog

What's concatarray?

concatarray is a YARV instruction that pops two arrays off the stack, concatenates them and pushes the concatenated array back onto the stack.

Here's a quick example of how it is used - we want to get the contents of x by splatting it in an array.

x = [2, 3]
[1, *x]

The generated bytecode below uses concatarray to combine 1 and the contents of x in a new array.

== disasm: #<ISeq:<main>@-e:1 (1,0)-(1,19)> (catch: FALSE)
local table (size: 1, argc: 0 [opts: 0, rest: -1, post: 0, block: -1, kw: -1@-1, kwrest: -1])
[ 1] x@0
0000 duparray                               [2, 3]                    (   1)[Li]
0002 setlocal_WC_0                          x@0
0004 duparray                               [1]
0006 getlocal_WC_0                          x@0
0008 concatarray
0009 leave

Why Bother?

concatarray hasn't been implemented in YJIT. So when YJIT encounters this instruction, the side exit operation is called. This makes the program slow because we'd have to deoptimise and fall back onto the interpreter from YJIT.

It also prevents YJIT from running consequent machine code after concatarray, wasting optimized machine code that may result in a faster runtime.

Verifying The Problem

But first, let's verify that YJIT is indeed running the exit operation when concatarray is encountered.

Let's add the relevant bits of our code into a loop that runs 20 times, enough to trigger YJIT to start working.

20.times do
  x = [2, 3]
  [1, *x]
end

Next, run ./miniruby --yjit-stats test.rb.

YJIT stats log tells us that concatarray is the instruction where the 100% of the exit operations occur (in this contrived example).

***YJIT: Printing YJIT statistics on exit***
...
Top-20 most frequent exit ops (100.0% of exits):
            concatarray:         11 (100.0%)
                    nop:          0 (0.0%)
               getlocal:          0 (0.0%)
               setlocal:          0 (0.0%)
...

How?

We want to re-use the existing C-function for concatarray. This means that we're just calling the compiled machine code from the C implementation of this instruction. There is potential to implement it fully from YJIT but for now we just want to avoid the exit operation.

The Implementation

There are some setup that needed to be done for this change to happen (such as exporting static C functions) that I won’t talk about.

The interesting bit is the function we wrote that generates assembly code for concatarray. I’ve numbered each change and talked about it below.

// concat two arrays
fn gen_concatarray(
    jit: &mut JITState,
    ctx: &mut Context,
    cb: &mut CodeBlock,
    _ocb: &mut OutlinedCb,
) -> CodegenStatus {
    // 1.
    jit_prepare_routine_call(jit, ctx, cb, REG0);

    // 2.
    let ary2st_opnd = ctx.stack_pop(1);
    let ary1_opnd = ctx.stack_pop(1);

    // 3.
    mov(cb, C_ARG_REGS[0], ary1_opnd);
    mov(cb, C_ARG_REGS[1], ary2st_opnd);

    // 4.
    call_ptr(cb, REG1, rb_vm_concat_array as *const u8);

    // 5.
    let stack_ret = ctx.stack_push(Type::Array);

    // 6.
    mov(cb, stack_ret, RAX);

    KeepCompiling
}

  1. I don't really know exactly what this bit of code does, but it ensures our stack is set up for what we're about to do next.

  2. Here, we get the operands from the stack! Notice how the first pop is for the second array, and the subsequent pop is for the first array - that's because of the order of the values on the stack. Something to note here is that the variables here are not the popped values, but a pointer for those values.

  3. We take the pointers to our arrays and generate assembly instructions to move them into a register. The C_ARG_REGS array is an abstraction of the integer registers for the System V AMD64 ABI calling convention.

  4. Here, we "jump" to the pointer of rb_vm_concat_array and execute the code there.

  5. We push an Array type value onto the stack. We get a pointer to that spot in the stack in return.

  6. Finally, we generate assembly to move the value in the RAX register into the stack. RAX is a scratch register (temp register) YJIT uses to store values, so presumably, the resulting pointer to the concatenated array is already stored in RAX.

Did It Work?

Aside from writing a test for our changes (which we promptly did), we ran ./miniruby --yjit-stats test.rb again to make sure the exit operations are gone. Of course, the number of exit operations called is now 0. Phew!

...
total_exits:                    0

Whoa, Assembly Code!

We also checked out the assembly code that was generated using the #disasm method from the RubyVM::YJIT library.

def foo
  x = [2, 3]
  [1, *x]
end

20.times do
  foo
end

puts RubyVM::YJIT.disasm(method(:foo))

We can see the assembly code that YJIT generated for concatarray.

...
  # concatarray
  0x10e430256: movabs rax, 0x600002e184f8
  0x10e430260: mov qword ptr [r13], rax
  0x10e430264: lea rax, [rbx + 0x10]
  0x10e430268: mov rbx, rax
  0x10e43026b: mov qword ptr [r13 + 8], rbx
  0x10e43026f: mov rdi, qword ptr [rbx - 0x10]
  0x10e430273: mov rsi, qword ptr [rbx - 8]
  0x10e430277: call 0x106b3b080
  0x10e43027c: mov qword ptr [rbx - 0x10], rax
...

I went ahead and annotated the parts of the code based on our implementation above:

...
  # concatarray

  # Set up (`jit_prepare_routine_call`)
  0x10e430256: movabs rax, 0x600002e184f8
  0x10e430260: mov qword ptr [r13], rax
  0x10e430264: lea rax, [rbx + 0x10]
  0x10e430268: mov rbx, rax
  0x10e43026b: mov qword ptr [r13 + 8], rbx

  # Move array pointers into registers
  0x10e43026f: mov rdi, qword ptr [rbx - 0x10]
  0x10e430273: mov rsi, qword ptr [rbx - 8]

  # Call C function `rb_vm_concat_array`
  0x10e430277: call 0x106b3b080

  # Move return value back into stack
  0x10e43027c: mov qword ptr [rbx - 0x10], rax
...

Final Thoughts

In our change, YJIT just calls the interpreter’s implementation of rb_vm_concat_array. In the future, it's possible to implement concatarray fully in YJIT, giving us access to type check optimizations.

This was my first YJIT change so I am thrilled about it! I ended up having to modify the original implementation to use API for the new backend intermediate representation (IR) that the YJIT team was working on. The new backend will allow the team to easily support additional CPUs. The YJIT team did a great job of documenting their work so I was able to read comments and follow along despite not knowing Rust.

Big thank you to John Hawthorn for suggesting this task, pairing with me, and answering all my questions thoughtfully. I learnt so much!