ZJIT: Support variadic C calls (ruby#14575)

* ZJIT: Support variadic C calls

This reduces the `dynamic_send_count` in `liquid-render` by ~21%

* ZJIT: Reuse gen_push_frame

* ZJIT: Avoid optimizing variadic C call when tracing is enabled

Author: st0012

test/ruby/test_zjit.rb

@@ -2710,6 +2710,44 @@ def test
     }, insns: [:invokeblock]
   end
 
+  def test_ccall_variadic_with_multiple_args
+    assert_compiles "[1, 2, 3]", %q{
+      def test
+        a = []
+        a.push(1, 2, 3)
+        a
+      end
+
+      test
+      test
+    }, insns: [:opt_send_without_block]
+  end
+
+  def test_ccall_variadic_with_no_args
+    assert_compiles "[1]", %q{
+      def test
+        a = [1]
+        a.push
+      end
+
+      test
+      test
+    }, insns: [:opt_send_without_block]
+  end
+
+  def test_ccall_variadic_with_no_args_causing_argument_error
+    assert_compiles ":error", %q{
+      def test
+        format
+      rescue ArgumentError
+        :error
+      end
+
+      test
+      test
+    }, insns: [:opt_send_without_block]
+  end
+
   private
 
   # Assert that every method call in `test_script` can be compiled by ZJIT

zjit.c

@@ -158,6 +158,23 @@ rb_zjit_defined_ivar(VALUE obj, ID id, VALUE pushval)
     return result ? pushval : Qnil;
 }
 
