YJIT: Lazily push a frame for specialized C funcs

Co-authored-by: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>

bootstraptest/test_yjit.rb

@@ -2046,6 +2046,57 @@ def getbyte(s, i)
   [getbyte("a", 0), getbyte("a", 1), getbyte("a", -1), getbyte("a", -2), getbyte("a", "a")]
 } unless rjit_enabled? # Not yet working on RJIT
 
+# Basic test for String#setbyte
+assert_equal 'AoZ', %q{
+  s = "foo"
+  s.setbyte(0, 65)
+  s.setbyte(-1, 90)
+  s
+}
+
+# String#setbyte IndexError
+assert_equal 'String#setbyte', %q{
+  def ccall = "".setbyte(1, 0)
+  begin
+    ccall
+  rescue => e
+    e.backtrace.first.split("'").last
+  end
+}
+
+# String#setbyte TypeError
+assert_equal 'String#setbyte', %q{
+  def ccall = "".setbyte(nil, 0)
+  begin
+    ccall
+  rescue => e
+    e.backtrace.first.split("'").last
+  end
+}
+
+# String#setbyte FrozenError
+assert_equal 'String#setbyte', %q{
+  def ccall = "a".freeze.setbyte(0, 0)
+  begin
+    ccall
+  rescue => e
+    e.backtrace.first.split("'").last
+  end
+}
+
+# non-leaf String#setbyte
+assert_equal 'String#setbyte', %q{
+  def to_int
+    @caller = caller
+    0
+  end
+
+  def ccall = "a".setbyte(self, 98)
+  ccall
+
+  @caller.first.split("'").last
+}
+
 # non-leaf String#byteslice
 assert_equal 'TypeError', %q{
   def ccall = "".byteslice(nil, nil)

common.mk

@@ -6706,6 +6706,7 @@ error.$(OBJEXT): {$(VPATH)}util.h
 error.$(OBJEXT): {$(VPATH)}vm_core.h
 error.$(OBJEXT): {$(VPATH)}vm_opts.h
 error.$(OBJEXT): {$(VPATH)}warning.rbinc
+error.$(OBJEXT): {$(VPATH)}yjit.h
 eval.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 eval.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
 eval.$(OBJEXT): $(CCAN_DIR)/list/list.h
@@ -11413,6 +11414,7 @@ object.$(OBJEXT): {$(VPATH)}vm_core.h
 object.$(OBJEXT): {$(VPATH)}vm_debug.h
 object.$(OBJEXT): {$(VPATH)}vm_opts.h
 object.$(OBJEXT): {$(VPATH)}vm_sync.h
+object.$(OBJEXT): {$(VPATH)}yjit.h
 pack.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 pack.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
 pack.$(OBJEXT): $(CCAN_DIR)/list/list.h

error.c

@@ -48,6 +48,7 @@
 #include "ruby/util.h"
 #include "ruby_assert.h"
 #include "vm_core.h"
+#include "yjit.h"
 
 #include "builtin.h"
 
@@ -1409,6 +1410,7 @@ rb_exc_new_cstr(VALUE etype, const char *s)
 VALUE
 rb_exc_new_str(VALUE etype, VALUE str)
 {
+    rb_yjit_check_pc(GET_EC()->cfp->pc);
     StringValue(str);
     return rb_class_new_instance(1, &str, etype);
 }
@@ -3827,6 +3829,7 @@ inspect_frozen_obj(VALUE obj, VALUE mesg, int recur)
 void
 rb_error_frozen_object(VALUE frozen_obj)
 {
+    rb_yjit_check_pc(GET_EC()->cfp->pc);
     VALUE debug_info;
     const ID created_info = id_debug_created_info;
     VALUE mesg = rb_sprintf("can't modify frozen %"PRIsVALUE": ",

object.c

@@ -42,6 +42,7 @@
 #include "ruby/assert.h"
 #include "builtin.h"
 #include "shape.h"
+#include "yjit.h"
 
 /* Flags of RObject
  *
@@ -3238,6 +3239,7 @@ rb_to_integer_with_id_exception(VALUE val, const char *method, ID mid, int raise
     VALUE v;
 
     if (RB_INTEGER_TYPE_P(val)) return val;
+    rb_yjit_check_pc(GET_EC()->cfp->pc);
     v = try_to_int(val, mid, raise);
     if (!raise && NIL_P(v)) return Qnil;
     if (!RB_INTEGER_TYPE_P(v)) {

string.c

@@ -6174,7 +6174,7 @@ rb_str_getbyte(VALUE str, VALUE index)
  *
  *  Related: String#getbyte.
  */
-static VALUE
+VALUE
 rb_str_setbyte(VALUE str, VALUE index, VALUE value)
 {
     long pos = NUM2LONG(index);

vm_insnhelper.c

@@ -3535,6 +3535,24 @@ vm_call_cfunc_with_frame_(rb_execution_context_t *ec, rb_control_frame_t *reg_cf
     return val;
 }
 
+// Push a C method frame for a given cme. This is called when JIT code skipped
+// pushing a frame but the C method reached a point where a frame is needed.
+void
+rb_vm_push_cfunc_frame(const rb_callable_method_entry_t *cme)
+{
+    VM_ASSERT(cme->def->type == VM_METHOD_TYPE_CFUNC);
+    rb_execution_context_t *ec = GET_EC();
+    VALUE *sp = ec->cfp->sp;
+    VALUE recv = *(sp - cme->def->body.cfunc.argc - 1);
+    VALUE frame_type = VM_FRAME_MAGIC_CFUNC | VM_FRAME_FLAG_CFRAME | VM_ENV_FLAG_LOCAL;
+    VALUE block_handler = VM_BLOCK_HANDLER_NONE;
+#if VM_CHECK_MODE > 0
+    // Clean up the stack canary since we're about to satisfy the "leaf or lazy push" assumption
+    *(GET_EC()->cfp->sp) = Qfalse;
+#endif
+    vm_push_frame(ec, NULL, frame_type, recv, block_handler, (VALUE)cme, 0, ec->cfp->sp, 0, 0);
+}
+
 // If true, cc->call needs to include `CALLER_SETUP_ARG` (i.e. can't be skipped in fastpath)
 bool
 rb_splat_or_kwargs_p(const struct rb_callinfo *restrict ci)

yjit.h

@@ -45,6 +45,7 @@ void rb_yjit_before_ractor_spawn(void);
 void rb_yjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic, unsigned insn_idx);
 void rb_yjit_tracing_invalidate_all(void);
 void rb_yjit_show_usage(int help, int highlight, unsigned int width, int columns);
+void rb_yjit_check_pc(const VALUE *pc);
 
 #else
 // !USE_YJIT
@@ -66,6 +67,7 @@ static inline void rb_yjit_iseq_free(void *payload) {}
 static inline void rb_yjit_before_ractor_spawn(void) {}
 static inline void rb_yjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic, unsigned insn_idx) {}
 static inline void rb_yjit_tracing_invalidate_all(void) {}
+static inline void rb_yjit_check_pc(const VALUE *pc) {}
 
 #endif // #if USE_YJIT
 

yjit.rb

@@ -328,6 +328,9 @@ def _print_stats(out: $stderr) # :nodoc:
       out.puts "num_getivar_megamorphic: " + format_number(11, stats[:num_getivar_megamorphic])
       out.puts "num_setivar_megamorphic: " + format_number(11, stats[:num_setivar_megamorphic])
       out.puts "num_opt_case_megamorphic: " + format_number(10, stats[:num_opt_case_dispatch_megamorphic])
+      out.puts "num_cfunc_lazy_frame:  " + format_number(13, stats[:num_cfunc_lazy_frame])
+      out.puts "num_cfunc_check_pc:    " + format_number(13, stats[:num_cfunc_check_pc])
+      out.puts "num_cfunc_check_pc_push: " + format_number_pct(11, stats[:num_cfunc_check_pc_push], stats[:num_cfunc_check_pc])
       out.puts "num_throw:             " + format_number(13, stats[:num_throw])
       out.puts "num_throw_break:       " + format_number_pct(13, stats[:num_throw_break], stats[:num_throw])
       out.puts "num_throw_retry:       " + format_number_pct(13, stats[:num_throw_retry], stats[:num_throw])

yjit/src/codegen.rs

@@ -411,6 +411,45 @@ fn gen_save_sp_with_offset(asm: &mut Assembler, offset: i8) {
     }
 }
 
+/// Basically jit_prepare_non_leaf_call(), but this registers the current PC
+/// to lazily push a C method frame when it's necessary.
+fn jit_prepare_lazy_cfunc_frame(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    cme: *const rb_callable_method_entry_t,
+    recv_opnd: YARVOpnd,
+) -> bool {
+    // Get the next PC. jit_save_pc() saves that PC.
+    let pc: *mut VALUE = unsafe {
+        let cur_insn_len = insn_len(jit.get_opcode()) as isize;
+        jit.get_pc().offset(cur_insn_len)
+    };
+
+    let pc_to_cfunc = CodegenGlobals::get_pc_to_cfunc();
+    match pc_to_cfunc.get(&pc) {
+        Some(&pc_cme) if pc_cme != cme => {
+            // Bail out if it's not the only cme on this callsite.
+            incr_counter!(num_cfunc_lazy_frame);
+            return false;
+        }
+        _ => {
+            // Let rb_yjit_check_pc() lazily push a C frame on this PC.
+            pc_to_cfunc.insert(pc, cme);
+        }
+    }
+
+    // Save the PC to trigger a lazy frame push, and save the SP to get the receiver.
+    // The C func may call a method that doesn't raise, so prepare for invalidation too.
+    jit_prepare_non_leaf_call(jit, asm);
+
+    // Make sure we're ready for calling vm_push_cfunc_frame().
+    let cfunc_argc = unsafe { get_mct_argc(get_cme_def_body_cfunc(cme)) };
+    assert_eq!(recv_opnd, StackOpnd(cfunc_argc as u8)); // We use this for cfp->self
+    assert!(asm.get_leaf_ccall()); // It checks the stack canary we set for known_cfunc_codegen.
+
+    true
+}
+
 /// jit_save_pc() + gen_save_sp(). Should be used before calling a routine that could:
 ///  - Perform GC allocation
 ///  - Take the VM lock through RB_VM_LOCK_ENTER()
@@ -5465,6 +5504,35 @@ fn jit_rb_str_getbyte(
     true
 }
 
+fn jit_rb_str_setbyte(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    _ocb: &mut OutlinedCb,
+    _ci: *const rb_callinfo,
+    cme: *const rb_callable_method_entry_t,
+    _block: Option<BlockHandler>,
+    _argc: i32,
+    _known_recv_class: Option<VALUE>,
+) -> bool {
+    // Raises when index is out of range. Lazily push a frame in that case.
+    if !jit_prepare_lazy_cfunc_frame(jit, asm, cme, StackOpnd(2)) {
+        return false;
+    }
+    asm_comment!(asm, "String#setbyte");
+
+    let value = asm.stack_opnd(0);
+    let index = asm.stack_opnd(1);
+    let recv = asm.stack_opnd(2);
+
+    let ret_opnd = asm.ccall(rb_str_setbyte as *const u8, vec![recv, index, value]);
+    asm.stack_pop(3); // Keep them on stack during ccall for GC
+
+    let out_opnd = asm.stack_push(Type::UnknownImm);
+    asm.mov(out_opnd, ret_opnd);
+
+    true
+}
+
 // Codegen for rb_str_to_s()
 // When String#to_s is called on a String instance, the method returns self and
 // most of the overhead comes from setting up the method call. We observed that
@@ -9693,6 +9761,7 @@ pub fn yjit_reg_method_codegen_fns() {
         yjit_reg_method(rb_cString, "size", jit_rb_str_length);
         yjit_reg_method(rb_cString, "bytesize", jit_rb_str_bytesize);
         yjit_reg_method(rb_cString, "getbyte", jit_rb_str_getbyte);
+        yjit_reg_method(rb_cString, "setbyte", jit_rb_str_setbyte);
         yjit_reg_method(rb_cString, "byteslice", jit_rb_str_byteslice);
         yjit_reg_method(rb_cString, "<<", jit_rb_str_concat);
         yjit_reg_method(rb_cString, "+@", jit_rb_str_uplus);
@@ -9769,6 +9838,9 @@ pub struct CodegenGlobals {
 
     /// Page indexes for outlined code that are not associated to any ISEQ.
     ocb_pages: Vec<usize>,
+
+    /// Lazily push a frame when it calls rb_yjit_check_pc() with a PC in this HashMap
+    pc_to_cfunc: HashMap<*mut VALUE, *const rb_callable_method_entry_t>,
 }
 
 /// For implementing global code invalidation. A position in the inline
@@ -9860,6 +9932,7 @@ impl CodegenGlobals {
             entry_stub_hit_trampoline,
             global_inval_patches: Vec::new(),
             ocb_pages,
+            pc_to_cfunc: HashMap::new(),
         };
 
         // Initialize the codegen globals instance
@@ -9938,6 +10011,10 @@ impl CodegenGlobals {
     pub fn get_ocb_pages() -> &'static Vec<usize> {
         &CodegenGlobals::get_instance().ocb_pages
     }
+
+    pub fn get_pc_to_cfunc() -> &'static mut HashMap<*mut VALUE, *const rb_callable_method_entry_t> {
+        &mut CodegenGlobals::get_instance().pc_to_cfunc
+    }
 }
 
 #[cfg(test)]

