Peel the peeler

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

@@ -68,7 +68,7 @@ impl<'a> RevAnalyse<'a> {
                 }
 
                 // Calculate inst_vals_alive_until
-                inst.map_operand_locals(self.m, &mut |x| {
+                inst.map_operand_vars(self.m, &mut |x| {
                     self.used_insts.set(usize::from(x), true);
                     if self.inst_vals_alive_until[usize::from(x)] < iidx {
                         self.inst_vals_alive_until[usize::from(x)] = iidx;

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

@@ -17,7 +17,7 @@ impl Module {
                 || inst.is_barrier(self)
             {
                 used.set(usize::from(iidx), true);
-                inst.map_operand_locals(self, &mut |x| {
+                inst.map_operand_vars(self, &mut |x| {
                     used.set(usize::from(x), true);
                 });
             } else {

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

@@ -169,10 +169,10 @@ pub(crate) struct Module {
     /// Live variables at the beginning of the root trace.
     root_entry_vars: Vec<VarLocation>,
     /// Live variables at the beginning of the trace body.
-    trace_body_start: Vec<PackedOperand>,
+    pub(crate) trace_body_start: Vec<PackedOperand>,
     /// The ordered sequence of operands at the end of the trace body: there will be one per
     /// [Operand] at the start of the loop.
-    trace_body_end: Vec<PackedOperand>,
+    pub(crate) trace_body_end: Vec<PackedOperand>,
     /// Live variables at the beginning the trace header.
     trace_header_start: Vec<PackedOperand>,
     /// Live variables at the end of the trace header.
@@ -1521,9 +1521,9 @@ impl Inst {
     }
 
     /// Apply the function `f` to each of this instruction's [Operand]s iff they are of type
-    /// [Operand::Local]. When an instruction references more than one [Operand], the order of
+    /// [Operand::Var]. When an instruction references more than one [Operand], the order of
     /// traversal is undefined.
-    pub(crate) fn map_operand_locals<F>(&self, m: &Module, f: &mut F)
+    pub(crate) fn map_operand_vars<F>(&self, m: &Module, f: &mut F)
     where
         F: FnMut(InstIdx),
     {
@@ -1628,11 +1628,15 @@ impl Inst {
         }
     }
 
-    /// Duplicate this [Inst] while applying function `f` to each operand.
-    pub(crate) fn dup_and_remap_locals<F>(
+    /// Duplicate this [Inst] and, during that process, apply the function `f` to each of this
+    /// instruction's [Operand]s iff they are of type [Operand::Var]. When an instruction
+    /// references more than one [Operand], the order of traversal is undefined.
+    ///
+    /// Note: for `TraceBody*` and `TraceHeader*` functions, this function is not idempotent.
+    pub(crate) fn dup_and_remap_vars<F>(
         &self,
         m: &mut Module,
-        f: &F,
+        f: F,
     ) -> Result<Self, CompilationError>
     where
         F: Fn(InstIdx) -> Operand,
@@ -1788,18 +1792,12 @@ impl Inst {
                 volatile: *volatile,
             }),
             Inst::Tombstone => Inst::Tombstone,
-            Inst::TraceHeaderStart => {
-                // Copy the header label into the body while remapping the operands.
-                m.trace_body_start = m
-                    .trace_header_start
-                    .iter()
-                    .map(|op| mapper(m, op))
-                    .collect();
+            Inst::TraceBodyStart => {
+                m.trace_body_start = m.trace_body_start.iter().map(|op| mapper(m, op)).collect();
                 Inst::TraceBodyStart
             }
-            Inst::TraceHeaderEnd => {
-                // Copy the header label into the body while remapping the operands.
-                m.trace_body_end = m.trace_header_end.iter().map(|op| mapper(m, op)).collect();
+            Inst::TraceBodyEnd => {
+                m.trace_body_end = m.trace_body_end.iter().map(|op| mapper(m, op)).collect();
                 Inst::TraceBodyEnd
             }
             Inst::Trunc(TruncInst { val, dest_tyidx }) => Inst::Trunc(TruncInst {

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

@@ -38,7 +38,7 @@ impl Module {
 
         let mut last_inst = None;
         for (iidx, inst) in self.iter_skipping_insts() {
-            inst.map_operand_locals(self, &mut |x| {
+            inst.map_operand_vars(self, &mut |x| {
                 if let Inst::Tombstone = self.insts[usize::from(x)] {
                     panic!(
                         "Instruction at position {iidx} uses undefined value (%{x})\n  {}",

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

@@ -35,12 +35,111 @@ impl Opt {
         Self { m, an, instll }
     }
 
+    fn opt(mut self) -> Result<Module, CompilationError> {
+        let base = self.m.insts_len();
+        let peel = match self.m.inst(self.m.last_inst_idx()) {
+            // If this is a sidetrace, we perform optimisations up to, but including, loop peeling.
+            Inst::SidetraceEnd => false,
+            Inst::TraceHeaderEnd => true,
+            #[cfg(test)]
+            // Not all tests create "fully valid" traces, in the sense that -- to keep things
+            // simple -- they don't end with `TraceHeaderEnd`. We don't want to peel such traces,
+            // but nor, in testing mode, do we consider them ill-formed.
+            _ => false,
+            #[cfg(not(test))]
+            _ => panic!(),
+        };
+
+        // Note that since we will apply loop peeling here, the list of instructions grows as this
+        // loop runs. Each instruction we process is (after optimisations were applied), duplicated
+        // and copied to the end of the module.
+        let skipping = self.m.iter_skipping_insts().collect::<Vec<_>>();
+        for (iidx, inst) in skipping.into_iter() {
+            match inst {
+                Inst::TraceHeaderStart => (),
+                Inst::TraceHeaderEnd => (),
+                _ => {
+                    self.opt_inst(iidx)?;
+                    self.cse(iidx);
+                }
+            }
+        }
+        // FIXME: When code generation supports backwards register allocation, we won't need to
+        // explicitly perform dead code elimination and this function can be made `#[cfg(test)]` only.
+        self.m.dead_code_elimination();
+
+        if !peel {
+            return Ok(self.m);
+        }
+
+        // Now that we've processed the trace header, duplicate it to create the loop body.
+        // FIXME: Do we need to call `iter_skipping_inst_idxs` again?
+        // Maps header instructions to their position in the body.
+        let mut iidx_map = vec![InstIdx::max(); base];
+        let skipping = self.m.iter_skipping_insts().collect::<Vec<_>>();
+        for (iidx, inst) in skipping.into_iter() {
+            match inst {
+                Inst::TraceHeaderStart => {
+                    self.m.trace_body_start = self.m.trace_header_start().to_vec();
+                    self.m.push(Inst::TraceBodyStart)?;
+                    // FIXME: We rely on `dup_and_remap_vars` not being idempotent here.
+                    let _ = Inst::TraceBodyStart
+                        .dup_and_remap_vars(&mut self.m, |op_iidx: InstIdx| {
+                            Operand::Var(iidx_map[usize::from(op_iidx)])
+                        })?;
+                    for (headop, bodyop) in self
+                        .m
+                        .trace_header_end()
+                        .iter()
+                        .zip(self.m.trace_body_start())
+                    {
+                        // Inform the analyser about any constants being passed from the header into
+                        // the body.
+                        if let Operand::Const(cidx) = headop.unpack(&self.m) {
+                            let Operand::Var(op_iidx) = bodyop.unpack(&self.m) else {
+                                panic!()
+                            };
+                            self.an.set_value(op_iidx, Value::Const(cidx));
+                        }
+                    }
+                }
+                Inst::TraceHeaderEnd => {
+                    self.m.trace_body_end = self.m.trace_header_end().to_vec();
+                    self.m.push(Inst::TraceBodyEnd)?;
+                    // FIXME: We rely on `dup_and_remap_vars` not being idempotent here.
+                    let _ = Inst::TraceBodyEnd
+                        .dup_and_remap_vars(&mut self.m, |op_iidx: InstIdx| {
+                            Operand::Var(iidx_map[usize::from(op_iidx)])
+                        })?;
+                }
+                _ => {
+                    let c = inst.dup_and_remap_vars(&mut self.m, |op_iidx: InstIdx| {
+                        Operand::Var(iidx_map[usize::from(op_iidx)])
+                    })?;
+                    let copy_iidx = self.m.push(c)?;
+                    iidx_map[usize::from(iidx)] = copy_iidx;
+                    self.opt_inst(copy_iidx)?;
+                }
+            }
+        }
+
+        // FIXME: Apply CSE and run another pass of optimisations on the peeled loop.
+        self.m.dead_code_elimination();
+        Ok(self.m)
+    }
+
     /// Optimise instruction `iidx`.
     fn opt_inst(&mut self, iidx: InstIdx) -> Result<(), CompilationError> {
         match self.m.inst(iidx) {
             #[cfg(test)]
             Inst::BlackBox(_) => (),
-            Inst::Const(_) | Inst::Copy(_) | Inst::Tombstone => unreachable!(),
+            Inst::Const(_)
+            | Inst::Copy(_)
+            | Inst::Tombstone
+            | Inst::TraceHeaderStart
+            | Inst::TraceHeaderEnd => {
+                unreachable!()
+            }
             Inst::BinOp(x) => match x.binop() {
                 BinOp::Add => match (
                     self.an.op_map(&self.m, x.lhs(&self.m)),
@@ -367,74 +466,6 @@ impl Opt {
         Ok(())
     }
 
-    fn opt(mut self) -> Result<Module, CompilationError> {
-        let base = self.m.insts_len();
-        let peel = match self.m.inst(self.m.last_inst_idx()) {
-            // If this is a sidetrace, we perform optimisations up to, but including, loop peeling.
-            Inst::SidetraceEnd => false,
-            Inst::TraceHeaderEnd => true,
-            #[cfg(test)]
-            // Not all tests create "fully valid" traces, in the sense that -- to keep things
-            // simple -- they don't end with `TraceHeaderEnd`. We don't want to peel such traces,
-            // but nor, in testing mode, do we consider them ill-formed.
-            _ => false,
-            #[cfg(not(test))]
-            _ => panic!(),
-        };
-
-        // Note that since we will apply loop peeling here, the list of instructions grows as this
-        // loop runs. Each instruction we process is (after optimisations were applied), duplicated
-        // and copied to the end of the module.
-        let skipping = self.m.iter_skipping_insts().collect::<Vec<_>>();
-        for (iidx, _inst) in skipping.into_iter() {
-            self.opt_inst(iidx)?;
-            self.cse(iidx);
-        }
-        // FIXME: When code generation supports backwards register allocation, we won't need to
-        // explicitly perform dead code elimination and this function can be made `#[cfg(test)]` only.
-        self.m.dead_code_elimination();
-
-        if !peel {
-            return Ok(self.m);
-        }
-
-        // Now that we've processed the trace header, duplicate it to create the loop body.
-        // FIXME: Do we need to call `iter_skipping_inst_idxs` again?
-        // Maps header instructions to their position in the body.
-        let mut iidx_map = vec![0; base];
-        let skipping = self.m.iter_skipping_insts().collect::<Vec<_>>();
-        for (iidx, inst) in skipping.into_iter() {
-            let c = inst.dup_and_remap_locals(&mut self.m, &|i: InstIdx| {
-                let newiidx = iidx_map[usize::from(i)];
-                Operand::Var(InstIdx::try_from(newiidx).unwrap())
-            })?;
-            let copyiidx = self.m.push(c)?;
-            iidx_map[usize::from(iidx)] = usize::from(copyiidx);
-            if let Inst::TraceHeaderStart = inst {
-                for (headop, bodyop) in self
-                    .m
-                    .trace_header_end()
-                    .iter()
-                    .zip(self.m.trace_body_start())
-                {
-                    // Inform the analyser about any constants being passed from the header into
-                    // the body.
-                    if let Operand::Const(cidx) = headop.unpack(&self.m) {
-                        let Operand::Var(op_iidx) = bodyop.unpack(&self.m) else {
-                            panic!()
-                        };
-                        self.an.set_value(op_iidx, Value::Const(cidx));
-                    }
-                }
-            }
-            self.opt_inst(copyiidx)?;
-        }
-
-        // FIXME: Apply CSE and run another pass of optimisations on the peeled loop.
-        self.m.dead_code_elimination();
-        Ok(self.m)
-    }
-
     /// Attempt common subexpression elimination on `iidx`, replacing it with a `Copy` or
     /// `Tombstone` if possible.
     fn cse(&mut self, iidx: InstIdx) {