ZJIT: Allow ccalls above 7 arguments (ruby#15312)

ZJIT: Add stack support for CCalls

Author: aidenfoxivey

zjit/src/asm/arm64/opnd.rs

@@ -98,6 +98,8 @@ pub const X2_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 2 };
 pub const X3_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 3 };
 pub const X4_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 4 };
 pub const X5_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 5 };
+pub const X6_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 6 };
+pub const X7_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 7 };
 
 // caller-save registers
 pub const X9_REG: A64Reg = A64Reg { num_bits: 64, reg_no: 9 };

zjit/src/backend/arm64/mod.rs

@@ -24,13 +24,15 @@ pub const EC: Opnd = Opnd::Reg(X20_REG);
 pub const SP: Opnd = Opnd::Reg(X21_REG);
 
 // C argument registers on this platform
-pub const C_ARG_OPNDS: [Opnd; 6] = [
+pub const C_ARG_OPNDS: [Opnd; 8] = [
     Opnd::Reg(X0_REG),
     Opnd::Reg(X1_REG),
     Opnd::Reg(X2_REG),
     Opnd::Reg(X3_REG),
     Opnd::Reg(X4_REG),
-    Opnd::Reg(X5_REG)
+    Opnd::Reg(X5_REG),
+    Opnd::Reg(X6_REG),
+    Opnd::Reg(X7_REG)
 ];
 
 // C return value register on this platform
@@ -199,6 +201,8 @@ pub const ALLOC_REGS: &[Reg] = &[
     X3_REG,
     X4_REG,
     X5_REG,
+    X6_REG,
+    X7_REG,
     X11_REG,
     X12_REG,
 ];
@@ -231,7 +235,7 @@ impl Assembler {
 
     /// Get a list of all of the caller-saved registers
     pub fn get_caller_save_regs() -> Vec<Reg> {
-        vec![X1_REG, X9_REG, X10_REG, X11_REG, X12_REG, X13_REG, X14_REG, X15_REG]
+        vec![X1_REG, X6_REG, X7_REG, X9_REG, X10_REG, X11_REG, X12_REG, X13_REG, X14_REG, X15_REG]
     }
 
     /// How many bytes a call and a [Self::frame_setup] would change native SP
@@ -488,31 +492,13 @@ impl Assembler {
                 }
                 */
                 Insn::CCall { opnds, .. } => {
-                    assert!(opnds.len() <= C_ARG_OPNDS.len());
-
-                    // Load each operand into the corresponding argument
-                    // register.
-                    // Note: the iteration order is reversed to avoid corrupting x0,
-                    // which is both the return value and first argument register
-                    if !opnds.is_empty() {
-                        let mut args: Vec<(Opnd, Opnd)> = vec![];
-                        for (idx, opnd) in opnds.iter_mut().enumerate().rev() {
-                            // If the value that we're sending is 0, then we can use
-                            // the zero register, so in this case we'll just send
-                            // a UImm of 0 along as the argument to the move.
-                            let value = match opnd {
-                                Opnd::UImm(0) | Opnd::Imm(0) => Opnd::UImm(0),
-                                Opnd::Mem(_) => split_memory_address(asm, *opnd),
-                                _ => *opnd
-                            };
-                            args.push((C_ARG_OPNDS[idx], value));
-                        }
-                        asm.parallel_mov(args);
-                    }
-
-                    // Now we push the CCall without any arguments so that it
-                    // just performs the call.
-                    *opnds = vec![];
+                    opnds.iter_mut().for_each(|opnd| {
+                        *opnd = match opnd {
+                            Opnd::UImm(0) | Opnd::Imm(0) => Opnd::UImm(0),
+                            Opnd::Mem(_) => split_memory_address(asm, *opnd),
+                            _ => *opnd
+                        };
+                    });
                     asm.push_insn(insn);
                 },
                 Insn::Cmp { left, right } => {
@@ -1857,17 +1843,19 @@ mod tests {
 
         assert_disasm_snapshot!(cb.disasm(), @r"
         0x0: str x1, [sp, #-0x10]!
-        0x4: str x9, [sp, #-0x10]!
-        0x8: str x10, [sp, #-0x10]!
-        0xc: str x11, [sp, #-0x10]!
-        0x10: str x12, [sp, #-0x10]!
-        0x14: str x13, [sp, #-0x10]!
-        0x18: str x14, [sp, #-0x10]!
-        0x1c: str x15, [sp, #-0x10]!
-        0x20: mrs x16, nzcv
-        0x24: str x16, [sp, #-0x10]!
+        0x4: str x6, [sp, #-0x10]!
+        0x8: str x7, [sp, #-0x10]!
+        0xc: str x9, [sp, #-0x10]!
+        0x10: str x10, [sp, #-0x10]!
+        0x14: str x11, [sp, #-0x10]!
+        0x18: str x12, [sp, #-0x10]!
+        0x1c: str x13, [sp, #-0x10]!
+        0x20: str x14, [sp, #-0x10]!
+        0x24: str x15, [sp, #-0x10]!
+        0x28: mrs x16, nzcv
+        0x2c: str x16, [sp, #-0x10]!
         ");
-        assert_snapshot!(cb.hexdump(), @"e10f1ff8e90f1ff8ea0f1ff8eb0f1ff8ec0f1ff8ed0f1ff8ee0f1ff8ef0f1ff810423bd5f00f1ff8");
+        assert_snapshot!(cb.hexdump(), @"e10f1ff8e60f1ff8e70f1ff8e90f1ff8ea0f1ff8eb0f1ff8ec0f1ff8ed0f1ff8ee0f1ff8ef0f1ff810423bd5f00f1ff8");
     }
 
     #[test]
@@ -1887,9 +1875,11 @@ mod tests {
         0x18: ldr x11, [sp], #0x10
         0x1c: ldr x10, [sp], #0x10
         0x20: ldr x9, [sp], #0x10
-        0x24: ldr x1, [sp], #0x10
+        0x24: ldr x7, [sp], #0x10
+        0x28: ldr x6, [sp], #0x10
+        0x2c: ldr x1, [sp], #0x10
         ");
-        assert_snapshot!(cb.hexdump(), @"10421bd5f00741f8ef0741f8ee0741f8ed0741f8ec0741f8eb0741f8ea0741f8e90741f8e10741f8");
+        assert_snapshot!(cb.hexdump(), @"10421bd5f00741f8ef0741f8ee0741f8ed0741f8ec0741f8eb0741f8ea0741f8e90741f8e70741f8e60741f8e10741f8");
     }
 
     #[test]
@@ -2658,14 +2648,14 @@ mod tests {
         ]);
         asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
 
-        assert_disasm_snapshot!(cb.disasm(), @"
-            0x0: mov x15, x0
-            0x4: mov x0, x1
-            0x8: mov x1, x15
-            0xc: mov x16, #0
-            0x10: blr x16
+        assert_disasm_snapshot!(cb.disasm(), @r"
+        0x0: mov x15, x1
+        0x4: mov x1, x0
+        0x8: mov x0, x15
+        0xc: mov x16, #0
+        0x10: blr x16
         ");
-        assert_snapshot!(cb.hexdump(), @"ef0300aae00301aae1030faa100080d200023fd6");
+        assert_snapshot!(cb.hexdump(), @"ef0301aae10300aae0030faa100080d200023fd6");
     }
 
     #[test]
@@ -2681,17 +2671,17 @@ mod tests {
         ]);
         asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
 
-        assert_disasm_snapshot!(cb.disasm(), @"
-            0x0: mov x15, x2
-            0x4: mov x2, x3
-            0x8: mov x3, x15
-            0xc: mov x15, x0
-            0x10: mov x0, x1
-            0x14: mov x1, x15
-            0x18: mov x16, #0
-            0x1c: blr x16
+        assert_disasm_snapshot!(cb.disasm(), @r"
+        0x0: mov x15, x1
+        0x4: mov x1, x0
+        0x8: mov x0, x15
+        0xc: mov x15, x3
+        0x10: mov x3, x2
+        0x14: mov x2, x15
+        0x18: mov x16, #0
+        0x1c: blr x16
         ");
-        assert_snapshot!(cb.hexdump(), @"ef0302aae20303aae3030faaef0300aae00301aae1030faa100080d200023fd6");
+        assert_snapshot!(cb.hexdump(), @"ef0301aae10300aae0030faaef0303aae30302aae2030faa100080d200023fd6");
     }
 
     #[test]
@@ -2706,15 +2696,15 @@ mod tests {
         ]);
         asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
 
-        assert_disasm_snapshot!(cb.disasm(), @"
-            0x0: mov x15, x0
-            0x4: mov x0, x1
-            0x8: mov x1, x2
-            0xc: mov x2, x15
-            0x10: mov x16, #0
-            0x14: blr x16
+        assert_disasm_snapshot!(cb.disasm(), @r"
+        0x0: mov x15, x1
+        0x4: mov x1, x2
+        0x8: mov x2, x0
+        0xc: mov x0, x15
+        0x10: mov x16, #0
+        0x14: blr x16
         ");
-        assert_snapshot!(cb.hexdump(), @"ef0300aae00301aae10302aae2030faa100080d200023fd6");
+        assert_snapshot!(cb.hexdump(), @"ef0301aae10302aae20300aae0030faa100080d200023fd6");
     }
 
     #[test]

zjit/src/backend/lir.rs

@@ -1559,9 +1559,8 @@ impl Assembler
                 frame_setup_idxs.push(asm.insns.len());
             }
 
-            let before_ccall = match (&insn, iterator.peek().map(|(_, insn)| insn)) {
-                (Insn::ParallelMov { .. }, Some(Insn::CCall { .. })) |
-                (Insn::CCall { .. }, _) if !pool.is_empty() => {
+            let before_ccall = match &insn {
+                Insn::CCall { .. } if !pool.is_empty() => {
                     // If C_RET_REG is in use, move it to another register.
                     // This must happen before last-use registers are deallocated.
                     if let Some(vreg_idx) = pool.vreg_for(&C_RET_REG) {
@@ -1612,6 +1611,25 @@ impl Assembler
                 if cfg!(target_arch = "x86_64") && saved_regs.len() % 2 == 1 {
                     asm.cpush(Opnd::Reg(saved_regs.last().unwrap().0));
                 }
+                if let Insn::ParallelMov { moves } = &insn {
+                    if moves.len() > C_ARG_OPNDS.len() {
+                        let difference = moves.len().saturating_sub(C_ARG_OPNDS.len());
+
+                        #[cfg(target_arch = "x86_64")]
+                        let offset = {
+                            // double quadword alignment
+                            ((difference + 3) / 4) * 4
+                        };
+
+                        #[cfg(target_arch = "aarch64")]
+                        let offset = {
+                            // quadword alignment
+                            if difference % 2 == 0 { difference } else { difference + 1 }
+                        };
+
+                        asm.sub_into(NATIVE_STACK_PTR, (offset * 8).into());
+                    }
+                }
             }
 
             // Allocate a register for the output operand if it exists
@@ -1703,7 +1721,23 @@ impl Assembler
             // Push instruction(s)
             let is_ccall = matches!(insn, Insn::CCall { .. });
             match insn {
-                Insn::ParallelMov { moves } => {
+                Insn::CCall { opnds, fptr, start_marker, end_marker, out } => {
+                    let mut moves: Vec<(Opnd, Opnd)> = vec![];
+                    let num_reg_args = opnds.len().min(C_ARG_OPNDS.len());
+                    let num_stack_args = opnds.len().saturating_sub(C_ARG_OPNDS.len());
+
+                    if num_stack_args > 0 {
+                        for (i, opnd) in opnds.iter().skip(num_reg_args).enumerate() {
+                            moves.push((Opnd::mem(64, NATIVE_STACK_PTR, 8 * i as i32), *opnd));
+                        }
+                    }
+
+                    if num_reg_args > 0 {
+                        for (i, opnd) in opnds.iter().take(num_reg_args).enumerate() {
+                            moves.push((C_ARG_OPNDS[i], *opnd));
+                        }
+                    }
+
                     // For trampolines that use scratch registers, attempt to lower ParallelMov without scratch_reg.
                     if let Some(moves) = Self::resolve_parallel_moves(&moves, None) {
                         for (dst, src) in moves {
@@ -1713,13 +1747,12 @@ impl Assembler
                         // If it needs a scratch_reg, leave it to *_split_with_scratch_regs to handle it.
                         asm.push_insn(Insn::ParallelMov { moves });
                     }
-                }
-                Insn::CCall { opnds, fptr, start_marker, end_marker, out } => {
+
                     // Split start_marker and end_marker here to avoid inserting push/pop between them.
                     if let Some(start_marker) = start_marker {
                         asm.push_insn(Insn::PosMarker(start_marker));
                     }
-                    asm.push_insn(Insn::CCall { opnds, fptr, start_marker: None, end_marker: None, out });
+                    asm.push_insn(Insn::CCall { opnds: vec![], fptr, start_marker: None, end_marker: None, out });
                     if let Some(end_marker) = end_marker {
                         asm.push_insn(Insn::PosMarker(end_marker));
                     }

zjit/src/backend/x86_64/mod.rs

@@ -351,21 +351,7 @@ impl Assembler {
                     };
                     asm.push_insn(insn);
                 },
-                Insn::CCall { opnds, .. } => {
-                    assert!(opnds.len() <= C_ARG_OPNDS.len());
-
-                    // Load each operand into the corresponding argument register.
-                    if !opnds.is_empty() {
-                        let mut args: Vec<(Opnd, Opnd)> = vec![];
-                        for (idx, opnd) in opnds.iter_mut().enumerate() {
-                            args.push((C_ARG_OPNDS[idx], *opnd));
-                        }
-                        asm.parallel_mov(args);
-                    }
-
-                    // Now we push the CCall without any arguments so that it
-                    // just performs the call.
-                    *opnds = vec![];
+                Insn::CCall { .. } => {
                     asm.push_insn(insn);
                 },
                 Insn::Lea { .. } => {

zjit/src/codegen.rs

@@ -398,15 +398,12 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::Send { cd, blockiseq, state, reason, .. } => gen_send(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
         &Insn::SendForward { cd, blockiseq, state, reason, .. } => gen_send_forward(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
         &Insn::SendWithoutBlock { cd, state, reason, .. } => gen_send_without_block(jit, asm, cd, &function.frame_state(state), reason),
-        // Give up SendWithoutBlockDirect for 6+ args since asm.ccall() doesn't support it.
+        // Give up SendWithoutBlockDirect for 6+ args since the callee doesn't know how to read arguments from the stack.
         Insn::SendWithoutBlockDirect { cd, state, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() => // +1 for self
             gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), SendFallbackReason::TooManyArgsForLir),
         Insn::SendWithoutBlockDirect { cme, iseq, recv, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(recv), opnds!(args), &function.frame_state(*state)),
         &Insn::InvokeSuper { cd, blockiseq, state, reason, .. } => gen_invokesuper(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
         &Insn::InvokeBlock { cd, state, reason, .. } => gen_invokeblock(jit, asm, cd, &function.frame_state(state), reason),
-        // Ensure we have enough room fit ec, self, and arguments
-        // TODO remove this check when we have stack args (we can use Time.new to test it)
-        Insn::InvokeBuiltin { bf, state, .. } if bf.argc + 2 > (C_ARG_OPNDS.len() as i32) => return Err(*state),
         Insn::InvokeBuiltin { bf, leaf, args, state, .. } => gen_invokebuiltin(jit, asm, &function.frame_state(*state), bf, *leaf, opnds!(args)),
         &Insn::EntryPoint { jit_entry_idx } => no_output!(gen_entry_point(jit, asm, jit_entry_idx)),
         Insn::Return { val } => no_output!(gen_return(asm, opnd!(val))),
@@ -453,10 +450,6 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::GuardGreaterEq { left, right, state } => gen_guard_greater_eq(jit, asm, opnd!(left), opnd!(right), &function.frame_state(state)),
         Insn::PatchPoint { invariant, state } => no_output!(gen_patch_point(jit, asm, invariant, &function.frame_state(*state))),
         Insn::CCall { cfunc, recv, args, name, return_type: _, elidable: _ } => gen_ccall(asm, *cfunc, *name, opnd!(recv), opnds!(args)),
-        // Give up CCallWithFrame for 7+ args since asm.ccall() supports at most 6 args (recv + args).
-        // There's no test case for this because no core cfuncs have this many parameters. But C extensions could have such methods.
-        Insn::CCallWithFrame { cd, state, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() =>
-            gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), SendFallbackReason::CCallWithFrameTooManyArgs),
         Insn::CCallWithFrame { cfunc, recv, name, args, cme, state, blockiseq, .. } =>
             gen_ccall_with_frame(jit, asm, *cfunc, *name, opnd!(recv), opnds!(args), *cme, *blockiseq, &function.frame_state(*state)),
         Insn::CCallVariadic { cfunc, recv, name, args, cme, state, blockiseq, return_type: _, elidable: _ } => {
@@ -722,10 +715,6 @@ fn gen_fixnum_bit_check(asm: &mut Assembler, val: Opnd, index: u8) -> Opnd {
 }
 
 fn gen_invokebuiltin(jit: &JITState, asm: &mut Assembler, state: &FrameState, bf: &rb_builtin_function, leaf: bool, args: Vec<Opnd>) -> lir::Opnd {
-    assert!(bf.argc + 2 <= C_ARG_OPNDS.len() as i32,
-            "gen_invokebuiltin should not be called for builtin function {} with too many arguments: {}",
-            unsafe { std::ffi::CStr::from_ptr(bf.name).to_str().unwrap() },
-            bf.argc);
     if leaf {
         gen_prepare_leaf_call_with_gc(asm, state);
     } else {