Remove an explicitly-set-aside scratch register per class.

fuzz/fuzz_targets/moves.rs

@@ -7,75 +7,121 @@
 use libfuzzer_sys::arbitrary::{Arbitrary, Result, Unstructured};
 use libfuzzer_sys::fuzz_target;
 
-use regalloc2::fuzzing::moves::ParallelMoves;
-use regalloc2::{Allocation, PReg, RegClass};
-use std::collections::HashSet;
+use regalloc2::fuzzing::moves::{MoveAndScratchResolver, ParallelMoves};
+use regalloc2::{Allocation, PReg, RegClass, SpillSlot};
+use std::collections::{HashMap, HashSet};
 
 #[derive(Clone, Debug)]
 struct TestCase {
     moves: Vec<(Allocation, Allocation)>,
+    available_pregs: Vec<Allocation>,
 }
 
 impl Arbitrary for TestCase {
     fn arbitrary(u: &mut Unstructured) -> Result<Self> {
-        let mut ret = TestCase { moves: vec![] };
+        let mut ret = TestCase {
+            moves: vec![],
+            available_pregs: vec![],
+        };
         let mut written = HashSet::new();
+        // An arbitrary sequence of moves between registers 0 to 29
+        // inclusive.
         while bool::arbitrary(u)? {
-            let reg1 = u.int_in_range(0..=30)?;
-            let reg2 = u.int_in_range(0..=30)?;
-            if written.contains(&reg2) {
+            let src = if bool::arbitrary(u)? {
+                let reg = u.int_in_range(0..=29)?;
+                Allocation::reg(PReg::new(reg, RegClass::Int))
+            } else {
+                let slot = u.int_in_range(0..=31)?;
+                Allocation::stack(SpillSlot::new(slot, RegClass::Int))
+            };
+            let dst = if bool::arbitrary(u)? {
+                let reg = u.int_in_range(0..=29)?;
+                Allocation::reg(PReg::new(reg, RegClass::Int))
+            } else {
+                let slot = u.int_in_range(0..=31)?;
+                Allocation::stack(SpillSlot::new(slot, RegClass::Int))
+            };
+
+            // Stop if we are going to write a reg more than once:
+            // that creates an invalid parallel move set.
+            if written.contains(&dst) {
                 break;
             }
-            written.insert(reg2);
-            ret.moves.push((
-                Allocation::reg(PReg::new(reg1, RegClass::Int)),
-                Allocation::reg(PReg::new(reg2, RegClass::Int)),
-            ));
+            written.insert(dst);
+
+            ret.moves.push((src, dst));
+        }
+
+        // We might have some unallocated registers free for scratch
+        // space...
+        for i in u.int_in_range(0..=2) {
+            let reg = PReg::new(30 + i, RegClass::Int);
+            ret.available_pregs.push(Allocation::reg(reg));
         }
         Ok(ret)
     }
 }
 
 fuzz_target!(|testcase: TestCase| {
     let _ = env_logger::try_init();
-    let scratch = Allocation::reg(PReg::new(31, RegClass::Int));
-    let mut par = ParallelMoves::new(scratch);
+    let mut par = ParallelMoves::new();
     for &(src, dst) in &testcase.moves {
         par.add(src, dst, ());
     }
+
     let moves = par.resolve();
+    log::trace!("raw resolved moves: {:?}", moves);
+
+    // Resolve uses of scratch reg and stack-to-stack moves with the
+    // scratch resolver.
+    let mut avail = testcase.available_pregs.clone();
+    let get_reg = || avail.pop();
+    let mut next_slot = 32;
+    let get_stackslot = || {
+        let slot = next_slot;
+        next_slot += 1;
+        Allocation::stack(SpillSlot::new(slot, RegClass::Int))
+    };
+    let preferred_victim = PReg::new(0, RegClass::Int);
+    let scratch_resolver = MoveAndScratchResolver::new(get_reg, get_stackslot, preferred_victim);
+    let moves = scratch_resolver.compute(moves);
+    log::trace!("resolved moves: {:?}", moves);
 
     // Compute the final source reg for each dest reg in the original
     // parallel-move set.
-    let mut final_src_per_dest: Vec<Option<usize>> = vec![None; 32];
+    let mut final_src_per_dest: HashMap<Allocation, Allocation> = HashMap::new();
     for &(src, dst) in &testcase.moves {
-        if let (Some(preg_src), Some(preg_dst)) = (src.as_reg(), dst.as_reg()) {
-            final_src_per_dest[preg_dst.hw_enc()] = Some(preg_src.hw_enc());
-        }
+        final_src_per_dest.insert(dst, src);
     }
+    log::trace!("expected final state: {:?}", final_src_per_dest);
 
     // Simulate the sequence of moves.
-    let mut regfile: Vec<Option<usize>> = vec![None; 32];
-    for i in 0..32 {
-        regfile[i] = Some(i);
-    }
+    let mut locations: HashMap<Allocation, Allocation> = HashMap::new();
     for (src, dst, _) in moves {
-        if let (Some(preg_src), Some(preg_dst)) = (src.as_reg(), dst.as_reg()) {
-            let data = regfile[preg_src.hw_enc()];
-            regfile[preg_dst.hw_enc()] = data;
-        } else {
-            panic!("Bad allocation in move list");
+        if src.is_stack() && dst.is_stack() {
+            panic!("Stack-to-stack move!");
         }
+
+        let data = locations.get(&src).cloned().unwrap_or(src);
+        locations.insert(dst, data);
     }
+    log::trace!("simulated final state: {:?}", locations);
 
     // Assert that the expected register-moves occurred.
-    // N.B.: range up to 31 (not 32) to skip scratch register.
-    for i in 0..31 {
-        if let Some(orig_src) = final_src_per_dest[i] {
-            assert_eq!(regfile[i], Some(orig_src));
+    for (reg, data) in locations {
+        if let Some(&expected_data) = final_src_per_dest.get(&reg) {
+            assert_eq!(expected_data, data);
         } else {
-            // Should be untouched.
-            assert_eq!(regfile[i], Some(i));
+            if data != reg {
+                // If not just the original value, then this location
+                // has been modified, but it was not part of the
+                // original parallel move. It must have been an
+                // available preg or a scratch stackslot.
+                assert!(
+                    testcase.available_pregs.contains(&reg)
+                        || (reg.is_stack() && reg.as_stack().unwrap().index() >= 32)
+                );
+            }
         }
     }
 });

src/fuzzing/func.rs

@@ -625,18 +625,15 @@ impl std::fmt::Debug for Func {
 }
 
 pub fn machine_env() -> MachineEnv {
-    // Reg 63 is the scratch reg.
     fn regs(r: std::ops::Range<usize>) -> Vec<PReg> {
         r.map(|i| PReg::new(i, RegClass::Int)).collect()
     }
     let preferred_regs_by_class: [Vec<PReg>; 2] = [regs(0..24), vec![]];
     let non_preferred_regs_by_class: [Vec<PReg>; 2] = [regs(24..32), vec![]];
-    let scratch_by_class: [PReg; 2] = [PReg::new(63, RegClass::Int), PReg::new(0, RegClass::Float)];
-    let fixed_stack_slots = regs(32..63);
+    let fixed_stack_slots = regs(32..64);
     MachineEnv {
         preferred_regs_by_class,
         non_preferred_regs_by_class,
-        scratch_by_class,
         fixed_stack_slots,
     }
 }

src/ion/data_structures.rs

@@ -351,7 +351,8 @@ pub struct Env<'a, F: Function> {
     pub spillslots: Vec<SpillSlotData>,
     pub slots_by_size: Vec<SpillSlotList>,
 
-    pub extra_spillslot: Vec<Option<Allocation>>,
+    pub extra_spillslots_by_class: [SmallVec<[Allocation; 2]>; 2],
+    pub preferred_victim_by_class: [PReg; 2],
 
     // Program moves: these are moves in the provided program that we
     // handle with our internal machinery, in order to avoid the

src/ion/liveranges.rs

@@ -109,6 +109,13 @@ impl<'a, F: Function> Env<'a, F> {
         for &preg in &self.env.fixed_stack_slots {
             self.pregs[preg.index()].is_stack = true;
         }
+        for class in 0..self.preferred_victim_by_class.len() {
+            self.preferred_victim_by_class[class] = self.env.non_preferred_regs_by_class[class]
+                .last()
+                .or(self.env.preferred_regs_by_class[class].last())
+                .cloned()
+                .unwrap_or(PReg::invalid());
+        }
         // Create VRegs from the vreg count.
         for idx in 0..self.func.num_vregs() {
             // We'll fill in the real details when we see the def.

src/ion/mod.rs

@@ -31,6 +31,7 @@ use liveranges::*;
 pub(crate) mod merge;
 pub(crate) mod process;
 use process::*;
+use smallvec::smallvec;
 pub(crate) mod dump;
 pub(crate) mod moves;
 pub(crate) mod spill;
@@ -66,7 +67,8 @@ impl<'a, F: Function> Env<'a, F> {
             slots_by_size: vec![],
             allocated_bundle_count: 0,
 
-            extra_spillslot: vec![None, None],
+            extra_spillslots_by_class: [smallvec![], smallvec![]],
+            preferred_victim_by_class: [PReg::invalid(), PReg::invalid()],
 
             prog_move_srcs: Vec::with_capacity(n / 2),
             prog_move_dsts: Vec::with_capacity(n / 2),

src/ion/moves.rs

@@ -16,12 +16,16 @@ use super::{
     Env, InsertMovePrio, InsertedMove, LiveRangeFlag, LiveRangeIndex, RedundantMoveEliminator,
     VRegIndex, SLOT_NONE,
 };
-use crate::ion::data_structures::{BlockparamIn, BlockparamOut, CodeRange, PosWithPrio};
-use crate::moves::ParallelMoves;
+use crate::ion::data_structures::{
+    BlockparamIn, BlockparamOut, CodeRange, LiveRangeKey, PosWithPrio,
+};
+use crate::ion::reg_traversal::RegTraversalIter;
+use crate::moves::{MoveAndScratchResolver, ParallelMoves};
 use crate::{
     Allocation, Block, Edit, Function, Inst, InstPosition, OperandConstraint, OperandKind,
-    OperandPos, PReg, ProgPoint, RegClass, VReg,
+    OperandPos, PReg, ProgPoint, RegClass, SpillSlot, VReg,
 };
+use fxhash::FxHashMap;
 use smallvec::{smallvec, SmallVec};
 use std::fmt::Debug;
 
@@ -965,8 +969,7 @@ impl<'a, F: Function> Env<'a, F> {
             // have two separate ParallelMove instances. They need to
             // be separate because moves between the two classes are
             // impossible. (We could enhance ParallelMoves to
-            // understand register classes and take multiple scratch
-            // regs, but this seems simpler.)
+            // understand register classes, but this seems simpler.)
             let mut int_moves: SmallVec<[InsertedMove; 8]> = smallvec![];
             let mut float_moves: SmallVec<[InsertedMove; 8]> = smallvec![];
 
@@ -993,8 +996,7 @@ impl<'a, F: Function> Env<'a, F> {
                 // All moves in `moves` semantically happen in
                 // parallel. Let's resolve these to a sequence of moves
                 // that can be done one at a time.
-                let scratch = self.env.scratch_by_class[regclass as u8 as usize];
-                let mut parallel_moves = ParallelMoves::new(Allocation::reg(scratch));
+                let mut parallel_moves = ParallelMoves::new();
                 trace!(
                     "parallel moves at pos {:?} prio {:?}",
                     pos_prio.pos,
@@ -1008,59 +1010,79 @@ impl<'a, F: Function> Env<'a, F> {
                 }
 
                 let resolved = parallel_moves.resolve();
-
-                // If (i) the scratch register is used, and (ii) a
-                // stack-to-stack move exists, then we need to
-                // allocate an additional scratch spillslot to which
-                // we can temporarily spill the scratch reg when we
-                // lower the stack-to-stack move to a
-                // stack-to-scratch-to-stack sequence.
-                let scratch_used = resolved.iter().any(|&(src, dst, _)| {
-                    src == Allocation::reg(scratch) || dst == Allocation::reg(scratch)
-                });
-                let stack_stack_move = resolved.iter().any(|&(src, dst, _)| {
-                    self.allocation_is_stack(src) && self.allocation_is_stack(dst)
+                let mut scratch_iter = RegTraversalIter::new(
+                    self.env,
+                    regclass,
+                    PReg::invalid(),
+                    PReg::invalid(),
+                    0,
+                    None,
+                );
+                let key = LiveRangeKey::from_range(&CodeRange {
+                    from: pos_prio.pos,
+                    to: pos_prio.pos.next(),
                 });
-                let extra_slot = if scratch_used && stack_stack_move {
-                    if self.extra_spillslot[regclass as u8 as usize].is_none() {
-                        let slot = self.allocate_spillslot(regclass);
-                        self.extra_spillslot[regclass as u8 as usize] = Some(slot);
+                let get_reg = || {
+                    while let Some(preg) = scratch_iter.next() {
+                        if !self.pregs[preg.index()]
+                            .allocations
+                            .btree
+                            .contains_key(&key)
+                        {
+                            let alloc = Allocation::reg(preg);
+                            if moves
+                                .iter()
+                                .any(|m| m.from_alloc == alloc || m.to_alloc == alloc)
+                            {
+                                // Skip pregs used by moves in this
+                                // parallel move set, even if not
+                                // marked used at progpoint: edge move
+                                // liveranges meet but don't overlap
+                                // so otherwise we may incorrectly
+                                // overwrite a source reg.
+                                continue;
+                            }
+                            return Some(alloc);
+                        }
                     }
-                    self.extra_spillslot[regclass as u8 as usize]
-                } else {
                     None
                 };
+                let mut stackslot_idx = 0;
+                let get_stackslot = || {
+                    let idx = stackslot_idx;
+                    stackslot_idx += 1;
+                    // We can't borrow `self` as mutable, so we create
+                    // these placeholders then allocate the actual
+                    // slots if needed with `self.allocate_spillslot`
+                    // below.
+                    Allocation::stack(SpillSlot::new(SpillSlot::MAX - idx, regclass))
+                };
+                let preferred_victim = self.preferred_victim_by_class[regclass as usize];
+
+                let scratch_resolver =
+                    MoveAndScratchResolver::new(get_reg, get_stackslot, preferred_victim);
+
+                let resolved = scratch_resolver.compute(resolved);
+
+                let mut rewrites = FxHashMap::default();
+                for i in 0..stackslot_idx {
+                    if i >= self.extra_spillslots_by_class[regclass as usize].len() {
+                        let slot = self.allocate_spillslot(regclass);
+                        self.extra_spillslots_by_class[regclass as usize].push(slot);
+                    }
+                    rewrites.insert(
+                        Allocation::stack(SpillSlot::new(SpillSlot::MAX - i, regclass)),
+                        self.extra_spillslots_by_class[regclass as usize][i],
+                    );
+                }
 
-                let mut scratch_used_yet = false;
                 for (src, dst, to_vreg) in resolved {
+                    let src = rewrites.get(&src).cloned().unwrap_or(src);
+                    let dst = rewrites.get(&dst).cloned().unwrap_or(dst);
                     trace!("  resolved: {} -> {} ({:?})", src, dst, to_vreg);
                     let action = redundant_moves.process_move(src, dst, to_vreg);
                     if !action.elide {
-                        if dst == Allocation::reg(scratch) {
-                            scratch_used_yet = true;
-                        }
-                        if self.allocation_is_stack(src) && self.allocation_is_stack(dst) {
-                            if !scratch_used_yet {
-                                self.add_move_edit(pos_prio, src, Allocation::reg(scratch));
-                                self.add_move_edit(pos_prio, Allocation::reg(scratch), dst);
-                            } else {
-                                debug_assert!(extra_slot.is_some());
-                                self.add_move_edit(
-                                    pos_prio,
-                                    Allocation::reg(scratch),
-                                    extra_slot.unwrap(),
-                                );
-                                self.add_move_edit(pos_prio, src, Allocation::reg(scratch));
-                                self.add_move_edit(pos_prio, Allocation::reg(scratch), dst);
-                                self.add_move_edit(
-                                    pos_prio,
-                                    extra_slot.unwrap(),
-                                    Allocation::reg(scratch),
-                                );
-                            }
-                        } else {
-                            self.add_move_edit(pos_prio, src, dst);
-                        }
+                        self.add_move_edit(pos_prio, src, dst);
                     } else {
                         trace!("    -> redundant move elided");
                     }
@@ -1081,7 +1103,7 @@ impl<'a, F: Function> Env<'a, F> {
                 let &(pos_prio, ref edit) = &self.edits[i];
                 match edit {
                     &Edit::Move { from, to } => {
-                        self.annotate(pos_prio.pos, format!("move {} -> {})", from, to));
+                        self.annotate(pos_prio.pos, format!("move {} -> {}", from, to));
                     }
                 }
             }