Add register repacking

CHANGELOG.md

@@ -33,6 +33,11 @@
       code using `anyhow` (or similar) may not need to change
 - Fix a bug in bulk evaluator argument checks where mismatched slices could be
   allowed under some circumstances
+- Add `VmData::asm` to get an immutable reference to the inner `RegTape`
+- Add `RegTape::repack_map` and `RegTape::repack` to repack registers by
+  frequency (making register 0 the most frequently used, etc)
+- Add `RegOp::visit_regs` and `RegOp::visit_regs_mut` to visit registers in an
+  operation
 
 # 0.4.3
 - Fixed bug in x86 interval `OR` function ([#395](https://github.com/mkeeter/fidget/pull/395)),

fidget-bytecode/src/lib.rs

@@ -189,10 +189,16 @@ impl Bytecode {
 
     /// Builds a new bytecode object from VM data
     ///
-    /// Returns an error if the reserved register (255) is in use
+    /// Registers are reordered by frequency of use, e.g. the most frequently
+    /// used register becomes register 0.
+    ///
+    /// Returns an error if the reserved register (255) is in use, which should
+    /// only happen if the incoming tape has 256 active registers.
     pub fn new<const N: usize>(
         t: &VmData<N>,
     ) -> Result<Self, ReservedRegister> {
+        // Build a map for repacking registers by frequency
+        let map = t.asm().repack_map();
         // The initial opcode is `OP_JUMP 0x0000_0000`
         let mut data = vec![u32::MAX, 0u32];
         let mut reg_count = 0u8;
@@ -202,6 +208,7 @@ impl Bytecode {
             let mut word = [0xFF; 4];
             let mut imm = None;
             let mut store_reg = |i, r| {
+                let r = map[&r];
                 if r == u8::MAX {
                     Err(ReservedRegister)
                 } else {

fidget-core/src/compiler/op.rs

@@ -288,3 +288,159 @@ opcodes!(
         Store(u8, u32),
     }
 );
+
+impl RegOp {
+    /// Apply a mutating function to every register in the op
+    ///
+    /// Both inputs and outputs are visited
+    pub fn visit_regs_mut<F: FnMut(&mut u8)>(&mut self, mut f: F) {
+        match self {
+            RegOp::CopyImm(out, imm) => {
+                let _: f32 = *imm;
+                f(out)
+            }
+            RegOp::NegReg(out, arg)
+            | RegOp::AbsReg(out, arg)
+            | RegOp::RecipReg(out, arg)
+            | RegOp::SqrtReg(out, arg)
+            | RegOp::SquareReg(out, arg)
+            | RegOp::FloorReg(out, arg)
+            | RegOp::CeilReg(out, arg)
+            | RegOp::RoundReg(out, arg)
+            | RegOp::CopyReg(out, arg)
+            | RegOp::SinReg(out, arg)
+            | RegOp::CosReg(out, arg)
+            | RegOp::TanReg(out, arg)
+            | RegOp::AsinReg(out, arg)
+            | RegOp::AcosReg(out, arg)
+            | RegOp::AtanReg(out, arg)
+            | RegOp::ExpReg(out, arg)
+            | RegOp::LnReg(out, arg)
+            | RegOp::NotReg(out, arg) => {
+                f(out);
+                f(arg);
+            }
+            RegOp::AddRegImm(out, arg, imm)
+            | RegOp::MulRegImm(out, arg, imm)
+            | RegOp::DivRegImm(out, arg, imm)
+            | RegOp::DivImmReg(out, arg, imm)
+            | RegOp::SubImmReg(out, arg, imm)
+            | RegOp::SubRegImm(out, arg, imm)
+            | RegOp::AtanRegImm(out, arg, imm)
+            | RegOp::AtanImmReg(out, arg, imm)
+            | RegOp::MinRegImm(out, arg, imm)
+            | RegOp::MaxRegImm(out, arg, imm)
+            | RegOp::CompareRegImm(out, arg, imm)
+            | RegOp::CompareImmReg(out, arg, imm)
+            | RegOp::ModRegImm(out, arg, imm)
+            | RegOp::ModImmReg(out, arg, imm)
+            | RegOp::AndRegImm(out, arg, imm)
+            | RegOp::OrRegImm(out, arg, imm) => {
+                let _: f32 = *imm; // type-checking pattern
+                f(out);
+                f(arg);
+            }
+
+            RegOp::AddRegReg(out, lhs, rhs)
+            | RegOp::MulRegReg(out, lhs, rhs)
+            | RegOp::DivRegReg(out, lhs, rhs)
+            | RegOp::SubRegReg(out, lhs, rhs)
+            | RegOp::AtanRegReg(out, lhs, rhs)
+            | RegOp::MinRegReg(out, lhs, rhs)
+            | RegOp::MaxRegReg(out, lhs, rhs)
+            | RegOp::CompareRegReg(out, lhs, rhs)
+            | RegOp::ModRegReg(out, lhs, rhs)
+            | RegOp::AndRegReg(out, lhs, rhs)
+            | RegOp::OrRegReg(out, lhs, rhs) => {
+                f(out);
+                f(lhs);
+                f(rhs);
+            }
+
+            RegOp::Output(reg, imm)
+            | RegOp::Input(reg, imm)
+            | RegOp::Store(reg, imm)
+            | RegOp::Load(reg, imm) => {
+                let _: u32 = *imm; // type-checking pattern
+                f(reg)
+            }
+        }
+    }
+
+    /// Apply a function to every register in the op
+    ///
+    /// Both inputs and outputs are visited
+    pub fn visit_regs<F: FnMut(u8)>(&self, mut f: F) {
+        match self {
+            RegOp::CopyImm(out, imm) => {
+                let _: f32 = *imm;
+                f(*out)
+            }
+            RegOp::NegReg(out, arg)
+            | RegOp::AbsReg(out, arg)
+            | RegOp::RecipReg(out, arg)
+            | RegOp::SqrtReg(out, arg)
+            | RegOp::SquareReg(out, arg)
+            | RegOp::FloorReg(out, arg)
+            | RegOp::CeilReg(out, arg)
+            | RegOp::RoundReg(out, arg)
+            | RegOp::CopyReg(out, arg)
+            | RegOp::SinReg(out, arg)
+            | RegOp::CosReg(out, arg)
+            | RegOp::TanReg(out, arg)
+            | RegOp::AsinReg(out, arg)
+            | RegOp::AcosReg(out, arg)
+            | RegOp::AtanReg(out, arg)
+            | RegOp::ExpReg(out, arg)
+            | RegOp::LnReg(out, arg)
+            | RegOp::NotReg(out, arg) => {
+                f(*out);
+                f(*arg);
+            }
+            RegOp::AddRegImm(out, arg, imm)
+            | RegOp::MulRegImm(out, arg, imm)
+            | RegOp::DivRegImm(out, arg, imm)
+            | RegOp::DivImmReg(out, arg, imm)
+            | RegOp::SubImmReg(out, arg, imm)
+            | RegOp::SubRegImm(out, arg, imm)
+            | RegOp::AtanRegImm(out, arg, imm)
+            | RegOp::AtanImmReg(out, arg, imm)
+            | RegOp::MinRegImm(out, arg, imm)
+            | RegOp::MaxRegImm(out, arg, imm)
+            | RegOp::CompareRegImm(out, arg, imm)
+            | RegOp::CompareImmReg(out, arg, imm)
+            | RegOp::ModRegImm(out, arg, imm)
+            | RegOp::ModImmReg(out, arg, imm)
+            | RegOp::AndRegImm(out, arg, imm)
+            | RegOp::OrRegImm(out, arg, imm) => {
+                let _: f32 = *imm; // type-checking pattern
+                f(*out);
+                f(*arg);
+            }
+
+            RegOp::AddRegReg(out, lhs, rhs)
+            | RegOp::MulRegReg(out, lhs, rhs)
+            | RegOp::DivRegReg(out, lhs, rhs)
+            | RegOp::SubRegReg(out, lhs, rhs)
+            | RegOp::AtanRegReg(out, lhs, rhs)
+            | RegOp::MinRegReg(out, lhs, rhs)
+            | RegOp::MaxRegReg(out, lhs, rhs)
+            | RegOp::CompareRegReg(out, lhs, rhs)
+            | RegOp::ModRegReg(out, lhs, rhs)
+            | RegOp::AndRegReg(out, lhs, rhs)
+            | RegOp::OrRegReg(out, lhs, rhs) => {
+                f(*out);
+                f(*lhs);
+                f(*rhs);
+            }
+
+            RegOp::Output(reg, imm)
+            | RegOp::Input(reg, imm)
+            | RegOp::Store(reg, imm)
+            | RegOp::Load(reg, imm) => {
+                let _: u32 = *imm; // type-checking pattern
+                f(*reg)
+            }
+        }
+    }
+}

fidget-core/src/compiler/reg_tape.rs

@@ -1,6 +1,7 @@
 //! Tape used for evaluation
 use crate::compiler::{RegOp, RegisterAllocator, SsaTape};
 use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
 
 /// Low-level tape for use with the Fidget virtual machine (or to be lowered
 /// further into machine instructions).
@@ -9,6 +10,9 @@ pub struct RegTape {
     tape: Vec<RegOp>,
 
     /// Total allocated slots
+    ///
+    /// This is a continuous space of registers (`0..N`) and memory (`N..`),
+    /// where `N` is the parameter in [`RegTape::new`].
     pub(super) slot_count: u32,
 }
 
@@ -27,6 +31,35 @@ impl RegTape {
         alloc.finalize()
     }
 
+    /// Repacks registers by frequency (so that register 0 is the most frequent)
+    pub fn repack(&mut self) {
+        let map = self.repack_map();
+        for op in &mut self.tape {
+            op.visit_regs_mut(|reg| *reg = map[reg]);
+        }
+    }
+
+    /// Returns a map for register repacking
+    ///
+    /// The map repacks registers in the tape by frequency, so that register 0
+    /// is the most frequent.
+    pub fn repack_map(&self) -> HashMap<u8, u8> {
+        let mut reg_counts: HashMap<u8, usize> = HashMap::new();
+        for op in &self.tape {
+            op.visit_regs(|reg| *reg_counts.entry(reg).or_default() += 1);
+        }
+        let mut sorted = reg_counts
+            .into_iter()
+            .map(|(reg, count)| (std::cmp::Reverse(count), reg))
+            .collect::<Vec<_>>();
+        sorted.sort_unstable();
+        sorted
+            .into_iter()
+            .enumerate()
+            .map(|(i, (_count, reg))| (reg, u8::try_from(i).unwrap()))
+            .collect()
+    }
+
     /// Builds a new empty tape
     pub(crate) fn empty() -> Self {
         Self {

fidget-core/src/vm/data.rs

@@ -318,6 +318,11 @@ impl<const N: usize> VmData<N> {
         self.asm.iter().cloned().rev()
     }
 
+    /// Returns a reference to the inner [`RegTape`]
+    pub fn asm(&self) -> &RegTape {
+        &self.asm
+    }
+
     /// Pretty-prints the inner SSA tape
     pub fn pretty_print(&self) {
         self.ssa.pretty_print();