Rust YJIT

In December 2021, we opened an [issue] to solicit feedback regarding the
porting of the YJIT codebase from C99 to Rust. There were some
reservations, but this project was given the go ahead by Ruby core
developers and Matz. Since then, we have successfully completed the port
of YJIT to Rust.

The new Rust version of YJIT has reached parity with the C version, in
that it passes all the CRuby tests, is able to run all of the YJIT
benchmarks, and performs similarly to the C version (because it works
the same way and largely generates the same machine code). We've even
incorporated some design improvements, such as a more fine-grained
constant invalidation mechanism which we expect will make a big
difference in Ruby on Rails applications.

Because we want to be careful, YJIT is guarded behind a configure
option:

```shell
./configure --enable-yjit # Build YJIT in release mode
./configure --enable-yjit=dev # Build YJIT in dev/debug mode
```

By default, YJIT does not get compiled and cargo/rustc is not required.
If YJIT is built in dev mode, then `cargo` is used to fetch development
dependencies, but when building in release, `cargo` is not required,
only `rustc`. At the moment YJIT requires Rust 1.60.0 or newer.

The YJIT command-line options remain mostly unchanged, and more details
about the build process are documented in `doc/yjit/yjit.md`.

The CI tests have been updated and do not take any more resources than
before.

The development history of the Rust port is available at the following
commit for interested parties:
https://github.com/Shopify/ruby/commit/1fd9573d8b4b65219f1c2407f30a0a60e537f8be

Our hope is that Rust YJIT will be compiled and included as a part of
system packages and compiled binaries of the Ruby 3.2 release. We do not
anticipate any major problems as Rust is well supported on every
platform which YJIT supports, but to make sure that this process works
smoothly, we would like to reach out to those who take care of building
systems packages before the 3.2 release is shipped and resolve any
issues that may come up.

[issue]: https://bugs.ruby-lang.org/issues/18481

Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>
Co-authored-by: Noah Gibbs <the.codefolio.guy@gmail.com>
Co-authored-by: Kevin Newton <kddnewton@gmail.com>

1 files changed, 218 insertions, 0 deletions

