Simplify stack-to-stack moves (#159)

* Simplify parallel stack-to-stack moves

Most of the changes here are simply refactoring.

Resolving stack-to-stack moves first fills in the scratch register for
the parallel moves, if one is needed. So we can split that into two
phases, inlining the `stack_to_stack` function which was not used
anywhere else. That lets us establish the invariant that we only
continue into the second phase if there is at least one stack-to-stack
move.

Because `compute` consumes `self`, we know that the lazily-allocated
stack-to-stack scratch register has not been allocated yet at entry to
the loop, but because the loop only runs if there are stack-to-stack
moves, we know the scratch register will definitely be allocated. So we
can remove the corresponding state from the `MoveVecWithScratch` struct
and keep it in locals with shorter names instead, eagerly initialized
before the start of the loop.

The above changes should have no effect on the output. However I then
added state tracking for whether the scratch register or the save slot
already contain the right value and therefore don't need to be the
target of another move. This should remove the need for the
redundant-move eliminator to clean up after parallel moves.

My first approach assumed the correct value was in exactly one of these
places, but both allocations can be correct, and we can save more copies
by acknowledging that.

* Remove MoveAndScratchResolver::new

This struct now contains exactly the fields whose values are passed to
`new`. I feel it's better API design in this case to use the named
fields than to expect callers to keep track of the positional arguments
to the `new` method.

* Clarify relationship between save_slot/scratch_dirty

* Review comments: factor out is_stack_to_stack_move

* Remove "victim" terminology from public API

Also update more comments I missed.

Author: jameysharp

fuzz/fuzz_targets/moves.rs

@@ -83,16 +83,20 @@ fuzz_target!(|testcase: TestCase| {
     // 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 find_free_reg = || avail.pop();
     let mut next_slot = 32;
     let get_stackslot = || {
         let slot = next_slot;
         next_slot += 1;
         Allocation::stack(SpillSlot::new(slot))
     };
     let preferred_victim = PReg::new(0, RegClass::Int);
-    let scratch_resolver =
-        MoveAndScratchResolver::new(get_reg, get_stackslot, is_stack_alloc, preferred_victim);
+    let scratch_resolver = MoveAndScratchResolver {
+        find_free_reg,
+        get_stackslot,
+        is_stack_alloc,
+        borrowed_scratch_reg: preferred_victim,
+    };
     let moves = scratch_resolver.compute(moves);
     log::trace!("resolved moves: {:?}", moves);
 

src/ion/moves.rs

@@ -909,7 +909,7 @@ impl<'a, F: Function> Env<'a, F> {
                     from: pos_prio.pos,
                     to: pos_prio.pos.next(),
                 });
-                let get_reg = || {
+                let find_free_reg = || {
                     // Use the dedicated scratch register first if it is
                     // available.
                     if let Some(reg) = dedicated_scratch.take() {
@@ -958,12 +958,12 @@ impl<'a, F: Function> Env<'a, F> {
                 };
                 let preferred_victim = self.preferred_victim_by_class[regclass as usize];
 
-                let scratch_resolver = MoveAndScratchResolver::new(
-                    get_reg,
+                let scratch_resolver = MoveAndScratchResolver {
+                    find_free_reg,
                     get_stackslot,
                     is_stack_alloc,
-                    preferred_victim,
-                );
+                    borrowed_scratch_reg: preferred_victim,
+                };
 
                 let resolved = scratch_resolver.compute(resolved);
 

src/moves.rs

@@ -277,15 +277,6 @@ impl<T> MoveVecWithScratch<T> {
             MoveVecWithScratch::Scratch(..) => true,
         }
     }
-
-    /// Do any moves go from stack to stack?
-    pub fn stack_to_stack(&self, is_stack_alloc: impl Fn(Allocation) -> bool) -> bool {
-        match self {
-            MoveVecWithScratch::NoScratch(moves) | MoveVecWithScratch::Scratch(moves) => moves
-                .iter()
-                .any(|&(src, dst, _)| is_stack_alloc(src) && is_stack_alloc(dst)),
-        }
-    }
 }
 
 /// Final stage of move resolution: finding or using scratch
@@ -306,37 +297,28 @@ impl<T> MoveVecWithScratch<T> {
 ///
 /// Then, we resolve stack-to-stack moves into stack-to-reg /
 /// reg-to-stack pairs. For this, we try to allocate a second free
-/// register. If unavailable, we create another scratch stackslot, and
-/// we pick a "victim" register in the appropriate class, and we
-/// resolve into: victim -> extra-stackslot; stack-src -> victim;
-/// victim -> stack-dst; extra-stackslot -> victim.
-///
-/// Sometimes move elision will be able to clean this up a bit. But,
-/// for simplicity reasons, let's keep the concerns separated! So we
-/// always do the full expansion above.
+/// register. If unavailable, we create a new scratch stackslot to
+/// serve as a backup of one of the in-use registers, then borrow that
+/// register as the scratch register in the middle of stack-to-stack
+/// moves.
 pub struct MoveAndScratchResolver<GetReg, GetStackSlot, IsStackAlloc>
 where
     GetReg: FnMut() -> Option<Allocation>,
     GetStackSlot: FnMut() -> Allocation,
     IsStackAlloc: Fn(Allocation) -> bool,
 {
-    /// Scratch register for stack-to-stack move expansion.
-    stack_stack_scratch_reg: Option<Allocation>,
-    /// Stackslot into which we need to save the stack-to-stack
-    /// scratch reg before doing any stack-to-stack moves, if we stole
-    /// the reg.
-    stack_stack_scratch_reg_save: Option<Allocation>,
     /// Closure that finds us a PReg at the current location.
-    find_free_reg: GetReg,
+    pub find_free_reg: GetReg,
     /// Closure that gets us a stackslot, if needed.
-    get_stackslot: GetStackSlot,
+    pub get_stackslot: GetStackSlot,
     /// Closure to determine whether an `Allocation` refers to a stack slot.
-    is_stack_alloc: IsStackAlloc,
-    /// The victim PReg to evict to another stackslot at every
-    /// stack-to-stack move if a free PReg is not otherwise
-    /// available. Provided by caller and statically chosen. This is a
-    /// very last-ditch option, so static choice is OK.
-    victim: PReg,
+    pub is_stack_alloc: IsStackAlloc,
+    /// Use this register if no free register is available to use as a
+    /// temporary in stack-to-stack moves. If we do use this register
+    /// for that purpose, its value will be restored by the end of the
+    /// move sequence. Provided by caller and statically chosen. This is
+    /// a very last-ditch option, so static choice is OK.
+    pub borrowed_scratch_reg: PReg,
 }
 
 impl<GetReg, GetStackSlot, IsStackAlloc> MoveAndScratchResolver<GetReg, GetStackSlot, IsStackAlloc>
@@ -345,88 +327,129 @@ where
     GetStackSlot: FnMut() -> Allocation,
     IsStackAlloc: Fn(Allocation) -> bool,
 {
-    pub fn new(
-        find_free_reg: GetReg,
-        get_stackslot: GetStackSlot,
-        is_stack_alloc: IsStackAlloc,
-        victim: PReg,
-    ) -> Self {
-        Self {
-            stack_stack_scratch_reg: None,
-            stack_stack_scratch_reg_save: None,
-            find_free_reg,
-            get_stackslot,
-            is_stack_alloc,
-            victim,
-        }
-    }
+    pub fn compute<T: Debug + Default + Copy>(
+        mut self,
+        moves: MoveVecWithScratch<T>,
+    ) -> MoveVec<T> {
+        let moves = if moves.needs_scratch() {
+            // Now, find a scratch allocation in order to resolve cycles.
+            let scratch = (self.find_free_reg)().unwrap_or_else(|| (self.get_stackslot)());
+            trace!("scratch resolver: scratch alloc {:?}", scratch);
+
+            moves.with_scratch(scratch)
+        } else {
+            moves.without_scratch().unwrap()
+        };
 
-    pub fn compute<T: Debug + Copy>(mut self, moves: MoveVecWithScratch<T>) -> MoveVec<T> {
-        // First, do we have a vec with no stack-to-stack moves or use
-        // of a scratch register? Fast return if so.
-        if !moves.needs_scratch() && !moves.stack_to_stack(&self.is_stack_alloc) {
-            return moves.without_scratch().unwrap();
+        // Do we have any stack-to-stack moves? Fast return if not.
+        let stack_to_stack = moves
+            .iter()
+            .any(|&(src, dst, _)| self.is_stack_to_stack_move(src, dst));
+        if !stack_to_stack {
+            return moves;
         }
 
-        let mut result = smallvec![];
-
-        // Now, find a scratch allocation in order to resolve cycles.
-        let scratch = (self.find_free_reg)().unwrap_or_else(|| (self.get_stackslot)());
-        trace!("scratch resolver: scratch alloc {:?}", scratch);
+        // Allocate a scratch register for stack-to-stack move expansion.
+        let (scratch_reg, save_slot) = if let Some(reg) = (self.find_free_reg)() {
+            trace!(
+                "scratch resolver: have free stack-to-stack scratch preg: {:?}",
+                reg
+            );
+            (reg, None)
+        } else {
+            let reg = Allocation::reg(self.borrowed_scratch_reg);
+            // Stackslot into which we need to save the stack-to-stack
+            // scratch reg before doing any stack-to-stack moves, if we stole
+            // the reg.
+            let save = (self.get_stackslot)();
+            trace!(
+                "scratch resolver: stack-to-stack borrowing {:?} with save stackslot {:?}",
+                reg,
+                save
+            );
+            (reg, Some(save))
+        };
+
+        // Mutually exclusive flags for whether either scratch_reg or
+        // save_slot need to be restored from the other. Initially,
+        // scratch_reg has a value we should preserve and save_slot
+        // has garbage.
+        let mut scratch_dirty = false;
+        let mut save_dirty = true;
 
-        let moves = moves.with_scratch(scratch);
+        let mut result = smallvec![];
         for &(src, dst, data) in &moves {
             // Do we have a stack-to-stack move? If so, resolve.
-            if (self.is_stack_alloc)(src) && (self.is_stack_alloc)(dst) {
+            if self.is_stack_to_stack_move(src, dst) {
                 trace!("scratch resolver: stack to stack: {:?} -> {:?}", src, dst);
-                // Lazily allocate a stack-to-stack scratch.
-                if self.stack_stack_scratch_reg.is_none() {
-                    if let Some(reg) = (self.find_free_reg)() {
-                        trace!(
-                            "scratch resolver: have free stack-to-stack scratch preg: {:?}",
-                            reg
-                        );
-                        self.stack_stack_scratch_reg = Some(reg);
-                    } else {
-                        self.stack_stack_scratch_reg = Some(Allocation::reg(self.victim));
-                        self.stack_stack_scratch_reg_save = Some((self.get_stackslot)());
-                        trace!("scratch resolver: stack-to-stack using victim {:?} with save stackslot {:?}",
-                                    self.stack_stack_scratch_reg,
-                                    self.stack_stack_scratch_reg_save);
+
+                // If the selected scratch register is stolen from the
+                // set of in-use registers, then we need to save the
+                // current contents of the scratch register before using
+                // it as a temporary.
+                if let Some(save_slot) = save_slot {
+                    // However we may have already done so for an earlier
+                    // stack-to-stack move in which case we don't need
+                    // to do it again.
+                    if save_dirty {
+                        debug_assert!(!scratch_dirty);
+                        result.push((scratch_reg, save_slot, T::default()));
+                        save_dirty = false;
                     }
                 }
 
-                // If we have a "victimless scratch", then do a
-                // stack-to-scratch / scratch-to-stack sequence.
-                if self.stack_stack_scratch_reg_save.is_none() {
-                    result.push((src, self.stack_stack_scratch_reg.unwrap(), data));
-                    result.push((self.stack_stack_scratch_reg.unwrap(), dst, data));
+                // We can't move directly from one stack slot to another
+                // on any architecture we care about, so stack-to-stack
+                // moves must go via a scratch register.
+                result.push((src, scratch_reg, data));
+                result.push((scratch_reg, dst, data));
+                scratch_dirty = true;
+            } else {
+                // This is not a stack-to-stack move, but we need to
+                // make sure that the scratch register is in the correct
+                // state if this move interacts with that register.
+                if src == scratch_reg && scratch_dirty {
+                    // We're copying from the scratch register so if
+                    // it was stolen for a stack-to-stack move then we
+                    // need to make sure it has the correct contents,
+                    // not whatever was temporarily copied into it. If
+                    // we got scratch_reg from find_free_reg then it
+                    // had better not have been used as the source of
+                    // a move. So if we're here it's because we fell
+                    // back to the caller-provided last-resort scratch
+                    // register, and we must therefore have a save-slot
+                    // allocated too.
+                    debug_assert!(!save_dirty);
+                    let save_slot = save_slot.expect("move source should not be a free register");
+                    result.push((save_slot, scratch_reg, T::default()));
+                    scratch_dirty = false;
                 }
-                // Otherwise, save the current value in the
-                // stack-to-stack scratch reg (which is our victim) to
-                // the extra stackslot, then do the stack-to-scratch /
-                // scratch-to-stack sequence, then restore it.
-                else {
-                    result.push((
-                        self.stack_stack_scratch_reg.unwrap(),
-                        self.stack_stack_scratch_reg_save.unwrap(),
-                        data,
-                    ));
-                    result.push((src, self.stack_stack_scratch_reg.unwrap(), data));
-                    result.push((self.stack_stack_scratch_reg.unwrap(), dst, data));
-                    result.push((
-                        self.stack_stack_scratch_reg_save.unwrap(),
-                        self.stack_stack_scratch_reg.unwrap(),
-                        data,
-                    ));
+                if dst == scratch_reg {
+                    // We are writing something to the scratch register
+                    // so it doesn't matter what was there before. We
+                    // can avoid restoring it, but we will need to save
+                    // it again before the next stack-to-stack move.
+                    scratch_dirty = false;
+                    save_dirty = true;
                 }
-            } else {
-                // Normal move.
                 result.push((src, dst, data));
             }
         }
 
+        // Now that all the stack-to-stack moves are done, restore the
+        // scratch register if necessary.
+        if let Some(save_slot) = save_slot {
+            if scratch_dirty {
+                debug_assert!(!save_dirty);
+                result.push((save_slot, scratch_reg, T::default()));
+            }
+        }
+
         trace!("scratch resolver: got {:?}", result);
         result
     }
+
+    fn is_stack_to_stack_move(&self, src: Allocation, dst: Allocation) -> bool {
+        (self.is_stack_alloc)(src) && (self.is_stack_alloc)(dst)
+    }
 }