From d8afaee75a445ee6af7a6da9eed06a00e03ae14a Mon Sep 17 00:00:00 2001 From: Laurence Tratt Date: Thu, 21 Nov 2024 10:08:23 +0000 Subject: [PATCH] Deal with input register aliasing. Currently our register allocator assumes that SSA variables can't alias. This is true for the parts under our control -- but isn't true of the inputs to a trace! For example both `%0` and `%1` might come from `RAX`. This caused a panic in the register allocator. I'd noted the problem in https://github.com/ykjit/yk/issues/1449 but thought it happened rarely. Edd noticed that it happens in other, more common situations (including as a result of common subexpression elimination), including in a yklua benchmark, and created a test for it. This commit is a "quick fix" (it's not pretty, or efficient): when we spot aliasing, we spill aliased variables onto the stack, implicitly dealiasing them. In other words, we avoid having to think about aliasing in the main bulk of the register allocator at the expense of having to spill (and potentially unspill) values unnecessarily. Co-authored-by: Edd Barrett --- .../compile/jitc_yk/codegen/x64/lsregalloc.rs | 30 ++++++++++++++---- ykrt/src/compile/jitc_yk/codegen/x64/mod.rs | 31 ++++++++++++++++++- 2 files changed, 54 insertions(+), 7 deletions(-) diff --git a/ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs b/ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs index 529d5c6fd..bf9d8aa58 100644 --- a/ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs +++ b/ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs @@ -252,12 +252,30 @@ impl<'a> LSRegAlloc<'a> { /// The parts of the register allocator needed for general purpose registers. impl LSRegAlloc<'_> { - /// Forcibly assign the machine register `reg`, which must be in the [RegState::Empty] state, - /// to the value produced by instruction `iidx`. - pub(crate) fn force_assign_inst_gp_reg(&mut self, iidx: InstIdx, reg: Rq) { - debug_assert!(!self.gp_regset.is_set(reg)); - self.gp_regset.set(reg); - self.gp_reg_states[usize::from(reg.code())] = RegState::FromInst(iidx); + /// Forcibly assign the machine register `reg` to the value produced by instruction `iidx`. + /// Note that if this register is already used, a spill will be generated instead. + pub(crate) fn force_assign_inst_gp_reg(&mut self, asm: &mut Assembler, iidx: InstIdx, reg: Rq) { + if self.gp_regset.is_set(reg) { + debug_assert_eq!(self.spills[usize::from(iidx)], SpillState::Empty); + // Input values alias to a single register. To avoid the rest of the register allocator + // having to think about this, we "dealias" the values by spilling. + let inst = self.m.inst_no_copies(iidx); + let size = inst.def_byte_size(self.m); + self.stack.align(size); // FIXME + let frame_off = self.stack.grow(size); + let off = i32::try_from(frame_off).unwrap(); + match size { + 1 => dynasm!(asm; mov BYTE [rbp - off], Rb(reg.code())), + 2 => dynasm!(asm; mov WORD [rbp - off], Rw(reg.code())), + 4 => dynasm!(asm; mov DWORD [rbp - off], Rd(reg.code())), + 8 => dynasm!(asm; mov QWORD [rbp - off], Rq(reg.code())), + _ => unreachable!(), + } + self.spills[usize::from(iidx)] = SpillState::Stack(off); + } else { + self.gp_regset.set(reg); + self.gp_reg_states[usize::from(reg.code())] = RegState::FromInst(iidx); + } } /// Forcibly assign the floating point register `reg`, which must be in the [RegState::Empty] state, diff --git a/ykrt/src/compile/jitc_yk/codegen/x64/mod.rs b/ykrt/src/compile/jitc_yk/codegen/x64/mod.rs index ffca2819c..78ff3ec4f 100644 --- a/ykrt/src/compile/jitc_yk/codegen/x64/mod.rs +++ b/ykrt/src/compile/jitc_yk/codegen/x64/mod.rs @@ -1154,7 +1154,7 @@ impl<'a> Assemble<'a> { debug_assert!(size <= REG64_BYTESIZE); match m { VarLocation::Register(reg_alloc::Register::GP(reg)) => { - self.ra.force_assign_inst_gp_reg(iidx, reg); + self.ra.force_assign_inst_gp_reg(&mut self.asm, iidx, reg); } VarLocation::Register(reg_alloc::Register::FP(reg)) => { self.ra.force_assign_inst_fp_reg(iidx, reg); @@ -4381,4 +4381,33 @@ mod tests { ", ); } + + #[test] + fn cg_aliasing_loadtis() { + let mut m = jit_ir::Module::new(0, 0).unwrap(); + + // Create two trace inputs whose locations alias. + let loc = yksmp::Location::Register(13, 1, 0, [].into()); + m.push_tiloc(loc); + let ti_inst = jit_ir::LoadTraceInputInst::new(0, m.int8_tyidx()); + let op1 = m.push_and_make_operand(ti_inst.clone().into()).unwrap(); + let op2 = m.push_and_make_operand(ti_inst.into()).unwrap(); + + let add_inst = jit_ir::BinOpInst::new(op1, jit_ir::BinOp::Add, op2); + m.push(add_inst.into()).unwrap(); + + let mt = MT::new().unwrap(); + let hl = HotLocation { + kind: HotLocationKind::Tracing, + tracecompilation_errors: 0, + }; + + Assemble::new(&m, None, None) + .unwrap() + .codegen(mt, Arc::new(Mutex::new(hl)), None) + .unwrap() + .as_any() + .downcast::() + .unwrap(); + } }