yjit/src/cruby.rs

@@ -117,6 +117,7 @@ extern "C" {
         ci: *const rb_callinfo,
     ) -> *const rb_callable_method_entry_t;
     pub fn rb_hash_empty_p(hash: VALUE) -> VALUE;
+    pub fn rb_str_setbyte(str: VALUE, index: VALUE, value: VALUE) -> VALUE;
     pub fn rb_vm_splat_array(flag: VALUE, ary: VALUE) -> VALUE;
     pub fn rb_vm_concat_array(ary1: VALUE, ary2st: VALUE) -> VALUE;
     pub fn rb_vm_concat_to_array(ary1: VALUE, ary2st: VALUE) -> VALUE;
@@ -142,6 +143,7 @@ extern "C" {
     ) -> VALUE;
     pub fn rb_vm_ic_hit_p(ic: IC, reg_ep: *const VALUE) -> bool;
     pub fn rb_vm_stack_canary() -> VALUE;
+    pub fn rb_vm_push_cfunc_frame(cme: *const rb_callable_method_entry_t);
 }
 
 // Renames

yjit/src/stats.rs

@@ -583,6 +583,10 @@ make_counters! {
     num_throw_retry,
     num_throw_return,
 
+    num_cfunc_lazy_frame,
+    num_cfunc_check_pc,
+    num_cfunc_check_pc_push,
+
     iseq_stack_too_large,
     iseq_too_long,
 

yjit/src/yjit.rs

@@ -207,3 +207,19 @@ pub extern "C" fn rb_yjit_simulate_oom_bang(_ec: EcPtr, _ruby_self: VALUE) -> VA
 
     return Qnil;
 }
+
+/// Push a C method frame if the given PC is supposed to lazily push one.
+/// This is called from rb_raise() (at rb_exc_new_str()) and other functions
+/// that may make a method call (e.g. rb_to_int()).
+#[no_mangle]
+pub extern "C" fn rb_yjit_check_pc(pc: *mut VALUE) {
+    if !yjit_enabled_p() {
+        return;
+    }
+
+    incr_counter!(num_cfunc_check_pc);
+    if let Some(&cme) = CodegenGlobals::get_pc_to_cfunc().get(&pc) {
+        incr_counter!(num_cfunc_check_pc_push);
+        unsafe { rb_vm_push_cfunc_frame(cme) }
+    }
+}