Skip to content

Commit f7e73ba

Browse files
39bytesXrXr
39bytes
authored and
XrXr
committed
ZJIT: A64: Avoid gaps in the stack when preserving registers for calls
Previously, we used a `str x, [sp, #-0x10]!` for each value, which left an 8-byte gap. Use STP to store a pair at a time instead.
1 parent e24b528 commit f7e73ba

4 files changed

Lines changed: 301 additions & 10 deletions

File tree

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

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1420,12 +1420,20 @@ impl Assembler {
14201420
Insn::CPush(opnd) => {
14211421
emit_push(cb, opnd.into());
14221422
},
1423+
Insn::CPushPair(opnd0, opnd1) => {
1424+
// Second operand ends up at the lower stack address
1425+
stp_pre(cb, opnd1.into(), opnd0.into(), A64Opnd::new_mem(64, C_SP_REG, -C_SP_STEP));
1426+
},
14231427
Insn::CPop { out } => {
14241428
emit_pop(cb, out.into());
14251429
},
14261430
Insn::CPopInto(opnd) => {
14271431
emit_pop(cb, opnd.into());
14281432
},
1433+
Insn::CPopPairInto(opnd0, opnd1) => {
1434+
// First operand is popped from the lower stack address
1435+
ldp_post(cb, opnd0.into(), opnd1.into(), A64Opnd::new_mem(64, C_SP_REG, C_SP_STEP));
1436+
},
14291437
Insn::CCall { fptr, .. } => {
14301438
match fptr {
14311439
Opnd::UImm(fptr) => {
@@ -2662,6 +2670,76 @@ mod tests {
26622670
assert_snapshot!(cb.hexdump(), @"ef0302aae20303aae3030faaef0300aae00301aae1030faa100080d200023fd6");
26632671
}
26642672

2673+
#[test]
2674+
fn test_ccall_register_preservation_even() {
2675+
let (mut asm, mut cb) = setup_asm();
2676+
2677+
let v0 = asm.load(1.into());
2678+
let v1 = asm.load(2.into());
2679+
let v2 = asm.load(3.into());
2680+
let v3 = asm.load(4.into());
2681+
asm.ccall(0 as _, vec![]);
2682+
_ = asm.add(v0, v1);
2683+
_ = asm.add(v2, v3);
2684+
2685+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
2686+
2687+
assert_disasm_snapshot!(cb.disasm(), @"
2688+
0x0: mov x0, #1
2689+
0x4: mov x1, #2
2690+
0x8: mov x2, #3
2691+
0xc: mov x3, #4
2692+
0x10: mov x4, x0
2693+
0x14: stp x2, x1, [sp, #-0x10]!
2694+
0x18: stp x4, x3, [sp, #-0x10]!
2695+
0x1c: mov x16, #0
2696+
0x20: blr x16
2697+
0x24: ldp x4, x3, [sp], #0x10
2698+
0x28: ldp x2, x1, [sp], #0x10
2699+
0x2c: adds x4, x4, x1
2700+
0x30: adds x2, x2, x3
2701+
");
2702+
assert_snapshot!(cb.hexdump(), @"200080d2410080d2620080d2830080d2e40300aae207bfa9e40fbfa9100080d200023fd6e40fc1a8e207c1a8840001ab420003ab");
2703+
}
2704+
2705+
#[test]
2706+
fn test_ccall_register_preservation_odd() {
2707+
let (mut asm, mut cb) = setup_asm();
2708+
2709+
let v0 = asm.load(1.into());
2710+
let v1 = asm.load(2.into());
2711+
let v2 = asm.load(3.into());
2712+
let v3 = asm.load(4.into());
2713+
let v4 = asm.load(5.into());
2714+
asm.ccall(0 as _, vec![]);
2715+
_ = asm.add(v0, v1);
2716+
_ = asm.add(v2, v3);
2717+
_ = asm.add(v2, v4);
2718+
2719+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
2720+
2721+
assert_disasm_snapshot!(cb.disasm(), @"
2722+
0x0: mov x0, #1
2723+
0x4: mov x1, #2
2724+
0x8: mov x2, #3
2725+
0xc: mov x3, #4
2726+
0x10: mov x4, #5
2727+
0x14: mov x5, x0
2728+
0x18: stp x2, x1, [sp, #-0x10]!
2729+
0x1c: stp x4, x3, [sp, #-0x10]!
2730+
0x20: str x5, [sp, #-0x10]!
2731+
0x24: mov x16, #0
2732+
0x28: blr x16
2733+
0x2c: ldr x5, [sp], #0x10
2734+
0x30: ldp x4, x3, [sp], #0x10
2735+
0x34: ldp x2, x1, [sp], #0x10
2736+
0x38: adds x5, x5, x1
2737+
0x3c: adds x0, x2, x3
2738+
0x40: adds x2, x2, x4
2739+
");
2740+
assert_snapshot!(cb.hexdump(), @"200080d2410080d2620080d2830080d2a40080d2e50300aae207bfa9e40fbfa9e50f1ff8100080d200023fd6e50741f8e40fc1a8e207c1a8a50001ab400003ab420004ab");
2741+
}
2742+
26652743
#[test]
26662744
fn test_ccall_resolve_parallel_moves_large_cycle() {
26672745
let (mut asm, mut cb) = setup_asm();
@@ -2685,6 +2763,38 @@ mod tests {
26852763
assert_snapshot!(cb.hexdump(), @"ef0300aae00301aae10302aae2030faa100080d200023fd6");
26862764
}
26872765

2766+
#[test]
2767+
fn test_cpush_pair() {
2768+
let (mut asm, mut cb) = setup_asm();
2769+
let v0 = asm.load(1.into());
2770+
let v1 = asm.load(2.into());
2771+
asm.cpush_pair(v0, v1);
2772+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
2773+
2774+
assert_disasm_snapshot!(cb.disasm(), @"
2775+
0x0: mov x0, #1
2776+
0x4: mov x1, #2
2777+
0x8: stp x1, x0, [sp, #-0x10]!
2778+
");
2779+
assert_snapshot!(cb.hexdump(), @"200080d2410080d2e103bfa9");
2780+
}
2781+
2782+
#[test]
2783+
fn test_cpop_pair_into() {
2784+
let (mut asm, mut cb) = setup_asm();
2785+
let v0 = asm.load(1.into());
2786+
let v1 = asm.load(2.into());
2787+
asm.cpop_pair_into(v0, v1);
2788+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
2789+
2790+
assert_disasm_snapshot!(cb.disasm(), @"
2791+
0x0: mov x0, #1
2792+
0x4: mov x1, #2
2793+
0x8: ldp x0, x1, [sp], #0x10
2794+
");
2795+
assert_snapshot!(cb.hexdump(), @"200080d2410080d2e007c1a8");
2796+
}
2797+
26882798
#[test]
26892799
fn test_split_spilled_lshift() {
26902800
let (mut asm, mut cb) = setup_asm();

‎zjit/src/backend/lir.rs‎

Lines changed: 48 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -380,9 +380,17 @@ pub enum Insn {
380380
/// Pop a register from the C stack and store it into another register
381381
CPopInto(Opnd),
382382

383+
/// Pop a pair of registers from the C stack and store it into a pair of registers.
384+
/// The registers are popped from left to right.
385+
CPopPairInto(Opnd, Opnd),
386+
383387
/// Push a register onto the C stack
384388
CPush(Opnd),
385389

390+
/// Push a pair of registers onto the C stack.
391+
/// The registers are pushed from left to right.
392+
CPushPair(Opnd, Opnd),
393+
386394
// C function call with N arguments (variadic)
387395
CCall {
388396
opnds: Vec<Opnd>,
@@ -609,7 +617,9 @@ impl Insn {
609617
Insn::Cmp { .. } => "Cmp",
610618
Insn::CPop { .. } => "CPop",
611619
Insn::CPopInto(_) => "CPopInto",
620+
Insn::CPopPairInto(_, _) => "CPopPairInto",
612621
Insn::CPush(_) => "CPush",
622+
Insn::CPushPair(_, _) => "CPushPair",
613623
Insn::CCall { .. } => "CCall",
614624
Insn::CRet(_) => "CRet",
615625
Insn::CSelE { .. } => "CSelE",
@@ -865,6 +875,8 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
865875
},
866876
Insn::Add { left: opnd0, right: opnd1, .. } |
867877
Insn::And { left: opnd0, right: opnd1, .. } |
878+
Insn::CPushPair(opnd0, opnd1) |
879+
Insn::CPopPairInto(opnd0, opnd1) |
868880
Insn::Cmp { left: opnd0, right: opnd1 } |
869881
Insn::CSelE { truthy: opnd0, falsy: opnd1, .. } |
870882
Insn::CSelG { truthy: opnd0, falsy: opnd1, .. } |
@@ -1034,6 +1046,8 @@ impl<'a> InsnOpndMutIterator<'a> {
10341046
},
10351047
Insn::Add { left: opnd0, right: opnd1, .. } |
10361048
Insn::And { left: opnd0, right: opnd1, .. } |
1049+
Insn::CPushPair(opnd0, opnd1) |
1050+
Insn::CPopPairInto(opnd0, opnd1) |
10371051
Insn::Cmp { left: opnd0, right: opnd1 } |
10381052
Insn::CSelE { truthy: opnd0, falsy: opnd1, .. } |
10391053
Insn::CSelG { truthy: opnd0, falsy: opnd1, .. } |
@@ -1592,9 +1606,19 @@ impl Assembler
15921606
saved_regs = pool.live_regs();
15931607

15941608
// Save live registers
1595-
for &(reg, _) in saved_regs.iter() {
1596-
asm.cpush(Opnd::Reg(reg));
1597-
pool.dealloc_opnd(&Opnd::Reg(reg));
1609+
for pair in saved_regs.chunks(2) {
1610+
match *pair {
1611+
[(reg0, _), (reg1, _)] => {
1612+
asm.cpush_pair(Opnd::Reg(reg0), Opnd::Reg(reg1));
1613+
pool.dealloc_opnd(&Opnd::Reg(reg0));
1614+
pool.dealloc_opnd(&Opnd::Reg(reg1));
1615+
}
1616+
[(reg, _)] => {
1617+
asm.cpush(Opnd::Reg(reg));
1618+
pool.dealloc_opnd(&Opnd::Reg(reg));
1619+
}
1620+
_ => unreachable!("chunks(2)")
1621+
}
15981622
}
15991623
// On x86_64, maintain 16-byte stack alignment
16001624
if cfg!(target_arch = "x86_64") && saved_regs.len() % 2 == 1 {
@@ -1725,9 +1749,19 @@ impl Assembler
17251749
asm.cpop_into(Opnd::Reg(saved_regs.last().unwrap().0));
17261750
}
17271751
// Restore saved registers
1728-
for &(reg, vreg_idx) in saved_regs.iter().rev() {
1729-
asm.cpop_into(Opnd::Reg(reg));
1730-
pool.take_reg(&reg, vreg_idx);
1752+
for pair in saved_regs.chunks(2).rev() {
1753+
match *pair {
1754+
[(reg, vreg_idx)] => {
1755+
asm.cpop_into(Opnd::Reg(reg));
1756+
pool.take_reg(&reg, vreg_idx);
1757+
}
1758+
[(reg0, vreg_idx0), (reg1, vreg_idx1)] => {
1759+
asm.cpop_pair_into(Opnd::Reg(reg1), Opnd::Reg(reg0));
1760+
pool.take_reg(&reg1, vreg_idx1);
1761+
pool.take_reg(&reg0, vreg_idx0);
1762+
}
1763+
_ => unreachable!("chunks(2)")
1764+
}
17311765
}
17321766
saved_regs.clear();
17331767
}
@@ -2125,10 +2159,18 @@ impl Assembler {
21252159
self.push_insn(Insn::CPopInto(opnd));
21262160
}
21272161

2162+
pub fn cpop_pair_into(&mut self, opnd0: Opnd, opnd1: Opnd) {
2163+
self.push_insn(Insn::CPopPairInto(opnd0, opnd1));
2164+
}
2165+
21282166
pub fn cpush(&mut self, opnd: Opnd) {
21292167
self.push_insn(Insn::CPush(opnd));
21302168
}
21312169

2170+
pub fn cpush_pair(&mut self, opnd0: Opnd, opnd1: Opnd) {
2171+
self.push_insn(Insn::CPushPair(opnd0, opnd1));
2172+
}
2173+
21322174
pub fn cret(&mut self, opnd: Opnd) {
21332175
self.push_insn(Insn::CRet(opnd));
21342176
}

‎zjit/src/backend/x86_64/mod.rs‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -851,12 +851,20 @@ impl Assembler {
851851
Insn::CPush(opnd) => {
852852
push(cb, opnd.into());
853853
},
854+
Insn::CPushPair(opnd0, opnd1) => {
855+
push(cb, opnd0.into());
856+
push(cb, opnd1.into());
857+
},
854858
Insn::CPop { out } => {
855859
pop(cb, out.into());
856860
},
857861
Insn::CPopInto(opnd) => {
858862
pop(cb, opnd.into());
859863
},
864+
Insn::CPopPairInto(opnd0, opnd1) => {
865+
pop(cb, opnd0.into());
866+
pop(cb, opnd1.into());
867+
},
860868

861869
// C function call
862870
Insn::CCall { fptr, .. } => {
@@ -1648,6 +1656,119 @@ mod tests {
16481656
assert_snapshot!(cb.hexdump(), @"b801000000b902000000ba030000004889c74889ce4989cb4889d14c89dab800000000ffd0");
16491657
}
16501658

1659+
#[test]
1660+
fn test_ccall_register_preservation_even() {
1661+
let (mut asm, mut cb) = setup_asm();
1662+
1663+
let v0 = asm.load(1.into());
1664+
let v1 = asm.load(2.into());
1665+
let v2 = asm.load(3.into());
1666+
let v3 = asm.load(4.into());
1667+
asm.ccall(0 as _, vec![]);
1668+
_ = asm.add(v0, v1);
1669+
_ = asm.add(v2, v3);
1670+
1671+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
1672+
1673+
assert_disasm_snapshot!(cb.disasm(), @"
1674+
0x0: mov edi, 1
1675+
0x5: mov esi, 2
1676+
0xa: mov edx, 3
1677+
0xf: mov ecx, 4
1678+
0x14: push rdi
1679+
0x15: push rsi
1680+
0x16: push rdx
1681+
0x17: push rcx
1682+
0x18: mov eax, 0
1683+
0x1d: call rax
1684+
0x1f: pop rcx
1685+
0x20: pop rdx
1686+
0x21: pop rsi
1687+
0x22: pop rdi
1688+
0x23: add rdi, rsi
1689+
0x26: add rdx, rcx
1690+
");
1691+
assert_snapshot!(cb.hexdump(), @"bf01000000be02000000ba03000000b90400000057565251b800000000ffd0595a5e5f4801f74801ca");
1692+
}
1693+
1694+
#[test]
1695+
fn test_ccall_register_preservation_odd() {
1696+
let (mut asm, mut cb) = setup_asm();
1697+
1698+
let v0 = asm.load(1.into());
1699+
let v1 = asm.load(2.into());
1700+
let v2 = asm.load(3.into());
1701+
let v3 = asm.load(4.into());
1702+
let v4 = asm.load(5.into());
1703+
asm.ccall(0 as _, vec![]);
1704+
_ = asm.add(v0, v1);
1705+
_ = asm.add(v2, v3);
1706+
_ = asm.add(v2, v4);
1707+
1708+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
1709+
1710+
assert_disasm_snapshot!(cb.disasm(), @"
1711+
0x0: mov edi, 1
1712+
0x5: mov esi, 2
1713+
0xa: mov edx, 3
1714+
0xf: mov ecx, 4
1715+
0x14: mov r8d, 5
1716+
0x1a: push rdi
1717+
0x1b: push rsi
1718+
0x1c: push rdx
1719+
0x1d: push rcx
1720+
0x1e: push r8
1721+
0x20: push r8
1722+
0x22: mov eax, 0
1723+
0x27: call rax
1724+
0x29: pop r8
1725+
0x2b: pop r8
1726+
0x2d: pop rcx
1727+
0x2e: pop rdx
1728+
0x2f: pop rsi
1729+
0x30: pop rdi
1730+
0x31: add rdi, rsi
1731+
0x34: mov rdi, rdx
1732+
0x37: add rdi, rcx
1733+
0x3a: add rdx, r8
1734+
");
1735+
assert_snapshot!(cb.hexdump(), @"bf01000000be02000000ba03000000b90400000041b8050000005756525141504150b800000000ffd041584158595a5e5f4801f74889d74801cf4c01c2");
1736+
}
1737+
1738+
#[test]
1739+
fn test_cpush_pair() {
1740+
let (mut asm, mut cb) = setup_asm();
1741+
let v0 = asm.load(1.into());
1742+
let v1 = asm.load(2.into());
1743+
asm.cpush_pair(v0, v1);
1744+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
1745+
1746+
assert_disasm_snapshot!(cb.disasm(), @"
1747+
0x0: mov edi, 1
1748+
0x5: mov esi, 2
1749+
0xa: push rdi
1750+
0xb: push rsi
1751+
");
1752+
assert_snapshot!(cb.hexdump(), @"bf01000000be020000005756");
1753+
}
1754+
1755+
#[test]
1756+
fn test_cpop_pair_into() {
1757+
let (mut asm, mut cb) = setup_asm();
1758+
let v0 = asm.load(1.into());
1759+
let v1 = asm.load(2.into());
1760+
asm.cpop_pair_into(v0, v1);
1761+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
1762+
1763+
assert_disasm_snapshot!(cb.disasm(), @"
1764+
0x0: mov edi, 1
1765+
0x5: mov esi, 2
1766+
0xa: pop rdi
1767+
0xb: pop rsi
1768+
");
1769+
assert_snapshot!(cb.hexdump(), @"bf01000000be020000005f5e");
1770+
}
1771+
16511772
#[test]
16521773
fn test_cmov_mem() {
16531774
let (mut asm, mut cb) = setup_asm();

0 commit comments

Comments
 (0)