Skip to content

Commit be55b77

Browse files
kddnewtonmaximecbjimmyhmiller
authored
Better b.cond usage on AArch64 (#6305)
* Better b.cond usage on AArch64 When we're lowering a conditional jump, we previously had a bit of a complicated setup where we could emit a conditional jump to skip over a jump that was the next instruction, and then write out the destination and use a branch register. Now instead we use the b.cond instruction if our offset fits (not common, but not unused either) and if it doesn't we write out an inverse condition to jump past loading the destination and branching directly. * Added an inverse fn for Condition (#443) Prevents the need to pass two params and potentially reduces errors. Co-authored-by: Jimmy Miller <jimmyhmiller@jimmys-mbp.lan> Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com> Co-authored-by: Jimmy Miller <jimmyhmiller@jimmys-mbp.lan>
1 parent 32a0591 commit be55b77

File tree

4 files changed

+100
-75
lines changed

4 files changed

+100
-75
lines changed

‎yjit/src/asm/arm64/arg/condition.rs

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,4 +19,34 @@ impl Condition {
1919
pub const GT: u8 = 0b1100; // greater than (signed)
2020
pub const LE: u8 = 0b1101; // less than or equal to (signed)
2121
pub const AL: u8 = 0b1110; // always
22-
}
22+
23+
pub const fn inverse(condition: u8) -> u8 {
24+
match condition {
25+
Condition::EQ => Condition::NE,
26+
Condition::NE => Condition::EQ,
27+
28+
Condition::CS => Condition::CC,
29+
Condition::CC => Condition::CS,
30+
31+
Condition::MI => Condition::PL,
32+
Condition::PL => Condition::MI,
33+
34+
Condition::VS => Condition::VC,
35+
Condition::VC => Condition::VS,
36+
37+
Condition::HI => Condition::LS,
38+
Condition::LS => Condition::HI,
39+
40+
Condition::LT => Condition::GE,
41+
Condition::GE => Condition::LT,
42+
43+
Condition::GT => Condition::LE,
44+
Condition::LE => Condition::GT,
45+
46+
Condition::AL => Condition::AL,
47+
48+
_ => panic!("Unknown condition")
49+
50+
}
51+
}
52+
}

‎yjit/src/asm/arm64/inst/branch_cond.rs

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,8 @@ pub struct BranchCond {
2020
impl BranchCond {
2121
/// B.cond
2222
/// https://developer.arm.com/documentation/ddi0596/2020-12/Base-Instructions/B-cond--Branch-conditionally-
23-
pub fn bcond(cond: u8, byte_offset: i32) -> Self {
24-
Self { cond, imm19: byte_offset >> 2 }
23+
pub fn bcond(cond: u8, imm19: i32) -> Self {
24+
Self { cond, imm19 }
2525
}
2626
}
2727

@@ -53,25 +53,25 @@ mod tests {
5353

5454
#[test]
5555
fn test_b_eq() {
56-
let result: u32 = BranchCond::bcond(Condition::EQ, 128).into();
56+
let result: u32 = BranchCond::bcond(Condition::EQ, 32).into();
5757
assert_eq!(0x54000400, result);
5858
}
5959

6060
#[test]
6161
fn test_b_vs() {
62-
let result: u32 = BranchCond::bcond(Condition::VS, 128).into();
62+
let result: u32 = BranchCond::bcond(Condition::VS, 32).into();
6363
assert_eq!(0x54000406, result);
6464
}
6565

6666
#[test]
6767
fn test_b_eq_max() {
68-
let result: u32 = BranchCond::bcond(Condition::EQ, (1 << 20) - 4).into();
68+
let result: u32 = BranchCond::bcond(Condition::EQ, (1 << 18) - 1).into();
6969
assert_eq!(0x547fffe0, result);
7070
}
7171

7272
#[test]
7373
fn test_b_eq_min() {
74-
let result: u32 = BranchCond::bcond(Condition::EQ, -(1 << 20)).into();
74+
let result: u32 = BranchCond::bcond(Condition::EQ, -(1 << 18)).into();
7575
assert_eq!(0x54800000, result);
7676
}
7777
}

‎yjit/src/asm/arm64/mod.rs

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -203,9 +203,10 @@ pub fn b(cb: &mut CodeBlock, imm26: A64Opnd) {
203203
cb.write_bytes(&bytes);
204204
}
205205

206-
/// Whether or not the offset between two instructions fits into the b.cond
207-
/// instruction. If it doesn't, then we have to load the value into a register
208-
/// first, then use the b.cond instruction to skip past a direct jump.
206+
/// Whether or not the offset in number of instructions between two instructions
207+
/// fits into the b.cond instruction. If it doesn't, then we have to load the
208+
/// value into a register first, then use the b.cond instruction to skip past a
209+
/// direct jump.
209210
pub const fn bcond_offset_fits_bits(offset: i64) -> bool {
210211
imm_fits_bits(offset, 21) && (offset & 0b11 == 0)
211212
}
@@ -216,7 +217,7 @@ pub fn bcond(cb: &mut CodeBlock, cond: u8, byte_offset: A64Opnd) {
216217
A64Opnd::Imm(imm) => {
217218
assert!(bcond_offset_fits_bits(imm), "The immediate operand must be 21 bits or less and be aligned to a 2-bit boundary.");
218219

219-
BranchCond::bcond(cond, imm as i32).into()
220+
BranchCond::bcond(cond, (imm / 4) as i32).into()
220221
},
221222
_ => panic!("Invalid operand combination to bcond instruction."),
222223
};

‎yjit/src/backend/arm64/mod.rs

Lines changed: 58 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -565,64 +565,42 @@ impl Assembler
565565
fn emit_conditional_jump<const CONDITION: u8>(cb: &mut CodeBlock, target: Target) {
566566
match target {
567567
Target::CodePtr(dst_ptr) => {
568-
let dst_addr = dst_ptr.into_u64();
569-
//let src_addr = cb.get_write_ptr().into_i64() + 4;
570-
//let offset = dst_addr - src_addr;
571-
572-
// If the condition is met, then we'll skip past the
573-
// next instruction, put the address in a register, and
574-
// jump to it.
575-
bcond(cb, CONDITION, A64Opnd::new_imm(8));
576-
577-
// If we get to this instruction, then the condition
578-
// wasn't met, in which case we'll jump past the
579-
// next instruction that perform the direct jump.
580-
581-
b(cb, A64Opnd::new_imm(2i64 + emit_load_size(dst_addr) as i64));
582-
let num_insns = emit_load_value(cb, Assembler::SCRATCH0, dst_addr);
583-
br(cb, Assembler::SCRATCH0);
584-
for _ in num_insns..4 {
585-
nop(cb);
586-
}
568+
let dst_addr = dst_ptr.into_i64();
569+
let src_addr = cb.get_write_ptr().into_i64();
570+
let offset = dst_addr - src_addr;
587571

588-
/*
589-
// If the jump offset fits into the conditional jump as an
590-
// immediate value and it's properly aligned, then we can
591-
// use the b.cond instruction directly. Otherwise, we need
592-
// to load the address into a register and use the branch
593-
// register instruction.
594-
if bcond_offset_fits_bits(offset) {
595-
bcond(cb, CONDITION, A64Opnd::new_imm(dst_addr - src_addr));
572+
let num_insns = if bcond_offset_fits_bits(offset) {
573+
// If the jump offset fits into the conditional jump as
574+
// an immediate value and it's properly aligned, then we
575+
// can use the b.cond instruction directly.
576+
bcond(cb, CONDITION, A64Opnd::new_imm(offset));
577+
578+
// Here we're going to return 1 because we've only
579+
// written out 1 instruction.
580+
1
596581
} else {
597-
// If the condition is met, then we'll skip past the
598-
// next instruction, put the address in a register, and
599-
// jump to it.
600-
bcond(cb, CONDITION, A64Opnd::new_imm(8));
601-
602-
// If the offset fits into a direct jump, then we'll use
603-
// that and the number of instructions will be shorter.
604-
// Otherwise we'll use the branch register instruction.
605-
if b_offset_fits_bits(offset) {
606-
// If we get to this instruction, then the condition
607-
// wasn't met, in which case we'll jump past the
608-
// next instruction that performs the direct jump.
609-
b(cb, A64Opnd::new_imm(1));
610-
611-
// Here we'll perform the direct jump to the target.
612-
let offset = dst_addr - cb.get_write_ptr().into_i64() + 4;
613-
b(cb, A64Opnd::new_imm(offset / 4));
614-
} else {
615-
// If we get to this instruction, then the condition
616-
// wasn't met, in which case we'll jump past the
617-
// next instruction that perform the direct jump.
618-
let value = dst_addr as u64;
619-
620-
b(cb, A64Opnd::new_imm(emit_load_size(value).into()));
621-
emit_load_value(cb, Assembler::SCRATCH0, value);
622-
br(cb, Assembler::SCRATCH0);
623-
}
624-
}
625-
*/
582+
// Otherwise, we need to load the address into a
583+
// register and use the branch register instruction.
584+
let dst_addr = dst_ptr.into_u64();
585+
let load_insns: i64 = emit_load_size(dst_addr).into();
586+
587+
// We're going to write out the inverse condition so
588+
// that if it doesn't match it will skip over the
589+
// instructions used for branching.
590+
bcond(cb, Condition::inverse(CONDITION), A64Opnd::new_imm((load_insns + 2) * 4));
591+
emit_load_value(cb, Assembler::SCRATCH0, dst_addr);
592+
br(cb, Assembler::SCRATCH0);
593+
594+
// Here we'll return the number of instructions that it
595+
// took to write out the destination address + 1 for the
596+
// b.cond and 1 for the br.
597+
load_insns + 2
598+
};
599+
600+
// We need to make sure we have at least 6 instructions for
601+
// every kind of jump for invalidation purposes, so we're
602+
// going to write out padding nop instructions here.
603+
for _ in num_insns..6 { nop(cb); }
626604
},
627605
Target::Label(label_idx) => {
628606
// Here we're going to save enough space for ourselves and
@@ -904,10 +882,10 @@ impl Assembler
904882
_ => unreachable!()
905883
};
906884
},
907-
Insn::Je(target) => {
885+
Insn::Je(target) | Insn::Jz(target) => {
908886
emit_conditional_jump::<{Condition::EQ}>(cb, *target);
909887
},
910-
Insn::Jne(target) => {
888+
Insn::Jne(target) | Insn::Jnz(target) => {
911889
emit_conditional_jump::<{Condition::NE}>(cb, *target);
912890
},
913891
Insn::Jl(target) => {
@@ -916,12 +894,6 @@ impl Assembler
916894
Insn::Jbe(target) => {
917895
emit_conditional_jump::<{Condition::LS}>(cb, *target);
918896
},
919-
Insn::Jz(target) => {
920-
emit_conditional_jump::<{Condition::EQ}>(cb, *target);
921-
},
922-
Insn::Jnz(target) => {
923-
emit_conditional_jump::<{Condition::NE}>(cb, *target);
924-
},
925897
Insn::Jo(target) => {
926898
emit_conditional_jump::<{Condition::VS}>(cb, *target);
927899
},
@@ -1053,6 +1025,28 @@ mod tests {
10531025
asm.compile_with_num_regs(&mut cb, 0);
10541026
}
10551027

1028+
#[test]
1029+
fn test_emit_je_fits_into_bcond() {
1030+
let (mut asm, mut cb) = setup_asm();
1031+
1032+
let offset = 80;
1033+
let target: CodePtr = ((cb.get_write_ptr().into_u64() + offset) as *mut u8).into();
1034+
1035+
asm.je(Target::CodePtr(target));
1036+
asm.compile_with_num_regs(&mut cb, 0);
1037+
}
1038+
1039+
#[test]
1040+
fn test_emit_je_does_not_fit_into_bcond() {
1041+
let (mut asm, mut cb) = setup_asm();
1042+
1043+
let offset = 1 << 21;
1044+
let target: CodePtr = ((cb.get_write_ptr().into_u64() + offset) as *mut u8).into();
1045+
1046+
asm.je(Target::CodePtr(target));
1047+
asm.compile_with_num_regs(&mut cb, 0);
1048+
}
1049+
10561050
#[test]
10571051
fn test_emit_lea_label() {
10581052
let (mut asm, mut cb) = setup_asm();

0 commit comments

Comments
 (0)