YJIT: Skip padding jumps to side exits on Arm (#6790)

YJIT: Skip padding jumps to side exits

Co-authored-by: Maxime Chevalier-Boisvert <[email protected]>
Co-authored-by: Alan Wu <[email protected]>

Co-authored-by: Maxime Chevalier-Boisvert <[email protected]>
Co-authored-by: Alan Wu <[email protected]>

7 files changed, 61 insertions, 47 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index cbda5eec05..5ff0ac3709 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -624,7 +624,7 @@ impl Assembler
         /// called when lowering any of the conditional jump instructions.
         fn emit_conditional_jump<const CONDITION: u8>(cb: &mut CodeBlock, target: Target) {
             match target {
-                Target::CodePtr(dst_ptr) => {
+                Target::CodePtr(dst_ptr) | Target::SideExitPtr(dst_ptr) => {
                     let dst_addr = dst_ptr.into_i64();
                     let src_addr = cb.get_write_ptr().into_i64();
 
@@ -659,10 +659,12 @@ impl Assembler
                         load_insns + 2
                     };
 
-                    // We need to make sure we have at least 6 instructions for
-                    // every kind of jump for invalidation purposes, so we're
-                    // going to write out padding nop instructions here.
-                    for _ in num_insns..6 { nop(cb); }
+                    if let Target::CodePtr(_) = target {
+                        // We need to make sure we have at least 6 instructions for
+                        // every kind of jump for invalidation purposes, so we're
+                        // going to write out padding nop instructions here.
+                        for _ in num_insns..6 { nop(cb); }
+                    }
                 },
                 Target::Label(label_idx) => {
                     // Here we're going to save enough space for ourselves and
@@ -690,7 +692,7 @@ impl Assembler
             ldr_post(cb, opnd, A64Opnd::new_mem(64, C_SP_REG, C_SP_STEP));
         }
 
-        fn emit_jmp_ptr(cb: &mut CodeBlock, dst_ptr: CodePtr) {
+        fn emit_jmp_ptr(cb: &mut CodeBlock, dst_ptr: CodePtr, padding: bool) {
             let src_addr = cb.get_write_ptr().into_i64();
             let dst_addr = dst_ptr.into_i64();
 
@@ -707,11 +709,13 @@ impl Assembler
                 num_insns + 1
             };
 
-            // Make sure it's always a consistent number of
-            // instructions in case it gets patched and has to
-            // use the other branch.
-            for _ in num_insns..(JMP_PTR_BYTES / 4) {
-                nop(cb);
+            if padding {
+                // Make sure it's always a consistent number of
+                // instructions in case it gets patched and has to
+                // use the other branch.
+                for _ in num_insns..(JMP_PTR_BYTES / 4) {
+                    nop(cb);
+                }
             }
         }
 
@@ -723,7 +727,7 @@ impl Assembler
         fn emit_jmp_ptr_with_invalidation(cb: &mut CodeBlock, dst_ptr: CodePtr) {
             #[cfg(not(test))]
             let start = cb.get_write_ptr();
-            emit_jmp_ptr(cb, dst_ptr);
+            emit_jmp_ptr(cb, dst_ptr, true);
             #[cfg(not(test))]
             {
                 let end = cb.get_write_ptr();
@@ -963,7 +967,10 @@ impl Assembler
                 Insn::Jmp(target) => {
                     match target {
                         Target::CodePtr(dst_ptr) => {
-                            emit_jmp_ptr(cb, *dst_ptr);
+                            emit_jmp_ptr(cb, *dst_ptr, true);
+                        },
+                        Target::SideExitPtr(dst_ptr) => {
+                            emit_jmp_ptr(cb, *dst_ptr, false);
                         },
                         Target::Label(label_idx) => {
                             // Here we're going to save enough space for
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 973fe89b6d..de8bac7465 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -254,9 +254,10 @@ impl From<VALUE> for Opnd {
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub enum Target
 {
-    CodePtr(CodePtr),   // Pointer to a piece of YJIT-generated code (e.g. side-exit)
-    FunPtr(*const u8),  // Pointer to a C function
-    Label(usize),       // A label within the generated code
+    CodePtr(CodePtr),     // Pointer to a piece of YJIT-generated code
+    SideExitPtr(CodePtr), // Pointer to a side exit code
+    FunPtr(*const u8),    // Pointer to a C function
+    Label(usize),         // A label within the generated code
 }
 
 impl Target
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 68cd35574f..46dcad036f 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -565,7 +565,7 @@ impl Assembler
                 // Conditional jump to a label
                 Insn::Jmp(target) => {
                     match *target {
-                        Target::CodePtr(code_ptr) => jmp_ptr(cb, code_ptr),
+                        Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jmp_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jmp_label(cb, label_idx),
                         _ => unreachable!()
                     }
@@ -573,7 +573,7 @@ impl Assembler
 
                 Insn::Je(target) => {
                     match *target {
-                        Target::CodePtr(code_ptr) => je_ptr(cb, code_ptr),
+                        Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => je_ptr(cb, code_ptr),
                         Target::Label(label_idx) => je_label(cb, label_idx),
                         _ => unreachable!()
                     }
@@ -581,7 +581,7 @@ impl Assembler
 
                 Insn::Jne(target) => {
                     match *target {
-                        Target::CodePtr(code_ptr) => jne_ptr(cb, code_ptr),
+                        Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jne_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jne_label(cb, label_idx),
                         _ => unreachable!()
                     }
@@ -589,7 +589,7 @@ impl Assembler
 
                 Insn::Jl(target) => {
                     match *target {
-                        Target::CodePtr(code_ptr) => jl_ptr(cb, code_ptr),
+                        Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jl_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jl_label(cb, label_idx),
                         _ => unreachable!()
                     }
@@ -597,7 +597,7 @@ impl Assembler
 
                 Insn::Jbe(target) => {
                     match *target {
-                        Target::CodePtr(code_ptr) => jbe_ptr(cb, code_ptr),
+                        Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jbe_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jbe_label(cb, label_idx),
                         _ => unreachable!()
                     }
@@ -605,7 +605,7 @@ impl Assembler
 
                 Insn::Jz(target) => {
                     match *target {
-                        Target::CodePtr(code_ptr) => jz_ptr(cb, code_ptr),
+                        Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jz_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jz_label(cb, label_idx),
                         _ => unreachable!()
                     }
@@ -613,7 +613,7 @@ impl Assembler
 
                 Insn::Jnz(target) => {
                     match *target {
-                        Target::CodePtr(code_ptr) => jnz_ptr(cb, code_ptr),
+                        Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jnz_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jnz_label(cb, label_idx),
                         _ => unreachable!()
                     }
@@ -621,7 +621,7 @@ impl Assembler
 
                 Insn::Jo(target) => {
                     match *target {
-                        Target::CodePtr(code_ptr) => jo_ptr(cb, code_ptr),
+                        Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jo_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jo_label(cb, label_idx),
                         _ => unreachable!()
                     }
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index c66a293322..9735f1dce9 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -212,7 +212,7 @@ macro_rules! counted_exit {
             gen_counter_incr!(ocb_asm, $counter_name);
 
             // Jump to the existing side exit
-            ocb_asm.jmp($existing_side_exit.into());
+            ocb_asm.jmp($existing_side_exit.as_side_exit());
             ocb_asm.compile(ocb);
 
             // Pointer to the side-exit code
@@ -649,7 +649,7 @@ fn gen_check_ints(asm: &mut Assembler, side_exit: CodePtr) {
         not_mask,
     );
 
-    asm.jnz(Target::CodePtr(side_exit));
+    asm.jnz(Target::SideExitPtr(side_exit));
 }
 
 // Generate a stubbed unconditional jump to the next bytecode instruction.
@@ -1118,7 +1118,7 @@ fn gen_opt_plus(
         // Add arg0 + arg1 and test for overflow
         let arg0_untag = asm.sub(arg0, Opnd::Imm(1));
         let out_val = asm.add(arg0_untag, arg1);
-        asm.jo(side_exit.into());
+        asm.jo(side_exit.as_side_exit());
 
         // Push the output on the stack
         let dst = ctx.stack_push(Type::Fixnum);
@@ -1301,11 +1301,11 @@ fn guard_object_is_heap(
 
     // Test that the object is not an immediate
     asm.test(object_opnd, (RUBY_IMMEDIATE_MASK as u64).into());
-    asm.jnz(side_exit.into());
+    asm.jnz(side_exit.as_side_exit());
 
     // Test that the object is not false or nil
     asm.cmp(object_opnd, Qnil.into());
-    asm.jbe(side_exit.into());
+    asm.jbe(side_exit.as_side_exit());
 }
 
 fn guard_object_is_array(
@@ -1325,7 +1325,7 @@ fn guard_object_is_array(
 
     // Compare the result with T_ARRAY
     asm.cmp(flags_opnd, (RUBY_T_ARRAY as u64).into());
-    asm.jne(side_exit.into());
+    asm.jne(side_exit.as_side_exit());
 }
 
 fn guard_object_is_string(
@@ -1347,7 +1347,7 @@ fn guard_object_is_string(
 
     // Compare the result with T_STRING
     asm.cmp(flags_reg, Opnd::UImm(RUBY_T_STRING as u64));
-    asm.jne(side_exit.into());
+    asm.jne(side_exit.as_side_exit());
 }
 
 // push enough nils onto the stack to fill out an array
@@ -1625,7 +1625,7 @@ fn gen_setlocal_wc0(
         let side_exit = get_side_exit(jit, ocb, ctx);
 
         // if (flags & VM_ENV_FLAG_WB_REQUIRED) != 0
-        asm.jnz(side_exit.into());
+        asm.jnz(side_exit.as_side_exit());
     }
 
     // Set the type of the local variable in the context
@@ -1670,7 +1670,7 @@ fn gen_setlocal_generic(
         let side_exit = get_side_exit(jit, ocb, ctx);
 
         // if (flags & VM_ENV_FLAG_WB_REQUIRED) != 0
-        asm.jnz(side_exit.into());
+        asm.jnz(side_exit.as_side_exit());
     }
 
     // Pop the value to write from the stack
@@ -2296,19 +2296,19 @@ fn guard_two_fixnums(
 
     if arg0_type.is_heap() || arg1_type.is_heap() {
         asm.comment("arg is heap object");
-        asm.jmp(side_exit.into());
+        asm.jmp(side_exit.as_side_exit());
         return;
     }
 
     if arg0_type != Type::Fixnum && arg0_type.is_specific() {
         asm.comment("arg0 not fixnum");
-        asm.jmp(side_exit.into());
+        asm.jmp(side_exit.as_side_exit());
         return;
     }
 
     if arg1_type != Type::Fixnum && arg1_type.is_specific() {
         asm.comment("arg1 not fixnum");
-        asm.jmp(side_exit.into());
+        asm.jmp(side_exit.as_side_exit());
         return;
     }
 
@@ -2929,7 +2929,7 @@ fn gen_opt_minus(
 
         // Subtract arg0 - arg1 and test for overflow
         let val_untag = asm.sub(arg0, arg1);
-        asm.jo(side_exit.into());
+        asm.jo(side_exit.as_side_exit());
         let val = asm.add(val_untag, Opnd::Imm(1));
 
         // Push the output on the stack
@@ -2996,7 +2996,7 @@ fn gen_opt_mod(
 
         // Check for arg0 % 0
         asm.cmp(arg1, Opnd::Imm(VALUE::fixnum_from_usize(0).as_i64()));
-        asm.je(side_exit.into());
+        asm.je(side_exit.as_side_exit());
 
         // Call rb_fix_mod_fix(VALUE recv, VALUE obj)
         let ret = asm.ccall(rb_fix_mod_fix as *const u8, vec![arg0, arg1]);
@@ -3774,9 +3774,9 @@ fn jit_rb_str_concat(
         if !arg_type.is_heap() {
             asm.comment("guard arg not immediate");
             asm.test(arg_opnd, Opnd::UImm(RUBY_IMMEDIATE_MASK as u64));
-            asm.jnz(side_exit.into());
+            asm.jnz(side_exit.as_side_exit());
             asm.cmp(arg_opnd, Qnil.into());
-            asm.jbe(side_exit.into());
+            asm.jbe(side_exit.as_side_exit());
         }
         guard_object_is_string(asm, arg_opnd, side_exit);
     }
@@ -3908,7 +3908,7 @@ fn jit_obj_respond_to(
     // This is necessary because we have no guarantee that sym_opnd is a constant
     asm.comment("guard known mid");
     asm.cmp(sym_opnd, mid_sym.into());
-    asm.jne(side_exit.into());
+    asm.jne(side_exit.as_side_exit());
 
     jit_putobject(jit, ctx, asm, result);
 
@@ -4197,7 +4197,7 @@ fn gen_send_cfunc(
     asm.comment("stack overflow check");
     let stack_limit = asm.lea(ctx.sp_opnd((SIZEOF_VALUE * 4 + 2 * RUBY_SIZEOF_CONTROL_FRAME) as isize));
     asm.cmp(CFP, stack_limit);
-    asm.jbe(counted_exit!(ocb, side_exit, send_se_cf_overflow).into());
+    asm.jbe(counted_exit!(ocb, side_exit, send_se_cf_overflow).as_side_exit());
 
     // Number of args which will be passed through to the callee
     // This is adjusted by the kwargs being combined into a hash.
@@ -4818,7 +4818,7 @@ fn gen_send_iseq(
         (SIZEOF_VALUE as i32) * (num_locals + stack_max) + 2 * (RUBY_SIZEOF_CONTROL_FRAME as i32);
     let stack_limit = asm.lea(ctx.sp_opnd(locals_offs as isize));
     asm.cmp(CFP, stack_limit);
-    asm.jbe(counted_exit!(ocb, side_exit, send_se_cf_overflow).into());
+    asm.jbe(counted_exit!(ocb, side_exit, send_se_cf_overflow).as_side_exit());
 
     // push_splat_args does stack manipulation so we can no longer side exit
     if flags & VM_CALL_ARGS_SPLAT != 0 {
@@ -6460,7 +6460,7 @@ fn gen_getblockparam(
     asm.test(flags_opnd, VM_ENV_FLAG_WB_REQUIRED.into());
 
     // if (flags & VM_ENV_FLAG_WB_REQUIRED) != 0
-    asm.jnz(side_exit.into());
+    asm.jnz(side_exit.as_side_exit());
 
     // Convert the block handler in to a proc
     // call rb_vm_bh_to_procval(const rb_execution_context_t *ec, VALUE block_handler)
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 869c0859b5..b45c3f010f 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -2195,7 +2195,7 @@ pub fn invalidate_block_version(blockref: &BlockRef) {
             cb.set_write_ptr(block_start);
 
             let mut asm = Assembler::new();
-            asm.jmp(block_entry_exit.into());
+            asm.jmp(block_entry_exit.as_side_exit());
             cb.set_dropped_bytes(false);
             asm.compile(&mut cb);
 
diff --git a/yjit/src/invariants.rs b/yjit/src/invariants.rs
index b911ff2c92..1542b40db8 100644
--- a/yjit/src/invariants.rs
+++ b/yjit/src/invariants.rs
@@ -503,7 +503,7 @@ pub extern "C" fn rb_yjit_tracing_invalidate_all() {
             assert!(last_patch_end <= patch.inline_patch_pos.raw_ptr(), "patches should not overlap");
 
             let mut asm = crate::backend::ir::Assembler::new();
-            asm.jmp(patch.outlined_target_pos.into());
+            asm.jmp(patch.outlined_target_pos.as_side_exit());
 
             cb.set_write_ptr(patch.inline_patch_pos);
             cb.set_dropped_bytes(false);
diff --git a/yjit/src/virtualmem.rs b/yjit/src/virtualmem.rs
index 668e3d6a7b..4755d0ac76 100644
--- a/yjit/src/virtualmem.rs
+++ b/yjit/src/virtualmem.rs
@@ -3,7 +3,7 @@
 // usize->pointer casts is viable. It seems like a lot of work for us to participate for not much
 // benefit.
 
-use crate::utils::IntoUsize;
+use crate::{utils::IntoUsize, backend::ir::Target};
 
 #[cfg(not(test))]
 pub type VirtualMem = VirtualMemory<sys::SystemAllocator>;
@@ -62,6 +62,12 @@ pub trait Allocator {
 #[repr(C, packed)]
 pub struct CodePtr(*const u8);
 
+impl CodePtr {
+    pub fn as_side_exit(self) -> Target {
+        Target::SideExitPtr(self)
+    }
+}
+
 /// Errors that can happen when writing to [VirtualMemory]
 #[derive(Debug, PartialEq)]
 pub enum WriteError {