Pull side-trace codegen-specific stuff out of the JIT module.

This commit is, again, a bit of a hack, but it still makes things a bit
better than before. In essence, before this commit JIT modules contained
`yksmp::Location`s for "normal" traces and `VarLocation`s for side
traces. The former are codegen agnostic: the latter are codegen
specific.

This commit hoiks the `VarLocation`s out of the JIT mod and puts them in
`TraceKind::Sidetrace`. Now that would end up having us add a `Register`
type parameter everywhere, which is too much churn, so for now
`TraceKind::Sidetrace` stores an `Arc<SidetraceInfo>` i.e. we store it
as a trait object. We then have to downcast to `YkSideTraceInfo` in lots
of places. This is pretty evil, because we're using dynamic typing where
we should be using static typing. However, overall, this is the lesser
of the evils.

By the end of this commit, it's clearer (even if not perfectly so) the
relationship between trace inputs and side-trace outputs. We can do
better in the future, and maybe we'll want to do so, but this does feel
like a nudge in the right direction.

Author: ltratt

ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs

@@ -1029,9 +1029,7 @@ impl LSRegAlloc<'_> {
                 RegConstraint::Output
                 | RegConstraint::OutputCanBeSameAsInput(_)
                 | RegConstraint::InputOutput(_) => {
-                    if let Some(reg_alloc::Register::FP(reg)) =
-                        self.rev_an.reg_hints[usize::from(iidx)]
-                    {
+                    if let Some(Register::FP(reg)) = self.rev_an.reg_hints[usize::from(iidx)] {
                         if !avoid.is_set(reg) {
                             *cnstr = match cnstr {
                                 RegConstraint::Output => RegConstraint::OutputFromReg(reg),

ykrt/src/compile/jitc_yk/codegen/x64/mod.rs

@@ -53,6 +53,7 @@ use dynasmrt::{
 use indexmap::IndexMap;
 use parking_lot::Mutex;
 use std::{
+    assert_matches::debug_assert_matches,
     cell::Cell,
     collections::HashMap,
     error::Error,
@@ -643,7 +644,7 @@ impl<'a> Assemble<'a> {
                 }
                 jit_ir::Inst::TraceBodyStart => self.cg_body_start(),
                 jit_ir::Inst::TraceBodyEnd => self.cg_body_end(iidx),
-                jit_ir::Inst::SidetraceEnd => self.cg_sidetrace_end(iidx, self.m.root_jump_addr()),
+                jit_ir::Inst::SidetraceEnd => self.cg_sidetrace_end(iidx),
                 jit_ir::Inst::SExt(i) => self.cg_sext(iidx, i),
                 jit_ir::Inst::ZExt(i) => self.cg_zext(iidx, i),
                 jit_ir::Inst::BitCast(i) => self.cg_bitcast(iidx, i),
@@ -1951,9 +1952,17 @@ impl<'a> Assemble<'a> {
     /// * `tgt_vars` - The target locations. If `None` use `self.loop_start_locs` instead.
     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()),
+            TraceKind::HeaderOnly => (self.header_start_locs.clone(), self.m.trace_header_end()),
+            TraceKind::HeaderAndBody => (self.body_start_locs.clone(), self.m.trace_body_end()),
+            TraceKind::Sidetrace(sti) => (
+                Arc::clone(sti)
+                    .as_any()
+                    .downcast::<YkSideTraceInfo<Register>>()
+                    .unwrap()
+                    .entry_vars
+                    .clone(),
+                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.
@@ -2046,7 +2055,7 @@ impl<'a> Assemble<'a> {
     }
 
     fn cg_body_end(&mut self, iidx: InstIdx) {
-        debug_assert_eq!(self.m.tracekind(), TraceKind::HeaderAndBody);
+        debug_assert_matches!(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");
@@ -2074,21 +2083,29 @@ 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);
-        self.ra.align_stack(SYSV_CALL_STACK_ALIGN);
+    fn cg_sidetrace_end(&mut self, iidx: InstIdx) {
+        match self.m.tracekind() {
+            TraceKind::Sidetrace(sti) => {
+                let sti = Arc::clone(sti)
+                    .as_any()
+                    .downcast::<YkSideTraceInfo<Register>>()
+                    .unwrap();
+                // 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);
+                self.ra.align_stack(SYSV_CALL_STACK_ALIGN);
 
-        dynasm!(self.asm
-            // Reset rsp to the root trace's frame.
-            ; mov rsp, rbp
-            ; sub rsp, i32::try_from(self.root_offset.unwrap()).unwrap()
-            ; mov rdi, QWORD addr as i64
-            // We can safely use RDI here, since the root trace won't expect live variables in this
-            // register since it's being used as an argument to the control point.
-            ; jmp rdi);
+                dynasm!(self.asm
+                    // Reset rsp to the root trace's frame.
+                    ; mov rsp, rbp
+                    ; sub rsp, i32::try_from(sti.root_offset()).unwrap()
+                    ; mov rdi, QWORD sti.root_addr() as i64
+                    // We can safely use RDI here, since the root trace won't expect live variables in this
+                    // register since it's being used as an argument to the control point.
+                    ; jmp rdi);
+            }
+            TraceKind::HeaderOnly | TraceKind::HeaderAndBody => panic!(),
+        }
     }
 
     fn cg_header_start(&mut self) {
@@ -2110,7 +2127,7 @@ impl<'a> Assemble<'a> {
             TraceKind::HeaderAndBody => {
                 dynasm!(self.asm; ->reentry:);
             }
-            TraceKind::Sidetrace => todo!(),
+            TraceKind::Sidetrace(_) => todo!(),
         }
         self.prologue_offset = self.asm.offset();
     }
@@ -2172,12 +2189,12 @@ impl<'a> Assemble<'a> {
                     }
                 }
             }
-            TraceKind::Sidetrace => todo!(),
+            TraceKind::Sidetrace(_) => panic!(),
         }
     }
 
     fn cg_body_start(&mut self) {
-        debug_assert_eq!(self.m.tracekind(), TraceKind::HeaderAndBody);
+        debug_assert_matches!(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

ykrt/src/compile/jitc_yk/codegen/x64/rev_analyse.rs

@@ -19,8 +19,10 @@ use crate::compile::jitc_yk::{
         BinOp, BinOpInst, DirectCallInst, DynPtrAddInst, ICmpInst, Inst, InstIdx, LoadInst, Module,
         Operand, PtrAddInst, SExtInst, SelectInst, StoreInst, TraceKind, TruncInst, Ty, ZExtInst,
     },
+    YkSideTraceInfo,
 };
 use dynasmrt::x64::Rq;
+use std::sync::Arc;
 use vob::Vob;
 
 pub(crate) struct RevAnalyse<'a> {
@@ -80,8 +82,12 @@ impl<'a> RevAnalyse<'a> {
                 // We don't care where the register allocator ends at the end of the header, so we
                 // don't propagate backwards from `TraceHeaderEnd`.
             }
-            TraceKind::Sidetrace => {
-                let vlocs = self.m.root_entry_vars();
+            TraceKind::Sidetrace(sti) => {
+                let sti = Arc::clone(sti)
+                    .as_any()
+                    .downcast::<YkSideTraceInfo<Register>>()
+                    .unwrap();
+                let vlocs = &sti.entry_vars;
                 // Side-traces don't have a trace body since we don't apply loop peeling and thus use
                 // `trace_header_end` to store the jump variables.
                 debug_assert_eq!(vlocs.len(), self.m.trace_header_end().len());
@@ -97,7 +103,7 @@ impl<'a> RevAnalyse<'a> {
         // ...and then we perform the rest of the reverse analysis.
         let mut iter = self.m.iter_skipping_insts().rev();
         match self.m.tracekind() {
-            TraceKind::HeaderOnly | TraceKind::Sidetrace => {
+            TraceKind::HeaderOnly | TraceKind::Sidetrace(_) => {
                 for (iidx, inst) in self.m.iter_skipping_insts().rev() {
                     self.analyse(iidx, inst);
                 }

ykrt/src/compile/jitc_yk/jit_ir/mod.rs

@@ -96,16 +96,17 @@ mod parser;
 #[cfg(any(debug_assertions, test))]
 mod well_formed;
 
+use super::aot_ir;
 #[cfg(debug_assertions)]
 use super::int_signs::Truncate;
-use super::{aot_ir, VarLocation};
-use crate::compile::CompilationError;
+use crate::compile::{CompilationError, SideTraceInfo};
 use indexmap::IndexSet;
 use std::{
     ffi::{c_void, CString},
     fmt,
     hash::Hash,
     mem,
+    sync::Arc,
 };
 use strum::{EnumCount, EnumDiscriminants};
 #[cfg(not(test))]
@@ -115,15 +116,15 @@ use ykaddr::addr::symbol_to_ptr;
 pub(crate) use super::aot_ir::{BinOp, FloatPredicate, FloatTy, Predicate};
 
 /// What kind of trace does this module represent?
-#[derive(Copy, Clone, Debug, PartialEq)]
+#[derive(Clone, Debug)]
 pub(crate) enum TraceKind {
     /// A trace which contains only a header: the trace must loop back to the very start every
     /// time.
     HeaderOnly,
     /// A trace with a header and a body: the trace must loop back to the start of the body.
     HeaderAndBody,
     /// A sidetrace: the trace must loop back to the root of the trace tree.
-    Sidetrace,
+    Sidetrace(Arc<dyn SideTraceInfo>),
 }
 
 /// The `Module` is the top-level container for JIT IR.
@@ -180,8 +181,6 @@ pub(crate) struct Module {
     guard_info: Vec<GuardInfo>,
     /// Indirect calls.
     indirect_calls: Vec<IndirectCallInst>,
-    /// Live variables at the beginning of the root trace.
-    root_entry_vars: Vec<VarLocation>,
     /// Live variables at the beginning of the trace body.
     pub(crate) trace_body_start: Vec<PackedOperand>,
     /// The ordered sequence of operands at the end of the trace body: there will be one per
@@ -215,15 +214,15 @@ impl Module {
 
     /// Returns this module's current [TraceKind]. Note: this can change as a result of calling
     /// [Self::set_tracekind]!
-    pub(crate) fn tracekind(&self) -> TraceKind {
-        self.tracekind
+    pub(crate) fn tracekind(&self) -> &TraceKind {
+        &self.tracekind
     }
 
     /// Returns this module's current [TraceKind]. Currently the only transition allowed is from
     /// [TraceKind::HeaderOnly] to [TraceKind::HeaderAndBody].
     pub(crate) fn set_tracekind(&mut self, tracekind: TraceKind) {
-        match (self.tracekind, tracekind) {
-            (TraceKind::HeaderOnly, TraceKind::HeaderAndBody) => (),
+        match (&self.tracekind, &tracekind) {
+            (&TraceKind::HeaderOnly, &TraceKind::HeaderAndBody) => (),
             (from, to) => panic!("Can't transition from a {from:?} trace to a {to:?} trace"),
         }
         self.tracekind = tracekind;
@@ -304,7 +303,6 @@ impl Module {
             global_decls: IndexSet::new(),
             guard_info: Vec::new(),
             indirect_calls: Vec::new(),
-            root_entry_vars: Vec::new(),
             trace_body_start: Vec::new(),
             trace_body_end: Vec::new(),
             trace_header_start: Vec::new(),
@@ -642,37 +640,14 @@ impl Module {
         self.trace_header_start.push(PackedOperand::new(&op));
     }
 
-    /// Store the entry live variables of the root traces so we can copy this side-trace's live
-    /// variables to the right place before jumping back to the root trace.
-    pub(crate) fn set_root_entry_vars(&mut self, entry_vars: &[VarLocation]) {
-        self.root_entry_vars.extend_from_slice(entry_vars);
-    }
-
     /// Return the loop jump operands.
     pub(crate) fn trace_body_end(&self) -> &[PackedOperand] {
         &self.trace_body_end
     }
 
-    /// Get the entry live variables of the root trace.
-    pub(crate) fn root_entry_vars(&self) -> &[VarLocation] {
-        &self.root_entry_vars
-    }
-
     pub(crate) fn push_header_end_var(&mut self, op: Operand) {
         self.trace_header_end.push(PackedOperand::new(&op));
     }
-
-    /// Get the address of the root trace. This is where we need jump to at the end of a
-    /// side-trace.
-    pub(crate) fn root_jump_addr(&self) -> *const libc::c_void {
-        self.root_jump_ptr
-    }
-
-    /// Set the entry address of the root trace. This is where we need jump to at the end of a
-    /// side-trace.
-    pub(crate) fn set_root_jump_addr(&mut self, ptr: *const libc::c_void) {
-        self.root_jump_ptr = ptr;
-    }
 }
 
 impl fmt::Display for Module {

ykrt/src/compile/jitc_yk/jit_ir/well_formed.rs

@@ -28,12 +28,17 @@ use super::{BinOp, BinOpInst, Const, GuardInst, Inst, Module, Operand, Ty};
 
 impl Module {
     pub(crate) fn assert_well_formed(&self) {
-        if !self.root_entry_vars.is_empty() {
-            if self.root_entry_vars.len() != self.trace_header_end.len() {
-                panic!("Loop start/end variables have different lengths.");
+        match self.tracekind() {
+            super::TraceKind::HeaderOnly | super::TraceKind::HeaderAndBody => {
+                if self.trace_header_start.len() != self.trace_header_end.len() {
+                    panic!(
+                        "Loop start ({}) / end ({}) variables have different length.",
+                        self.trace_header_start.len(),
+                        self.trace_header_end.len()
+                    );
+                }
             }
-        } else if self.trace_header_start.len() != self.trace_header_end.len() {
-            panic!("Loop start/end variables have different lengths.");
+            super::TraceKind::Sidetrace(_) => (),
         }
 
         let mut last_inst = None;