+bool
+rb_zjit_method_tracing_currently_enabled(void)
+{
+    rb_event_flag_t tracing_events;
+    if (rb_multi_ractor_p()) {
+        tracing_events = ruby_vm_event_enabled_global_flags;
+    }
+    else {
+        // At the time of writing, events are never removed from
+        // ruby_vm_event_enabled_global_flags so always checking using it would
+        // mean we don't compile even after tracing is disabled.
+        tracing_events = rb_ec_ractor_hooks(GET_EC())->events;
+    }
+
+    return tracing_events & (RUBY_EVENT_C_CALL | RUBY_EVENT_C_RETURN);
+}
+
 bool
 rb_zjit_insn_leaf(int insn, const VALUE *opes)
 {

zjit/bindgen/src/main.rs

@@ -274,6 +274,7 @@ fn main() {
         .allowlist_function("rb_zjit_local_id")
         .allowlist_function("rb_set_cfp_(pc|sp)")
         .allowlist_function("rb_c_method_tracing_currently_enabled")
+        .allowlist_function("rb_zjit_method_tracing_currently_enabled")
         .allowlist_function("rb_full_cfunc_return")
         .allowlist_function("rb_assert_(iseq|cme)_handle")
         .allowlist_function("rb_IMEMO_TYPE_P")

zjit/src/codegen.rs

@@ -397,6 +397,9 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::GuardBlockParamProxy { level, state } => no_output!(gen_guard_block_param_proxy(jit, asm, level, &function.frame_state(state))),
         Insn::PatchPoint { invariant, state } => no_output!(gen_patch_point(jit, asm, invariant, &function.frame_state(*state))),
         Insn::CCall { cfun, args, name: _, return_type: _, elidable: _ } => gen_ccall(asm, *cfun, opnds!(args)),
+        Insn::CCallVariadic { cfun, recv, args, name: _, cme, state } => {
+            gen_ccall_variadic(jit, asm, *cfun, opnd!(recv), opnds!(args), *cme, &function.frame_state(*state))
+        }
         Insn::GetIvar { self_val, id, state: _ } => gen_getivar(asm, opnd!(self_val), *id),
         Insn::SetGlobal { id, val, state } => no_output!(gen_setglobal(jit, asm, *id, opnd!(val), &function.frame_state(*state))),
         Insn::GetGlobal { id, state } => gen_getglobal(jit, asm, *id, &function.frame_state(*state)),
@@ -638,6 +641,52 @@ fn gen_ccall(asm: &mut Assembler, cfun: *const u8, args: Vec<Opnd>) -> lir::Opnd
     asm.ccall(cfun, args)
 }
 
+/// Generate code for a variadic C function call
+/// func(int argc, VALUE *argv, VALUE recv)
+fn gen_ccall_variadic(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    cfun: *const u8,
+    recv: Opnd,
+    args: Vec<Opnd>,
+    cme: *const rb_callable_method_entry_t,
+    state: &FrameState,
+) -> lir::Opnd {
+    gen_prepare_non_leaf_call(jit, asm, state);
+
+    gen_push_frame(asm, args.len(), state, ControlFrame {
+        recv,
+        iseq: None,
+        cme,
+        frame_type: VM_FRAME_MAGIC_CFUNC | VM_FRAME_FLAG_CFRAME | VM_ENV_FLAG_LOCAL,
+    });
+
+    asm_comment!(asm, "switch to new SP register");
+    let sp_offset = (state.stack().len() - args.len() + VM_ENV_DATA_SIZE.as_usize()) * SIZEOF_VALUE;
+    let new_sp = asm.add(SP, sp_offset.into());
+    asm.mov(SP, new_sp);
+
+    asm_comment!(asm, "switch to new CFP");
+    let new_cfp = asm.sub(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
+    asm.mov(CFP, new_cfp);
+    asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);
+
+    let argv_ptr = gen_push_opnds(jit, asm, &args);
+    let result = asm.ccall(cfun, vec![args.len().into(), argv_ptr, recv]);
+    gen_pop_opnds(asm, &args);
+
+    asm_comment!(asm, "pop C frame");
+    let new_cfp = asm.add(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
+    asm.mov(CFP, new_cfp);
+    asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);
+
+    asm_comment!(asm, "restore SP register for the caller");
+    let new_sp = asm.sub(SP, sp_offset.into());
+    asm.mov(SP, new_sp);
+
+    result
+}
+
 /// Emit an uncached instance variable lookup
 fn gen_getivar(asm: &mut Assembler, recv: Opnd, id: ID) -> Opnd {
     asm_ccall!(asm, rb_ivar_get, recv, id.0.into())
@@ -1054,7 +1103,7 @@ fn gen_send_without_block_direct(
     // TODO: Lazily materialize caller frames on side exits or when needed
     gen_push_frame(asm, args.len(), state, ControlFrame {
         recv,
-        iseq,
+        iseq: Some(iseq),
         cme,
         frame_type: VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL,
     });
@@ -1582,7 +1631,7 @@ fn gen_prepare_non_leaf_call(jit: &JITState, asm: &mut Assembler, state: &FrameS
 /// Frame metadata written by gen_push_frame()
 struct ControlFrame {
     recv: Opnd,
-    iseq: IseqPtr,
+    iseq: Option<IseqPtr>,
     cme: *const rb_callable_method_entry_t,
     frame_type: u32,
 }
@@ -1594,7 +1643,11 @@ fn gen_push_frame(asm: &mut Assembler, argc: usize, state: &FrameState, frame: C
     // See vm_push_frame() for details
     asm_comment!(asm, "push cme, specval, frame type");
     // ep[-2]: cref of cme
-    let local_size = unsafe { get_iseq_body_local_table_size(frame.iseq) } as i32;
+    let local_size = if let Some(iseq) = frame.iseq {
+        (unsafe { get_iseq_body_local_table_size(iseq) }) as i32
+    } else {
+        0
+    };
     let ep_offset = state.stack().len() as i32 + local_size - argc as i32 + VM_ENV_DATA_SIZE as i32 - 1;
     asm.store(Opnd::mem(64, SP, (ep_offset - 2) * SIZEOF_VALUE_I32), VALUE::from(frame.cme).into());
     // ep[-1]: block_handler or prev EP
@@ -1609,9 +1662,19 @@ fn gen_push_frame(asm: &mut Assembler, argc: usize, state: &FrameState, frame: C
     }
 
     asm_comment!(asm, "push callee control frame");
-    // cfp_opnd(RUBY_OFFSET_CFP_PC): written by the callee frame on side-exits or non-leaf calls
-    // cfp_opnd(RUBY_OFFSET_CFP_SP): written by the callee frame on side-exits or non-leaf calls
-    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_ISEQ), VALUE::from(frame.iseq).into());
+
+    if let Some(iseq) = frame.iseq {
+        // cfp_opnd(RUBY_OFFSET_CFP_PC): written by the callee frame on side-exits or non-leaf calls
+        // cfp_opnd(RUBY_OFFSET_CFP_SP): written by the callee frame on side-exits or non-leaf calls
+        asm.mov(cfp_opnd(RUBY_OFFSET_CFP_ISEQ), VALUE::from(iseq).into());
+    } else {
+        // C frames don't have a PC and ISEQ
+        asm.mov(cfp_opnd(RUBY_OFFSET_CFP_PC), 0.into());
+        let new_sp = asm.lea(Opnd::mem(64, SP, (ep_offset + 1) * SIZEOF_VALUE_I32));
+        asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SP), new_sp);
+        asm.mov(cfp_opnd(RUBY_OFFSET_CFP_ISEQ), 0.into());
+    }
+
     asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SELF), frame.recv);
     let ep = asm.lea(Opnd::mem(64, SP, ep_offset * SIZEOF_VALUE_I32));
     asm.mov(cfp_opnd(RUBY_OFFSET_CFP_EP), ep);

zjit/src/cruby_bindings.inc.rs

