From 29e0713a1272cb63f1e3cebfab85dec2424ead0f Mon Sep 17 00:00:00 2001 From: Kevin Newton Date: Thu, 25 Aug 2022 21:19:56 -0400 Subject: [PATCH] TBZ and TBNZ for AArch64 (https://github.com/Shopify/ruby/pull/434) --- yjit/src/asm/arm64/inst/mod.rs | 2 + yjit/src/asm/arm64/inst/test_bit.rs | 135 ++++++++++++++++++++++++++++ yjit/src/asm/arm64/mod.rs | 34 +++++++ 3 files changed, 171 insertions(+) create mode 100644 yjit/src/asm/arm64/inst/test_bit.rs diff --git a/yjit/src/asm/arm64/inst/mod.rs b/yjit/src/asm/arm64/inst/mod.rs index f4c27a5102..b3a77e73c9 100644 --- a/yjit/src/asm/arm64/inst/mod.rs +++ b/yjit/src/asm/arm64/inst/mod.rs @@ -22,6 +22,7 @@ mod reg_pair; mod sbfm; mod shift_imm; mod sys_reg; +mod test_bit; pub use atomic::Atomic; pub use branch::Branch; @@ -44,3 +45,4 @@ pub use reg_pair::RegisterPair; pub use sbfm::SBFM; pub use shift_imm::ShiftImm; pub use sys_reg::SysReg; +pub use test_bit::TestBit; diff --git a/yjit/src/asm/arm64/inst/test_bit.rs b/yjit/src/asm/arm64/inst/test_bit.rs new file mode 100644 index 0000000000..1961e65949 --- /dev/null +++ b/yjit/src/asm/arm64/inst/test_bit.rs @@ -0,0 +1,135 @@ +/// The upper bit of the bit number to test. +#[derive(Debug)] +enum B5 { + /// When the bit number is below 32. + B532 = 0, + + /// When the bit number is equal to or above 32. + B564 = 1 +} + +/// A convenience function so that we can convert the bit number directly into a +/// B5 variant. +impl From for B5 { + fn from(bit_num: u8) -> Self { + match bit_num { + 0..=31 => B5::B532, + 32..=63 => B5::B564, + _ => panic!("Invalid bit number: {}", bit_num) + } + } +} + +/// The operation to perform for this instruction. +enum Op { + /// The test bit zero operation. + TBZ = 0, + + /// The test bit not zero operation. + TBNZ = 1 +} + +/// The struct that represents an A64 test bit instruction that can be encoded. +/// +/// TBNZ/TBZ +/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+ +/// | 31 30 29 28 | 27 26 25 24 | 23 22 21 20 | 19 18 17 16 | 15 14 13 12 | 11 10 09 08 | 07 06 05 04 | 03 02 01 00 | +/// | 0 1 1 0 1 1 | +/// | b5 op b40............. imm14.......................................... rt.............. | +/// +-------------+-------------+-------------+-------------+-------------+-------------+-------------+-------------+ +/// +pub struct TestBit { + /// The number of the register to test. + rt: u8, + + /// The PC-relative offset to the target instruction in term of number of + /// instructions. + imm14: i16, + + /// The lower 5 bits of the bit number to be tested. + b40: u8, + + /// The operation to perform for this instruction. + op: Op, + + /// The upper bit of the bit number to test. + b5: B5 +} + +impl TestBit { + /// TBNZ + /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/TBNZ--Test-bit-and-Branch-if-Nonzero-?lang=en + pub fn tbnz(rt: u8, bit_num: u8, offset: i16) -> Self { + Self { rt, imm14: offset, b40: bit_num & 0b11111, op: Op::TBNZ, b5: bit_num.into() } + } + + /// TBZ + /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/TBZ--Test-bit-and-Branch-if-Zero-?lang=en + pub fn tbz(rt: u8, bit_num: u8, offset: i16) -> Self { + Self { rt, imm14: offset, b40: bit_num & 0b11111, op: Op::TBZ, b5: bit_num.into() } + } +} + +/// https://developer.arm.com/documentation/ddi0602/2022-03/Index-by-Encoding/Branches--Exception-Generating-and-System-instructions?lang=en +const FAMILY: u32 = 0b11011; + +impl From for u32 { + /// Convert an instruction into a 32-bit value. + fn from(inst: TestBit) -> Self { + let b40 = (inst.b40 & 0b11111) as u32; + let mut imm14 = (inst.imm14 & ((1 << 13) - 1)) as u32; + + if inst.imm14 < 0 { + imm14 |= (1 << 13); + } + + 0 + | ((inst.b5 as u32) << 31) + | (FAMILY << 25) + | ((inst.op as u32) << 24) + | (b40 << 19) + | (imm14 << 5) + | inst.rt as u32 + } +} + +impl From for [u8; 4] { + /// Convert an instruction into a 4 byte array. + fn from(inst: TestBit) -> [u8; 4] { + let result: u32 = inst.into(); + result.to_le_bytes() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_tbnz() { + let inst = TestBit::tbnz(0, 0, 0); + let result: u32 = inst.into(); + assert_eq!(0x37000000, result); + } + + #[test] + fn test_tbnz_negative() { + let inst = TestBit::tbnz(0, 0, -1); + let result: u32 = inst.into(); + assert_eq!(0x3707ffe0, result); + } + + #[test] + fn test_tbz() { + let inst = TestBit::tbz(0, 0, 0); + let result: u32 = inst.into(); + assert_eq!(0x36000000, result); + } + + #[test] + fn test_tbz_negative() { + let inst = TestBit::tbz(0, 0, -1); + let result: u32 = inst.into(); + assert_eq!(0x3607ffe0, result); + } +} diff --git a/yjit/src/asm/arm64/mod.rs b/yjit/src/asm/arm64/mod.rs index cf898d2b5a..a6aa8ffcbb 100644 --- a/yjit/src/asm/arm64/mod.rs +++ b/yjit/src/asm/arm64/mod.rs @@ -934,6 +934,30 @@ pub fn ret(cb: &mut CodeBlock, rn: A64Opnd) { cb.write_bytes(&bytes); } +/// TBNZ - test bit and branch if not zero +pub fn tbnz(cb: &mut CodeBlock, rt: A64Opnd, bit_num: A64Opnd, offset: A64Opnd) { + let bytes: [u8; 4] = match (rt, bit_num, offset) { + (A64Opnd::Reg(rt), A64Opnd::UImm(bit_num), A64Opnd::Imm(offset)) => { + TestBit::tbnz(rt.reg_no, bit_num.try_into().unwrap(), offset.try_into().unwrap()).into() + }, + _ => panic!("Invalid operand combination to tbnz instruction.") + }; + + cb.write_bytes(&bytes); +} + +/// TBZ - test bit and branch if zero +pub fn tbz(cb: &mut CodeBlock, rt: A64Opnd, bit_num: A64Opnd, offset: A64Opnd) { + let bytes: [u8; 4] = match (rt, bit_num, offset) { + (A64Opnd::Reg(rt), A64Opnd::UImm(bit_num), A64Opnd::Imm(offset)) => { + TestBit::tbz(rt.reg_no, bit_num.try_into().unwrap(), offset.try_into().unwrap()).into() + }, + _ => panic!("Invalid operand combination to tbz instruction.") + }; + + cb.write_bytes(&bytes); +} + /// TST - test the bits of a register against a mask, then update flags pub fn tst(cb: &mut CodeBlock, rn: A64Opnd, rm: A64Opnd) { let bytes: [u8; 4] = match (rn, rm) { @@ -1393,6 +1417,16 @@ mod tests { check_bytes("6a7d4093", |cb| sxtw(cb, X10, W11)); } + #[test] + fn test_tbnz() { + check_bytes("4a005037", |cb| tbnz(cb, X10, A64Opnd::UImm(10), A64Opnd::Imm(2))); + } + + #[test] + fn test_tbz() { + check_bytes("4a005036", |cb| tbz(cb, X10, A64Opnd::UImm(10), A64Opnd::Imm(2))); + } + #[test] fn test_tst_register() { check_bytes("1f0001ea", |cb| tst(cb, X0, X1));