From a6aac0695ad863ad208faf126297d8afa698af35 Mon Sep 17 00:00:00 2001 From: Takashi Kokubun Date: Tue, 15 Jul 2025 13:24:37 -0700 Subject: [PATCH 1/3] ZJIT: Fix SP alignment on JIT entry for x86_64 --- test/ruby/test_zjit.rb | 20 ++++++++++++++++++++ zjit/src/codegen.rs | 23 +++++++++++------------ 2 files changed, 31 insertions(+), 12 deletions(-) diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb index fdcee3cafd858e..299600eef95d53 100644 --- a/test/ruby/test_zjit.rb +++ b/test/ruby/test_zjit.rb @@ -999,6 +999,26 @@ def test_require_rubygems_with_auto_compact }, call_threshold: 2 end + def test_profile_under_nested_jit_call + assert_compiles '[nil, nil, 3]', %q{ + def profile + 1 + 2 + end + + def jit_call(flag) + if flag + profile + end + end + + def entry(flag) + jit_call(flag) + end + + [entry(false), entry(false), entry(true)] + }, call_threshold: 2 + end + def test_bop_redefinition assert_runs '[3, :+, 100]', %q{ def test diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs index ce3decd56f319b..917528ca0a5c9f 100644 --- a/zjit/src/codegen.rs +++ b/zjit/src/codegen.rs @@ -156,10 +156,6 @@ fn gen_entry(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function, function_pt // Restore registers for CFP, EC, and SP after use asm_comment!(asm, "exit to the interpreter"); - // On x86_64, maintain 16-byte stack alignment - if cfg!(target_arch = "x86_64") { - asm.cpop_into(SP); - } asm.cpop_into(SP); asm.cpop_into(EC); asm.cpop_into(CFP); @@ -511,10 +507,6 @@ fn gen_entry_prologue(asm: &mut Assembler, iseq: IseqPtr) { asm.cpush(CFP); asm.cpush(EC); asm.cpush(SP); - // On x86_64, maintain 16-byte stack alignment - if cfg!(target_arch = "x86_64") { - asm.cpush(SP); - } // EC and CFP are passed as arguments asm.mov(EC, C_ARG_OPNDS[0]); @@ -1157,12 +1149,19 @@ fn max_num_params(function: &Function) -> usize { /// the function needs to allocate on the stack for the stack frame. fn aligned_stack_bytes(num_slots: usize) -> usize { // Both x86_64 and arm64 require the stack to be aligned to 16 bytes. - // Since SIZEOF_VALUE is 8 bytes, we need to round up the size to the nearest even number. - let num_slots = if num_slots % 2 == 0 { - num_slots - } else { + let num_slots = if cfg!(target_arch = "x86_64") && num_slots % 2 == 0 { + // On x86_64, since the call instruction bumps the stack pointer by 8 bytes on entry, + // we need to round up `num_slots` to an odd number. num_slots + 1 + } else if cfg!(target_arch = "aarch64") && num_slots % 2 == 1 { + // On arm64, the stack pointer is always aligned to 16 bytes, so we need to round up + // `num_slots`` to an even number. + num_slots + 1 + } else { + num_slots }; + + const _: () = assert!(SIZEOF_VALUE == 8, "aligned_stack_bytes() assumes SIZEOF_VALUE == 8"); num_slots * SIZEOF_VALUE } From 388cfee02be173b58d42d40c05badbecb2abfb10 Mon Sep 17 00:00:00 2001 From: Takashi Kokubun Date: Tue, 15 Jul 2025 13:24:43 -0700 Subject: [PATCH 2/3] ZJIT: Restore SP on side-exit chains --- zjit/src/backend/lir.rs | 68 +++++++++++++++++++++++++---------------- zjit/src/codegen.rs | 15 +++++---- zjit/src/hir.rs | 1 + 3 files changed, 51 insertions(+), 33 deletions(-) diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs index e5453e4d554b67..94c53569b4e95c 100644 --- a/zjit/src/backend/lir.rs +++ b/zjit/src/backend/lir.rs @@ -270,6 +270,14 @@ impl From for Opnd { } } +/// Set of things we need to restore for side exits. +#[derive(Clone, Debug)] +pub struct SideExitContext { + pub pc: *const VALUE, + pub stack: Vec, + pub locals: Vec, +} + /// Branch target (something that we can jump to) /// for branch instructions #[derive(Clone, Debug)] @@ -281,12 +289,14 @@ pub enum Target Label(Label), /// Side exit to the interpreter SideExit { - pc: *const VALUE, - stack: Vec, - locals: Vec, - c_stack_bytes: usize, + /// Context to restore on regular side exits. None for side exits right + /// after JIT-to-JIT calls because we restore them before the JIT call. + context: Option, + /// We use this to enrich asm comments. reason: SideExitReason, - // Some if the side exit should write this label. We use it for patch points. + /// The number of bytes we need to adjust the C stack pointer by. + c_stack_bytes: usize, + /// Some if the side exit should write this label. We use it for patch points. label: Option