@@ -621,6 +621,17 @@ pub enum Insn {
     /// `name` is for printing purposes only
     CCall { cfun: *const u8, args: Vec<InsnId>, name: ID, return_type: Type, elidable: bool },
 
+    /// Call a variadic C function with signature: func(int argc, VALUE *argv, VALUE recv)
+    /// This handles frame setup, argv creation, and frame teardown all in one
+    CCallVariadic {
+        cfun: *const u8,
+        recv: InsnId,
+        args: Vec<InsnId>,
+        cme: *const rb_callable_method_entry_t,
+        name: ID,
+        state: InsnId,
+    },
+
     /// Un-optimized fallback implementation (dynamic dispatch) for send-ish instructions
     /// Ignoring keyword arguments etc for now
     SendWithoutBlock {
@@ -940,6 +951,13 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             },
+            Insn::CCallVariadic { cfun,  recv, args, name, .. } => {
+                write!(f, "CCallVariadic {}@{:p}, {recv}", name.contents_lossy(), self.ptr_map.map_ptr(cfun))?;
+                for arg in args {
+                    write!(f, ", {arg}")?;
+                }
+                Ok(())
+            },
             Insn::Snapshot { state } => write!(f, "Snapshot {}", state.print(self.ptr_map)),
             Insn::Defined { op_type, v, .. } => {
                 // op_type (enum defined_type) printing logic from iseq.c.
@@ -1422,6 +1440,9 @@ impl Function {
             &ObjectAlloc { val, state } => ObjectAlloc { val: find!(val), state },
             &ObjectAllocClass { class, state } => ObjectAllocClass { class, state: find!(state) },
             &CCall { cfun, ref args, name, return_type, elidable } => CCall { cfun, args: find_vec!(args), name, return_type, elidable },
+            &CCallVariadic { cfun, recv, ref args, cme, name, state } => CCallVariadic {
+                cfun, recv: find!(recv), args: find_vec!(args), cme, name, state
+            },
             &Defined { op_type, obj, pushval, v, state } => Defined { op_type, obj, pushval, v: find!(v), state: find!(state) },
             &DefinedIvar { self_val, pushval, id, state } => DefinedIvar { self_val: find!(self_val), pushval, id, state },
             &NewArray { ref elements, state } => NewArray { elements: find_vec!(elements), state: find!(state) },
@@ -1509,6 +1530,7 @@ impl Function {
             Insn::ObjectAlloc { .. } => types::HeapObject,
             Insn::ObjectAllocClass { class, .. } => Type::from_class(*class),
             Insn::CCall { return_type, .. } => *return_type,
+            Insn::CCallVariadic { .. } => types::BasicObject,
             Insn::GuardType { val, guard_type, .. } => self.type_of(*val).intersection(*guard_type),
             Insn::GuardTypeNot { .. } => types::BasicObject,
             Insn::GuardBitEquals { val, expected, .. } => self.type_of(*val).intersection(Type::from_value(*expected)),
@@ -2076,9 +2098,39 @@ impl Function {
                         return Ok(());
                     }
                 }
+                // Variadic method
                 -1 => {
-                    // (argc, argv, self) parameter form
-                    // Falling through for now
+                    if unsafe { rb_zjit_method_tracing_currently_enabled() } {
+                        return Err(());
+                    }
+                    // The method gets a pointer to the first argument
+                    // func(int argc, VALUE *argv, VALUE recv)
+                    let ci_flags = unsafe { vm_ci_flag(call_info) };
+                    if ci_flags & VM_CALL_ARGS_SIMPLE != 0 {
+                        fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoTracePoint, state });
+                        fun.push_insn(block, Insn::PatchPoint {
+                            invariant: Invariant::MethodRedefined {
+                                klass: recv_class,
+                                method: method_id,
+                                cme: method
+                            },
+                            state
+                        });
+
+                        let cfun = unsafe { get_mct_func(cfunc) }.cast();
+                        let ccall = fun.push_insn(block, Insn::CCallVariadic {
+                            cfun,
+                            recv,
+                            args,
+                            cme: method,
+                            name: method_id,
+                            state,
+                        });
+
+                        fun.make_equal_to(send_insn_id, ccall);
+                        return Ok(());
+                    }
+                    // Fall through for complex cases (splat, kwargs, etc.)
                 }
                 -2 => {
                     // (self, args_ruby_array) parameter form
@@ -2379,6 +2431,7 @@ impl Function {
             }
             &Insn::Send { recv, ref args, state, .. }
             | &Insn::SendWithoutBlock { recv, ref args, state, .. }
+            | &Insn::CCallVariadic { recv, ref args, state, .. }
             | &Insn::SendWithoutBlockDirect { recv, ref args, state, .. }
             | &Insn::InvokeSuper { recv, ref args, state, .. } => {
                 worklist.push_back(recv);
@@ -6895,6 +6948,26 @@ mod opt_tests {
         ");
     }
 
+    #[test]
+    fn test_optimize_variadic_ccall() {
+        eval("
+            def test
+              puts 'Hello'
+            end
+            test; test
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0(v0:BasicObject):
+          v4:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+          v6:StringExact = StringCopy v4
+          PatchPoint MethodRedefined(Object@0x1008, puts@0x1010, cme:0x1018)
+          v16:BasicObject = CCallVariadic puts@0x1040, v0, v6
+          CheckInterrupts
+          Return v16
+        ");
+    }
+
     #[test]
     fn test_dont_optimize_fixnum_add_if_redefined() {
         eval("