YJIT: Inline basic Ruby methods (#8855)

* YJIT: Inline basic Ruby methods

* YJIT: Fix "InsnOut operand made it past register allocation"

checktype should not generate a useless instruction.

3 files changed, 59 insertions, 4 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 8f668bcf86..eba1de713c 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -2752,7 +2752,6 @@ fn gen_checktype(
     if let RUBY_T_STRING | RUBY_T_ARRAY | RUBY_T_HASH = type_val {
         let val_type = asm.ctx.get_opnd_type(StackOpnd(0));
         let val = asm.stack_pop(1);
-        let val = asm.load(val);
 
         // Check if we know from type information
         match val_type.known_value_type() {
@@ -2770,6 +2769,7 @@ fn gen_checktype(
 
         let ret = asm.new_label("ret");
 
+        let val = asm.load(val);
         if !val_type.is_heap() {
             // if (SPECIAL_CONST_P(val)) {
             // Return Qfalse via REG1 if not on heap
@@ -5417,6 +5417,7 @@ fn gen_send_cfunc(
         if let Some(known_cfunc_codegen) = codegen_p {
             if known_cfunc_codegen(jit, asm, ocb, ci, cme, block, argc, recv_known_klass) {
                 assert_eq!(expected_stack_after, asm.ctx.get_stack_size() as i32);
+                gen_counter_incr(asm, Counter::num_send_known_cfunc);
                 // cfunc codegen generated code. Terminate the block so
                 // there isn't multiple calls in the same block.
                 jump_to_next_insn(jit, asm, ocb);
@@ -5852,6 +5853,31 @@ fn gen_send_bmethod(
     gen_send_iseq(jit, asm, ocb, iseq, ci, frame_type, Some(capture.ep), cme, block, flags, argc, None)
 }
 
+/// Return the ISEQ's return value if it consists of only putnil/putobject and leave.
+fn iseq_get_return_value(iseq: IseqPtr) -> Option<VALUE> {
+    // Expect only two instructions and one possible operand
+    let iseq_size = unsafe { get_iseq_encoded_size(iseq) };
+    if !(2..=3).contains(&iseq_size) {
+        return None;
+    }
+
+    // Get the first two instructions
+    let first_insn = iseq_opcode_at_idx(iseq, 0);
+    let second_insn = iseq_opcode_at_idx(iseq, insn_len(first_insn as usize));
+
+    // Extract the return value if known
+    if second_insn != YARVINSN_leave {
+        return None;
+    }
+    match first_insn {
+        YARVINSN_putnil => Some(Qnil),
+        YARVINSN_putobject => unsafe { Some(*rb_iseq_pc_at_idx(iseq, 1)) },
+        YARVINSN_putobject_INT2FIX_0_ => Some(VALUE::fixnum_from_usize(0)),
+        YARVINSN_putobject_INT2FIX_1_ => Some(VALUE::fixnum_from_usize(1)),
+        _ => None,
+    }
+}
+
 fn gen_send_iseq(
     jit: &mut JITState,
     asm: &mut Assembler,
@@ -6112,8 +6138,6 @@ fn gen_send_iseq(
     if let (None, Some(builtin_info), true, false) = (block, builtin_func, builtin_attrs & BUILTIN_ATTR_LEAF != 0, opt_send_call) {
         let builtin_argc = unsafe { (*builtin_info).argc };
         if builtin_argc + 1 < (C_ARG_OPNDS.len() as i32) {
-            asm_comment!(asm, "inlined leaf builtin");
-
             // We pop the block arg without using it because:
             //  - the builtin is leaf, so it promises to not `yield`.
             //  - no leaf builtins have block param at the time of writing, and
@@ -6126,6 +6150,9 @@ fn gen_send_iseq(
                 asm.stack_pop(1);
             }
 
+            asm_comment!(asm, "inlined leaf builtin");
+            gen_counter_incr(asm, Counter::num_send_leaf_builtin);
+
             // Skip this if it doesn't trigger GC
             if builtin_attrs & BUILTIN_ATTR_NO_GC == 0 {
                 // The callee may allocate, e.g. Integer#abs on a Bignum.
@@ -6152,10 +6179,29 @@ fn gen_send_iseq(
             // Note: assuming that the leaf builtin doesn't change local variables here.
             // Seems like a safe assumption.
 
-            return Some(KeepCompiling);
+            // Let guard chains share the same successor
+            jump_to_next_insn(jit, asm, ocb);
+            return Some(EndBlock);
         }
     }
 
+    // Inline simple ISEQs whose return value is known at compile time
+    if let (Some(value), None, false) = (iseq_get_return_value(iseq), block_arg_type, opt_send_call) {
+        asm_comment!(asm, "inlined simple ISEQ");
+        gen_counter_incr(asm, Counter::num_send_inline);
+
+        // Pop receiver and arguments
+        asm.stack_pop(argc as usize + if captured_opnd.is_some() { 0 } else { 1 });
+
+        // Push the return value
+        let stack_ret = asm.stack_push(Type::from(value));
+        asm.mov(stack_ret, value.into());
+
+        // Let guard chains share the same successor
+        jump_to_next_insn(jit, asm, ocb);
+        return Some(EndBlock);
+    }
+
     // Stack overflow check
     // Note that vm_push_frame checks it against a decremented cfp, hence the multiply by 2.
     // #define CHECK_VM_STACK_OVERFLOW0(cfp, sp, margin)
diff --git a/yjit/src/cruby.rs b/yjit/src/cruby.rs
index 6896ae9fce..a9e35cdbdd 100644
--- a/yjit/src/cruby.rs
+++ b/yjit/src/cruby.rs
@@ -252,6 +252,12 @@ pub fn iseq_pc_to_insn_idx(iseq: IseqPtr, pc: *mut VALUE) -> Option<u16> {
     unsafe { pc.offset_from(pc_zero) }.try_into().ok()
 }
 
+/// Given an ISEQ pointer and an instruction index, return an opcode.
+pub fn iseq_opcode_at_idx(iseq: IseqPtr, insn_idx: u32) -> u32 {
+    let pc = unsafe { rb_iseq_pc_at_idx(iseq, insn_idx) };
+    unsafe { rb_iseq_opcode_at_pc(iseq, pc) as u32 }
+}
+
 /// Opaque execution-context type from vm_core.h
 #[repr(C)]
 pub struct rb_execution_context_struct {
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 4cd7c1979d..4fc1825991 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -480,6 +480,9 @@ make_counters! {
     num_send_x86_rel32,
     num_send_x86_reg,
     num_send_dynamic,
+    num_send_inline,
+    num_send_leaf_builtin,
+    num_send_known_cfunc,
 
     num_getivar_megamorphic,
     num_setivar_megamorphic,