Skip to content

Commit f2715af

Browse files
tekknolagitekknolagi
authored
ZJIT: Allocate register for VRegs that begin and end at the same index (#14270)
If the LiveRange looks like (idx, idx), we will currently not allocate a register. This change allocates a register and then immediately deallocates it. Fix Shopify#614
1 parent 7ac16ef commit f2715af

File tree

3 files changed

+116
-9
lines changed

3 files changed

+116
-9
lines changed

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

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2059,4 +2059,93 @@ mod tests {
20592059
0x4: adds x1, x0, #1
20602060
"});
20612061
}
2062+
2063+
#[test]
2064+
fn test_reorder_c_args_no_cycle() {
2065+
crate::options::rb_zjit_prepare_options();
2066+
let (mut asm, mut cb) = setup_asm();
2067+
2068+
asm.ccall(0 as _, vec![
2069+
C_ARG_OPNDS[0], // mov x0, x0 (optimized away)
2070+
C_ARG_OPNDS[1], // mov x1, x1 (optimized away)
2071+
]);
2072+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
2073+
2074+
assert_disasm!(cb, "100080d200023fd6", {"
2075+
0x0: mov x16, #0
2076+
0x4: blr x16
2077+
"});
2078+
}
2079+
2080+
#[test]
2081+
fn test_reorder_c_args_single_cycle() {
2082+
crate::options::rb_zjit_prepare_options();
2083+
let (mut asm, mut cb) = setup_asm();
2084+
2085+
// x0 and x1 form a cycle
2086+
asm.ccall(0 as _, vec![
2087+
C_ARG_OPNDS[1], // mov x0, x1
2088+
C_ARG_OPNDS[0], // mov x1, x0
2089+
C_ARG_OPNDS[2], // mov x2, x2 (optimized away)
2090+
]);
2091+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
2092+
2093+
assert_disasm!(cb, "f00300aae00301aae10310aa100080d200023fd6", {"
2094+
0x0: mov x16, x0
2095+
0x4: mov x0, x1
2096+
0x8: mov x1, x16
2097+
0xc: mov x16, #0
2098+
0x10: blr x16
2099+
"});
2100+
}
2101+
2102+
#[test]
2103+
fn test_reorder_c_args_two_cycles() {
2104+
crate::options::rb_zjit_prepare_options();
2105+
let (mut asm, mut cb) = setup_asm();
2106+
2107+
// x0 and x1 form a cycle, and x2 and rcx form another cycle
2108+
asm.ccall(0 as _, vec![
2109+
C_ARG_OPNDS[1], // mov x0, x1
2110+
C_ARG_OPNDS[0], // mov x1, x0
2111+
C_ARG_OPNDS[3], // mov x2, rcx
2112+
C_ARG_OPNDS[2], // mov rcx, x2
2113+
]);
2114+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
2115+
2116+
assert_disasm!(cb, "f00302aae20303aae30310aaf00300aae00301aae10310aa100080d200023fd6", {"
2117+
0x0: mov x16, x2
2118+
0x4: mov x2, x3
2119+
0x8: mov x3, x16
2120+
0xc: mov x16, x0
2121+
0x10: mov x0, x1
2122+
0x14: mov x1, x16
2123+
0x18: mov x16, #0
2124+
0x1c: blr x16
2125+
"});
2126+
}
2127+
2128+
#[test]
2129+
fn test_reorder_c_args_large_cycle() {
2130+
crate::options::rb_zjit_prepare_options();
2131+
let (mut asm, mut cb) = setup_asm();
2132+
2133+
// x0, x1, and x2 form a cycle
2134+
asm.ccall(0 as _, vec![
2135+
C_ARG_OPNDS[1], // mov x0, x1
2136+
C_ARG_OPNDS[2], // mov x1, x2
2137+
C_ARG_OPNDS[0], // mov x2, x0
2138+
]);
2139+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
2140+
2141+
assert_disasm!(cb, "f00300aae00301aae10302aae20310aa100080d200023fd6", {"
2142+
0x0: mov x16, x0
2143+
0x4: mov x0, x1
2144+
0x8: mov x1, x2
2145+
0xc: mov x2, x16
2146+
0x10: mov x16, #0
2147+
0x14: blr x16
2148+
"});
2149+
}
2150+
20622151
}

‎zjit/src/backend/lir.rs

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1202,7 +1202,7 @@ impl Assembler
12021202
/// Append an instruction onto the current list of instructions and update
12031203
/// the live ranges of any instructions whose outputs are being used as
12041204
/// operands to this instruction.
1205-
pub fn push_insn(&mut self, mut insn: Insn) {
1205+
pub fn push_insn(&mut self, insn: Insn) {
12061206
// Index of this instruction
12071207
let insn_idx = self.insns.len();
12081208

@@ -1214,7 +1214,7 @@ impl Assembler
12141214
}
12151215

12161216
// If we find any VReg from previous instructions, extend the live range to insn_idx
1217-
let mut opnd_iter = insn.opnd_iter_mut();
1217+
let mut opnd_iter = insn.opnd_iter();
12181218
while let Some(opnd) = opnd_iter.next() {
12191219
match *opnd {
12201220
Opnd::VReg { idx, .. } |
@@ -1380,13 +1380,15 @@ impl Assembler
13801380
}
13811381
}
13821382

1383-
// If the output VReg of this instruction is used by another instruction,
1384-
// we need to allocate a register to it
1383+
// Allocate a register for the output operand if it exists
13851384
let vreg_idx = match insn.out_opnd() {
13861385
Some(Opnd::VReg { idx, .. }) => Some(*idx),
13871386
_ => None,
13881387
};
1389-
if vreg_idx.is_some() && live_ranges[vreg_idx.unwrap()].end() != index {
1388+
if vreg_idx.is_some() {
1389+
if live_ranges[vreg_idx.unwrap()].end() == index {
1390+
debug!("Allocating a register for VReg({}) at instruction index {} even though it does not live past this index", vreg_idx.unwrap(), index);
1391+
}
13901392
// This is going to be the output operand that we will set on the
13911393
// instruction. CCall and LiveReg need to use a specific register.
13921394
let mut out_reg = match insn {
@@ -1466,6 +1468,18 @@ impl Assembler
14661468
}
14671469
}
14681470

1471+
// If we have an output that dies at its definition (it is unused), free up the
1472+
// register
1473+
if let Some(idx) = vreg_idx {
1474+
if live_ranges[idx].end() == index {
1475+
if let Some(reg) = reg_mapping[idx] {
1476+
pool.dealloc_reg(&reg);
1477+
} else {
1478+
unreachable!("no register allocated for insn {:?}", insn);
1479+
}
1480+
}
1481+
}
1482+
14691483
// Push instruction(s)
14701484
let is_ccall = matches!(insn, Insn::CCall { .. });
14711485
match insn {

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1247,13 +1247,14 @@ mod tests {
12471247

12481248
#[test]
12491249
fn test_reorder_c_args_no_cycle() {
1250+
crate::options::rb_zjit_prepare_options();
12501251
let (mut asm, mut cb) = setup_asm();
12511252

12521253
asm.ccall(0 as _, vec![
12531254
C_ARG_OPNDS[0], // mov rdi, rdi (optimized away)
12541255
C_ARG_OPNDS[1], // mov rsi, rsi (optimized away)
12551256
]);
1256-
asm.compile_with_num_regs(&mut cb, 0);
1257+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
12571258

12581259
assert_disasm!(cb, "b800000000ffd0", {"
12591260
0x0: mov eax, 0
@@ -1263,6 +1264,7 @@ mod tests {
12631264

12641265
#[test]
12651266
fn test_reorder_c_args_single_cycle() {
1267+
crate::options::rb_zjit_prepare_options();
12661268
let (mut asm, mut cb) = setup_asm();
12671269

12681270
// rdi and rsi form a cycle
@@ -1271,7 +1273,7 @@ mod tests {
12711273
C_ARG_OPNDS[0], // mov rsi, rdi
12721274
C_ARG_OPNDS[2], // mov rdx, rdx (optimized away)
12731275
]);
1274-
asm.compile_with_num_regs(&mut cb, 0);
1276+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
12751277

12761278
assert_disasm!(cb, "4989f34889fe4c89dfb800000000ffd0", {"
12771279
0x0: mov r11, rsi
@@ -1284,6 +1286,7 @@ mod tests {
12841286

12851287
#[test]
12861288
fn test_reorder_c_args_two_cycles() {
1289+
crate::options::rb_zjit_prepare_options();
12871290
let (mut asm, mut cb) = setup_asm();
12881291

12891292
// rdi and rsi form a cycle, and rdx and rcx form another cycle
@@ -1293,7 +1296,7 @@ mod tests {
12931296
C_ARG_OPNDS[3], // mov rdx, rcx
12941297
C_ARG_OPNDS[2], // mov rcx, rdx
12951298
]);
1296-
asm.compile_with_num_regs(&mut cb, 0);
1299+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
12971300

12981301
assert_disasm!(cb, "4989f34889fe4c89df4989cb4889d14c89dab800000000ffd0", {"
12991302
0x0: mov r11, rsi
@@ -1309,6 +1312,7 @@ mod tests {
13091312

13101313
#[test]
13111314
fn test_reorder_c_args_large_cycle() {
1315+
crate::options::rb_zjit_prepare_options();
13121316
let (mut asm, mut cb) = setup_asm();
13131317

13141318
// rdi, rsi, and rdx form a cycle
@@ -1317,7 +1321,7 @@ mod tests {
13171321
C_ARG_OPNDS[2], // mov rsi, rdx
13181322
C_ARG_OPNDS[0], // mov rdx, rdi
13191323
]);
1320-
asm.compile_with_num_regs(&mut cb, 0);
1324+
asm.compile_with_num_regs(&mut cb, ALLOC_REGS.len());
13211325

13221326
assert_disasm!(cb, "4989f34889d64889fa4c89dfb800000000ffd0", {"
13231327
0x0: mov r11, rsi

0 commit comments

Comments
 (0)