Revert "ZJIT: Allow ccalls above 7 arguments"

This reverts commit 2f151e7.

The SP decrement (push) before the call do not match up with
the pops after the call, so registers were restored incorrectly.

Code from:

    ./miniruby --zjit-call-threshold=1 --zjit-dump-disasm -e 'p Time.new(1992, 9, 23, 23, 0, 0, :std)'

    str x11, [sp, #-0x10]!
    str x12, [sp, #-0x10]!
    stur x7, [sp]              # last argument
    mov x0, x20
    mov x7, x6
    mov x6, x5
    mov x5, x4
    mov x4, x3
    mov x3, x2
    mov x2, x1
    ldur x1, [x29, #-0x20]
    mov x16, #0xccfc
    movk x16, #0x2e7, lsl #16
    movk x16, #1, lsl #32
    blr x16
    ldr x12, [sp], #0x10       # supposed to match str x12, [sp, #-0x10]!, but got last argument
    ldr x11, [sp], #0x10

Author: XrXr

zjit/src/asm/arm64/opnd.rs

@@ -98,8 +98,6 @@ 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,15 +24,13 @@ 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; 8] = [
+pub const C_ARG_OPNDS: [Opnd; 6] = [
     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(X6_REG),
-    Opnd::Reg(X7_REG)
+    Opnd::Reg(X5_REG)
 ];
 
 // C return value register on this platform
@@ -201,8 +199,6 @@ pub const ALLOC_REGS: &[Reg] = &[
     X3_REG,
     X4_REG,
     X5_REG,
-    X6_REG,
-    X7_REG,
     X11_REG,
     X12_REG,
 ];
@@ -235,7 +231,7 @@ impl Assembler {
 
     /// Get a list of all of the caller-saved registers
     pub fn get_caller_save_regs() -> Vec<Reg> {
-        vec![X1_REG, X6_REG, X7_REG, X9_REG, X10_REG, X11_REG, X12_REG, X13_REG, X14_REG, X15_REG]
+        vec![X1_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
@@ -492,13 +488,31 @@ impl Assembler {
                 }
                 */
                 Insn::CCall { opnds, .. } => {
-                    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
-                        };
-                    });
+                    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![];
                     asm.push_insn(insn);
                 },
                 Insn::Cmp { left, right } => {
@@ -1843,19 +1857,17 @@ mod tests {
 
         assert_disasm_snapshot!(cb.disasm(), @r"
         0x0: str x1, [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]!
+        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]!
         ");
-        assert_snapshot!(cb.hexdump(), @"e10f1ff8e60f1ff8e70f1ff8e90f1ff8ea0f1ff8eb0f1ff8ec0f1ff8ed0f1ff8ee0f1ff8ef0f1ff810423bd5f00f1ff8");
+        assert_snapshot!(cb.hexdump(), @"e10f1ff8e90f1ff8ea0f1ff8eb0f1ff8ec0f1ff8ed0f1ff8ee0f1ff8ef0f1ff810423bd5f00f1ff8");
     }
 
     #[test]
@@ -1875,11 +1887,9 @@ mod tests {
         0x18: ldr x11, [sp], #0x10
         0x1c: ldr x10, [sp], #0x10
         0x20: ldr x9, [sp], #0x10
-        0x24: ldr x7, [sp], #0x10
-        0x28: ldr x6, [sp], #0x10
-        0x2c: ldr x1, [sp], #0x10
+        0x24: ldr x1, [sp], #0x10
         ");
-        assert_snapshot!(cb.hexdump(), @"10421bd5f00741f8ef0741f8ee0741f8ed0741f8ec0741f8eb0741f8ea0741f8e90741f8e70741f8e60741f8e10741f8");
+        assert_snapshot!(cb.hexdump(), @"10421bd5f00741f8ef0741f8ee0741f8ed0741f8ec0741f8eb0741f8ea0741f8e90741f8e10741f8");
     }
 
     #[test]
@@ -2648,14 +2658,14 @@ mod tests {
         ]);
         asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
 
-        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_disasm_snapshot!(cb.disasm(), @"
+            0x0: mov x15, x0
+            0x4: mov x0, x1
+            0x8: mov x1, x15
+            0xc: mov x16, #0
+            0x10: blr x16
         ");
-        assert_snapshot!(cb.hexdump(), @"ef0301aae10300aae0030faa100080d200023fd6");
+        assert_snapshot!(cb.hexdump(), @"ef0300aae00301aae1030faa100080d200023fd6");
     }
 
     #[test]
@@ -2671,17 +2681,17 @@ mod tests {
         ]);
         asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
 
-        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_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_snapshot!(cb.hexdump(), @"ef0301aae10300aae0030faaef0303aae30302aae2030faa100080d200023fd6");
+        assert_snapshot!(cb.hexdump(), @"ef0302aae20303aae3030faaef0300aae00301aae1030faa100080d200023fd6");
     }
 
     #[test]
@@ -2696,15 +2706,15 @@ mod tests {
         ]);
         asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
 
-        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_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_snapshot!(cb.hexdump(), @"ef0301aae10302aae20300aae0030faa100080d200023fd6");
+        assert_snapshot!(cb.hexdump(), @"ef0300aae00301aae10302aae2030faa100080d200023fd6");
     }
 
     #[test]

zjit/src/backend/lir.rs

@@ -1559,8 +1559,9 @@ impl Assembler
                 frame_setup_idxs.push(asm.insns.len());
             }
 
-            let before_ccall = match &insn {
-                Insn::CCall { .. } if !pool.is_empty() => {
+            let before_ccall = match (&insn, iterator.peek().map(|(_, insn)| insn)) {
+                (Insn::ParallelMov { .. }, Some(Insn::CCall { .. })) |
+                (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) {
@@ -1611,25 +1612,6 @@ 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
@@ -1721,23 +1703,7 @@ impl Assembler
             // Push instruction(s)
             let is_ccall = matches!(insn, Insn::CCall { .. });
             match insn {
-                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));
-                        }
-                    }
-
+                Insn::ParallelMov { moves } => {
                     // 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 {
@@ -1747,12 +1713,13 @@ 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: vec![], fptr, start_marker: None, end_marker: None, out });
+                    asm.push_insn(Insn::CCall { opnds, 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,7 +351,21 @@ impl Assembler {
                     };
                     asm.push_insn(insn);
                 },
-                Insn::CCall { .. } => {
+                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![];
                     asm.push_insn(insn);
                 },
                 Insn::Lea { .. } => {

zjit/src/codegen.rs

@@ -398,12 +398,15 @@ 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 the callee doesn't know how to read arguments from the stack.
+        // Give up SendWithoutBlockDirect for 6+ args since asm.ccall() doesn't support it.
         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))),
@@ -450,6 +453,10 @@ 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: _ } => {
@@ -715,6 +722,10 @@ 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 {