Skip to content

Commit

Permalink
Unbreak YKD_OPT=0.
Browse files Browse the repository at this point in the history
When adding support for loop peeling, we forgot that loops are only
peeled when optimisations are turned on. When they were turned off, we
generated code with no backjump e.g.:

```
; %294: ptr = load %293
mov rsi, [rdi+0x588]
; header_end [...]
; Deopt ID for guard 0
push rsi
mov rsi, 0x00
```

Notice no `jmp` between `header_end` and `Deopt`. That meant that we
fell through to whichever guard's deopt happened to come first, with
confusing results. For example, before this commit the new test
`ykd_opt_off` would print:

```
10
9
8
7: 7
```

even though the code shows this is clearly "impossible":

```rust
if (i < 7)
  printf("7: ");
```

This commit not only fixes that, but brings a bit more order to three
kinds of traces we can compile: "header only" traces (most tests, and
when `YKD_OPT=0`); "header and body" traces (with loop peeling); and
"side traces". A new `TraceKind` captures these, and allows us to
simplify quite a bit of code in the x64 backend, which no longer has to
guess what kind of trace it has.

Perhaps inevitably this also showed up a few tests where we'd muddled
header and body.
  • Loading branch information
ltratt committed Dec 20, 2024
1 parent 40f5cad commit 3508eca
Show file tree
Hide file tree
Showing 6 changed files with 192 additions and 78 deletions.
47 changes: 47 additions & 0 deletions tests/c/ykd_opt_off.c
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
// Run-time:
// env-var: YKD_OPT=0
// env-var: YKD_SERIALISE_COMPILATION=1
// stdout:
// 10
// 9
// 8
// 7
// 7: 6
// 7: 5
// 7: 4
// 7: 4: 3
// 7: 4: 2
// 7: 4: 1


// Check that basic trace compilation works.

#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <yk.h>
#include <yk_testing.h>