diff --git a/yjit/src/disasm.rs b/yjit/src/disasm.rs
new file mode 100644
index 0000000000..97edc786bc
--- /dev/null
+++ b/yjit/src/disasm.rs
@@ -0,0 +1,218 @@
+use crate::asm::*;
+use crate::codegen::*;
+use crate::core::*;
+use crate::cruby::*;
+use crate::yjit::yjit_enabled_p;
+use std::fmt::Write;
+
+/// Primitive called in yjit.rb
+/// Produce a string representing the disassembly for an ISEQ
+#[no_mangle]
+pub extern "C" fn rb_yjit_disasm_iseq(_ec: EcPtr, _ruby_self: VALUE, iseqw: VALUE) -> VALUE {
+    #[cfg(not(feature = "disasm"))]
+    {
+        let _ = iseqw;
+        return Qnil;
+    }
+
+    #[cfg(feature = "disasm")]
+    {
+        // TODO:
+        //if unsafe { CLASS_OF(iseqw) != rb_cISeq } {
+        //    return Qnil;
+        //}
+
+        if !yjit_enabled_p() {
+            return Qnil;
+        }
+
+        // Get the iseq pointer from the wrapper
+        let iseq = unsafe { rb_iseqw_to_iseq(iseqw) };
+
+        let out_string = disasm_iseq(iseq);
+
+        return rust_str_to_ruby(&out_string);
+    }
+}
+
+#[cfg(feature = "disasm")]
+fn disasm_iseq(iseq: IseqPtr) -> String {
+    let mut out = String::from("");
+
+    // Get a list of block versions generated for this iseq
+    let mut block_list = get_iseq_block_list(iseq);
+
+    // Get a list of codeblocks relevant to this iseq
+    let global_cb = CodegenGlobals::get_inline_cb();
+
+    // Sort the blocks by increasing start addresses
+    block_list.sort_by(|a, b| {
+        use std::cmp::Ordering;
+
+        // Get the start addresses for each block
+        let addr_a = a.borrow().get_start_addr().unwrap().raw_ptr();
+        let addr_b = b.borrow().get_start_addr().unwrap().raw_ptr();
+
+        if addr_a < addr_b {
+            Ordering::Less
+        } else if addr_a == addr_b {
+            Ordering::Equal
+        } else {
+            Ordering::Greater
+        }
+    });
+
+    // Compute total code size in bytes for all blocks in the function
+    let mut total_code_size = 0;
+    for blockref in &block_list {
+        total_code_size += blockref.borrow().code_size();
+    }
+
+    // Initialize capstone
+    extern crate capstone;
+    use capstone::prelude::*;
+    let cs = Capstone::new()
+        .x86()
+        .mode(arch::x86::ArchMode::Mode64)
+        .syntax(arch::x86::ArchSyntax::Intel)
+        .build()
+        .unwrap();
+
+    out.push_str(&format!("NUM BLOCK VERSIONS: {}\n", block_list.len()));
+    out.push_str(&format!(
+        "TOTAL INLINE CODE SIZE: {} bytes\n",
+        total_code_size
+    ));
+
+    // For each block, sorted by increasing start address
+    for block_idx in 0..block_list.len() {
+        let block = block_list[block_idx].borrow();
+        let blockid = block.get_blockid();
+        let end_idx = block.get_end_idx();
+        let start_addr = block.get_start_addr().unwrap().raw_ptr();
+        let end_addr = block.get_end_addr().unwrap().raw_ptr();
+        let code_size = block.code_size();
+
+        // Write some info about the current block
+        let block_ident = format!(
+            "BLOCK {}/{}, ISEQ RANGE [{},{}), {} bytes ",
+            block_idx + 1,
+            block_list.len(),
+            blockid.idx,
+            end_idx,
+            code_size
+        );
+        out.push_str(&format!("== {:=<60}\n", block_ident));
+
+        // Disassemble the instructions
+        let code_slice = unsafe { std::slice::from_raw_parts(start_addr, code_size) };
+        let insns = cs.disasm_all(code_slice, start_addr as u64).unwrap();
+
+        // For each instruction in this block
+        for insn in insns.as_ref() {
+            // Comments for this block
+            if let Some(comment_list) = global_cb.comments_at(insn.address() as usize) {
+                for comment in comment_list {
+                    out.push_str(&format!("  \x1b[1m# {}\x1b[0m\n", comment));
+                }
+            }
+            out.push_str(&format!("  {}\n", insn));
+        }
+
+        // If this is not the last block
+        if block_idx < block_list.len() - 1 {
+            // Compute the size of the gap between this block and the next
+            let next_block = block_list[block_idx + 1].borrow();
+            let next_start_addr = next_block.get_start_addr().unwrap().raw_ptr();
+            let gap_size = (next_start_addr as usize) - (end_addr as usize);
+
+            // Log the size of the gap between the blocks if nonzero
+            if gap_size > 0 {
+                out.push_str(&format!("... {} byte gap ...\n", gap_size));
+            }
+        }
+    }
+
+    return out;
+}
+
+/// Primitive called in yjit.rb
+/// Produce a list of instructions compiled for an isew
+#[no_mangle]
+pub extern "C" fn rb_yjit_insns_compiled(_ec: EcPtr, _ruby_self: VALUE, iseqw: VALUE) -> VALUE {
+    {
+        // TODO:
+        //if unsafe { CLASS_OF(iseqw) != rb_cISeq } {
+        //    return Qnil;
+        //}
+
+        if !yjit_enabled_p() {
+            return Qnil;
+        }
+
+        // Get the iseq pointer from the wrapper
+        let iseq = unsafe { rb_iseqw_to_iseq(iseqw) };
+
+        // Get the list of instructions compiled
+        let insn_vec = insns_compiled(iseq);
+
+        unsafe {
+            let insn_ary = rb_ary_new_capa((insn_vec.len() * 2) as i64);
+
+            // For each instruction compiled
+            for idx in 0..insn_vec.len() {
+                let op_name = &insn_vec[idx].0;
+                let insn_idx = insn_vec[idx].1;
+
+                let op_sym = rust_str_to_sym(&op_name);
+
+                // Store the instruction index and opcode symbol
+                rb_ary_store(
+                    insn_ary,
+                    (2 * idx + 0) as i64,
+                    VALUE::fixnum_from_usize(insn_idx as usize),
+                );
+                rb_ary_store(insn_ary, (2 * idx + 1) as i64, op_sym);
+            }
+
+            insn_ary
+        }
+    }
+}
+
+fn insns_compiled(iseq: IseqPtr) -> Vec<(String, u32)> {
+    let mut insn_vec = Vec::new();
+
+    // Get a list of block versions generated for this iseq
+    let block_list = get_iseq_block_list(iseq);
+
+    // For each block associated with this iseq
+    for blockref in &block_list {
+        let block = blockref.borrow();
+        let start_idx = block.get_blockid().idx;
+        let end_idx = block.get_end_idx();
+        assert!(end_idx <= unsafe { get_iseq_encoded_size(iseq) });
+
+        // For each YARV instruction in the block
+        let mut insn_idx = start_idx;
+        while insn_idx < end_idx {
+            // Get the current pc and opcode
+            let pc = unsafe { rb_iseq_pc_at_idx(iseq, insn_idx) };
+            // try_into() call below is unfortunate. Maybe pick i32 instead of usize for opcodes.
+            let opcode: usize = unsafe { rb_iseq_opcode_at_pc(iseq, pc) }
+                .try_into()
+                .unwrap();
+
+            // Get the mnemonic for this opcode
+            let op_name = insn_name(opcode);
+
+            // Add the instruction to the list
+            insn_vec.push((op_name, insn_idx));
+
+            // Move to the next instruction
+            insn_idx += insn_len(opcode);
+        }
+    }
+
+    return insn_vec;
+}