int main(int argc, char **argv) {
YkMT *mt = yk_mt_new(NULL);
yk_mt_hot_threshold_set(mt, 0);
yk_mt_sidetrace_threshold_set(mt, 1);
YkLocation loc = yk_location_new();

int i = 10;
NOOPT_VAL(loc);
NOOPT_VAL(i);
while (i > 0) {
yk_mt_control_point(mt, &loc);
if (i < 7)
printf("7: ");
if (i < 4)
printf("4: ");
printf("%d\n", i);
i--;
}
yk_location_drop(loc);
yk_mt_shutdown(mt);
return (EXIT_SUCCESS);
}
160 changes: 91 additions & 69 deletions ykrt/src/compile/jitc_yk/codegen/x64/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
use super::{
super::{
int_signs::{SignExtend, Truncate},
jit_ir::{self, BinOp, FloatTy, Inst, InstIdx, Module, Operand, PtrAddInst, Ty},
jit_ir::{self, BinOp, FloatTy, Inst, InstIdx, Module, Operand, PtrAddInst, TraceKind, Ty},
CompilationError,
},
reg_alloc::{self, VarLocation},
Expand Down Expand Up @@ -519,7 +519,7 @@ impl<'a> Assemble<'a> {
jit_ir::Inst::Guard(i) => self.cg_guard(iidx, i),
jit_ir::Inst::TraceHeaderStart => self.cg_header_start(),
jit_ir::Inst::TraceHeaderEnd => {
self.cg_header_end();
self.cg_header_end(iidx);
in_header = false;
}
jit_ir::Inst::TraceBodyStart => self.cg_body_start(),
Expand Down Expand Up @@ -1791,12 +1791,11 @@ impl<'a> Assemble<'a> {
/// # Arguments
///
/// * `tgt_vars` - The target locations. If `None` use `self.loop_start_locs` instead.
fn write_jump_vars(&mut self, iidx: InstIdx, is_sidetrace: bool) {
let (tgt_vars, src_ops) = if is_sidetrace {
// Side-traces don't have a body and store these variables in `trace_header_end`.
(self.m.root_entry_vars(), self.m.trace_header_end())
} else {
(self.body_start_locs.as_slice(), self.m.trace_body_end())
fn write_jump_vars(&mut self, iidx: InstIdx) {
let (tgt_vars, src_ops) = match self.m.tracekind() {
TraceKind::HeaderOnly => (self.header_start_locs.as_slice(), self.m.trace_header_end()),
TraceKind::HeaderAndBody => (self.body_start_locs.as_slice(), self.m.trace_body_end()),
TraceKind::Sidetrace => (self.m.root_entry_vars(), self.m.trace_header_end()),
};
// If we pass in `None` use `self.loop_start_locs` instead. We need to do this since we
// can't pass in `&self.loop_start_locs` directly due to borrowing restrictions.
Expand Down Expand Up @@ -1889,13 +1888,14 @@ impl<'a> Assemble<'a> {
}

fn cg_body_end(&mut self, iidx: InstIdx) {
debug_assert_eq!(self.m.tracekind(), TraceKind::HeaderAndBody);
// Loop the JITted code if the `tloop_start` label is present (not relevant for IR created
// by a test or a side-trace).
let label = StaticLabel::global("tloop_start");
match self.asm.labels().resolve_static(&label) {
Ok(_) => {
// Found the label, emit a jump to it.
self.write_jump_vars(iidx, false);
self.write_jump_vars(iidx);
dynasm!(self.asm; jmp ->tloop_start);
}
Err(DynasmError::UnknownLabel(_)) => {
Expand All @@ -1917,9 +1917,10 @@ impl<'a> Assemble<'a> {
}

fn cg_sidetrace_end(&mut self, iidx: InstIdx, addr: *const libc::c_void) {
debug_assert_eq!(self.m.tracekind(), TraceKind::Sidetrace);
// The end of a side-trace. Map live variables of this side-trace to the entry variables of
// the root parent trace, then jump to it.
self.write_jump_vars(iidx, true);
self.write_jump_vars(iidx);
self.ra.align_stack(SYSV_CALL_STACK_ALIGN);

dynasm!(self.asm
Expand All @@ -1944,60 +1945,80 @@ impl<'a> Assemble<'a> {
};
self.header_start_locs.push(loc);
}
dynasm!(self.asm; ->reentry:);
match self.m.tracekind() {
TraceKind::HeaderOnly => {
dynasm!(self.asm; ->tloop_start:);
}
TraceKind::HeaderAndBody => {
dynasm!(self.asm; ->reentry:);
}
TraceKind::Sidetrace => todo!(),
}
self.prologue_offset = self.asm.offset();
}

fn cg_header_end(&mut self) {
// FIXME: This is a bit of a roundabout way of doing things. Especially, since it means
// that the [ParamInst]s in the trace body are just placeholders. While, since a recent
// change, the register allocator makes sure the values automatically end up in the
// [VarLocation]s expected by the loop start, this only works for registers right now. We
// can extend this to spill locations as well, but won't be able to do so for variables
// that have become constants during the trace header. So we will always have to either
// update the [ParamInst]s of the trace body, which isn't ideal since it requires the
// [Module] the be mutable. Or we do what we do below just for constants.
let mut varlocs = Vec::new();
for var in self.m.trace_header_end().iter() {
let varloc = self.op_to_var_location(var.unpack(self.m));
varlocs.push(varloc);
}
// Reset the register allocator before priming it with information about the trace body
// inputs.
self.ra.reset();
for (i, op) in self.m.trace_body_start().iter().enumerate() {
// By definition these can only be variables.
let iidx = match op.unpack(self.m) {
Operand::Var(iidx) => iidx,
_ => panic!(),
};
let varloc = varlocs[i];

// Write the varlocations from the head jump to the body start.
// FIXME: This is copied verbatim from `cg_param` and can be reused.
match varloc {
VarLocation::Register(reg_alloc::Register::GP(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);
}
VarLocation::Direct { frame_off, size: _ } => {
self.ra.force_assign_inst_direct(iidx, frame_off);
}
VarLocation::Stack { frame_off, size: _ } => {
self.ra
.force_assign_inst_indirect(iidx, i32::try_from(frame_off).unwrap());
fn cg_header_end(&mut self, iidx: InstIdx) {
match self.m.tracekind() {
TraceKind::HeaderOnly => {
self.write_jump_vars(iidx);
dynasm!(self.asm; jmp ->tloop_start);
}
TraceKind::HeaderAndBody => {
// FIXME: This is a bit of a roundabout way of doing things. Especially, since it means
// that the [ParamInst]s in the trace body are just placeholders. While, since a recent
// change, the register allocator makes sure the values automatically end up in the
// [VarLocation]s expected by the loop start, this only works for registers right now. We
// can extend this to spill locations as well, but won't be able to do so for variables
// that have become constants during the trace header. So we will always have to either
// update the [ParamInst]s of the trace body, which isn't ideal since it requires the
// [Module] the be mutable. Or we do what we do below just for constants.
let mut varlocs = Vec::new();
for var in self.m.trace_header_end().iter() {
let varloc = self.op_to_var_location(var.unpack(self.m));
varlocs.push(varloc);
}
VarLocation::ConstInt { bits, v } => {
self.ra.assign_const(iidx, bits, v);
// Reset the register allocator before priming it with information about the trace body
// inputs.
self.ra.reset();
for (i, op) in self.m.trace_body_start().iter().enumerate() {
// By definition these can only be variables.
let iidx = match op.unpack(self.m) {
Operand::Var(iidx) => iidx,
_ => panic!(),
};
let varloc = varlocs[i];

// Write the varlocations from the head jump to the body start.
// FIXME: This is copied verbatim from `cg_param` and can be reused.
match varloc {
VarLocation::Register(reg_alloc::Register::GP(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);
}
VarLocation::Direct { frame_off, size: _ } => {
self.ra.force_assign_inst_direct(iidx, frame_off);
}
VarLocation::Stack { frame_off, size: _ } => {
self.ra.force_assign_inst_indirect(
iidx,
i32::try_from(frame_off).unwrap(),
);
}
VarLocation::ConstInt { bits, v } => {
self.ra.assign_const(iidx, bits, v);
}
e => panic!("{:?}", e),
}
}
e => panic!("{:?}", e),
}
TraceKind::Sidetrace => todo!(),
}
}

fn cg_body_start(&mut self) {
debug_assert_eq!(self.m.tracekind(), TraceKind::HeaderAndBody);
debug_assert_eq!(self.body_start_locs.len(), 0);
// Remember the locations of the live variables at the beginning of the trace loop. When we
// loop back around here we need to write the live variables back into these same
Expand Down Expand Up @@ -2495,7 +2516,7 @@ enum Immediate {
mod tests {
use super::{Assemble, X64CompiledTrace};
use crate::compile::{
jitc_yk::jit_ir::{self, Inst, Module, ParamIdx},
jitc_yk::jit_ir::{self, Inst, Module, ParamIdx, TraceKind},
CompiledTrace,
};
use crate::location::{HotLocation, HotLocationKind};
Expand Down Expand Up @@ -3984,12 +4005,13 @@ mod tests {
);
}

#[should_panic]
#[test]
fn unterminated_trace() {
codegen_and_test(
"
entry:
body_end []
header_end []
",
"
...
Expand All @@ -4006,13 +4028,13 @@ mod tests {
codegen_and_test(
"
entry:
body_start []
body_end []
header_start []
header_end []
",
"
...
; body_start []
; body_end []
; header_start []
; header_end []
jmp {{target}}
",
false,
Expand All @@ -4025,20 +4047,20 @@ mod tests {
"
entry:
%0: i8 = param 0
body_start [%0]
header_start [%0]
%2: i8 = add %0, %0
black_box %2
body_end [%0]
header_end [%0]
",
"
...
; %0: i8 = param ...
...
; body_start [%0]
; header_start [%0]
; %2: i8 = add %0, %0
{{_}} {{off}}: ...
...
; body_end [%0]
; header_end [%0]
...
{{_}} {{_}}: jmp 0x00000000{{off}}
",
Expand Down Expand Up @@ -4531,16 +4553,16 @@ mod tests {
"
entry:
%0: i8 = param 0
body_start [%0]
header_start [%0]
%2: i8 = 42i8
body_end [%2]
header_end [%2]
",
"
...
; %0: i8 = param ...
...
; body_start [%0]
; body_end [42i8]
; header_start [%0]
; header_end [42i8]
mov r.64.x, 0x2a
jmp ...
",
Expand Down Expand Up @@ -4578,7 +4600,7 @@ mod tests {

#[test]
fn cg_aliasing_params() {
let mut m = jit_ir::Module::new(0, 0).unwrap();
let mut m = jit_ir::Module::new(TraceKind::HeaderOnly, 0, 0).unwrap();

// Create two trace paramaters whose locations alias.
let loc = yksmp::Location::Register(13, 1, [].into());
Expand Down
Loading

0 comments on commit 3508eca

Please sign in to comment.