| /* |
| * Released under the terms of the Apache 2.0 license with LLVM |
| * exception. See `LICENSE` for details. |
| */ |
| |
| use crate::{ion::data_structures::u64_key, Allocation, PReg}; |
| use core::fmt::Debug; |
| use smallvec::{smallvec, SmallVec}; |
| |
| /// A list of moves to be performed in sequence, with auxiliary data |
| /// attached to each. |
| pub type MoveVec<T> = SmallVec<[(Allocation, Allocation, T); 16]>; |
| |
| /// A list of moves to be performance in sequence, like a |
| /// `MoveVec<T>`, except that an unchosen scratch space may occur as |
| /// well, represented by `Allocation::none()`. |
| #[derive(Clone, Debug)] |
| pub enum MoveVecWithScratch<T> { |
| /// No scratch was actually used. |
| NoScratch(MoveVec<T>), |
| /// A scratch space was used. |
| Scratch(MoveVec<T>), |
| } |
| |
| /// A `ParallelMoves` represents a list of alloc-to-alloc moves that |
| /// must happen in parallel -- i.e., all reads of sources semantically |
| /// happen before all writes of destinations, and destinations are |
| /// allowed to overwrite sources. It can compute a list of sequential |
| /// moves that will produce the equivalent data movement, possibly |
| /// using a scratch register if one is necessary. |
| pub struct ParallelMoves<T: Clone + Copy + Default> { |
| parallel_moves: MoveVec<T>, |
| } |
| |
| impl<T: Clone + Copy + Default + PartialEq> ParallelMoves<T> { |
| pub fn new() -> Self { |
| Self { |
| parallel_moves: smallvec![], |
| } |
| } |
| |
| pub fn add(&mut self, from: Allocation, to: Allocation, t: T) { |
| self.parallel_moves.push((from, to, t)); |
| } |
| |
| fn sources_overlap_dests(&self) -> bool { |
| // Assumes `parallel_moves` has already been sorted by `dst` |
| // in `resolve()` below. The O(n log n) cost of this loop is no |
| // worse than the sort we already did. |
| for &(src, _, _) in &self.parallel_moves { |
| if self |
| .parallel_moves |
| .binary_search_by_key(&src, |&(_, dst, _)| dst) |
| .is_ok() |
| { |
| return true; |
| } |
| } |
| false |
| } |
| |
| /// Resolve the parallel-moves problem to a sequence of separate |
| /// moves, such that the combined effect of the sequential moves |
| /// is as-if all of the moves added to this `ParallelMoves` |
| /// resolver happened in parallel. |
| /// |
| /// Sometimes, if there is a cycle, a scratch register is |
| /// necessary to allow the moves to occur sequentially. In this |
| /// case, `Allocation::none()` is returned to represent the |
| /// scratch register. The caller may choose to always hold a |
| /// separate scratch register unused to allow this to be trivially |
| /// rewritten; or may dynamically search for or create a free |
| /// register as needed, if none are available. |
| pub fn resolve(mut self) -> MoveVecWithScratch<T> { |
| // Easy case: zero or one move. Just return our vec. |
| if self.parallel_moves.len() <= 1 { |
| return MoveVecWithScratch::NoScratch(self.parallel_moves); |
| } |
| |
| // Sort moves so that we can efficiently test for presence. |
| // For that purpose it doesn't matter whether we sort by |
| // source or destination, but later we'll want them sorted |
| // by destination. |
| self.parallel_moves |
| .sort_by_key(|&(src, dst, _)| u64_key(dst.bits(), src.bits())); |
| |
| // Duplicate moves cannot change the semantics of this |
| // parallel move set, so remove them. This is cheap since we |
| // just sorted the list. |
| self.parallel_moves.dedup(); |
| |
| // General case: some moves overwrite dests that other moves |
| // read as sources. We'll use a general algorithm. |
| // |
| // *Important property*: because we expect that each register |
| // has only one writer (otherwise the effect of the parallel |
| // move is undefined), each move can only block one other move |
| // (with its one source corresponding to the one writer of |
| // that source). Thus, we *can only have simple cycles* (those |
| // that are a ring of nodes, i.e., with only one path from a |
| // node back to itself); there are no SCCs that are more |
| // complex than that. We leverage this fact below to avoid |
| // having to do a full Tarjan SCC DFS (with lowest-index |
| // computation, etc.): instead, as soon as we find a cycle, we |
| // know we have the full cycle and we can do a cyclic move |
| // sequence and continue. |
| |
| // Check that each destination has only one writer. |
| if cfg!(debug_assertions) { |
| let mut last_dst = None; |
| for &(_, dst, _) in &self.parallel_moves { |
| if last_dst.is_some() { |
| debug_assert!(last_dst.unwrap() != dst); |
| } |
| last_dst = Some(dst); |
| } |
| } |
| |
| // Moving an allocation into itself is technically a cycle but |
| // should have no effect, as long as there are no other writes |
| // into that destination. |
| self.parallel_moves.retain(|&mut (src, dst, _)| src != dst); |
| |
| // Do any dests overlap sources? If not, we can also just |
| // return the list. |
| if !self.sources_overlap_dests() { |
| return MoveVecWithScratch::NoScratch(self.parallel_moves); |
| } |
| |
| // Construct a mapping from move indices to moves they must |
| // come before. Any given move must come before a move that |
| // overwrites its destination; we have moves sorted by dest |
| // above so we can efficiently find such a move, if any. |
| const NONE: usize = usize::MAX; |
| let must_come_before: SmallVec<[usize; 16]> = self |
| .parallel_moves |
| .iter() |
| .map(|&(src, _, _)| { |
| self.parallel_moves |
| .binary_search_by_key(&src, |&(_, dst, _)| dst) |
| .unwrap_or(NONE) |
| }) |
| .collect(); |
| |
| // Do a simple stack-based DFS and emit moves in postorder, |
| // then reverse at the end for RPO. Unlike Tarjan's SCC |
| // algorithm, we can emit a cycle as soon as we find one, as |
| // noted above. |
| #[derive(Clone, Copy, Debug, Eq, PartialEq)] |
| enum State { |
| /// Not on stack, not visited |
| ToDo, |
| /// On stack, not yet visited |
| Pending, |
| /// Visited |
| Done, |
| } |
| let mut ret: MoveVec<T> = smallvec![]; |
| let mut stack: SmallVec<[usize; 16]> = smallvec![]; |
| let mut state: SmallVec<[State; 16]> = smallvec![State::ToDo; self.parallel_moves.len()]; |
| let mut scratch_used = false; |
| |
| while let Some(next) = state.iter().position(|&state| state == State::ToDo) { |
| stack.push(next); |
| state[next] = State::Pending; |
| |
| while let Some(&top) = stack.last() { |
| debug_assert_eq!(state[top], State::Pending); |
| let next = must_come_before[top]; |
| if next == NONE || state[next] == State::Done { |
| ret.push(self.parallel_moves[top]); |
| state[top] = State::Done; |
| stack.pop(); |
| while let Some(top) = stack.pop() { |
| ret.push(self.parallel_moves[top]); |
| state[top] = State::Done; |
| } |
| } else if state[next] == State::ToDo { |
| stack.push(next); |
| state[next] = State::Pending; |
| } else { |
| // Found a cycle -- emit a cyclic-move sequence |
| // for the cycle on the top of stack, then normal |
| // moves below it. Recall that these moves will be |
| // reversed in sequence, so from the original |
| // parallel move set |
| // |
| // { B := A, C := B, A := B } |
| // |
| // we will generate something like: |
| // |
| // A := scratch |
| // B := A |
| // C := B |
| // scratch := C |
| // |
| // which will become: |
| // |
| // scratch := C |
| // C := B |
| // B := A |
| // A := scratch |
| debug_assert_ne!(top, next); |
| state[top] = State::Done; |
| stack.pop(); |
| |
| let (scratch_src, dst, dst_t) = self.parallel_moves[top]; |
| scratch_used = true; |
| |
| ret.push((Allocation::none(), dst, dst_t)); |
| while let Some(move_idx) = stack.pop() { |
| state[move_idx] = State::Done; |
| ret.push(self.parallel_moves[move_idx]); |
| |
| if move_idx == next { |
| break; |
| } |
| } |
| ret.push((scratch_src, Allocation::none(), T::default())); |
| } |
| } |
| } |
| |
| ret.reverse(); |
| |
| if scratch_used { |
| MoveVecWithScratch::Scratch(ret) |
| } else { |
| MoveVecWithScratch::NoScratch(ret) |
| } |
| } |
| } |
| |
| impl<T> MoveVecWithScratch<T> { |
| /// Fills in the scratch space, if needed, with the given |
| /// register/allocation and returns a final list of moves. The |
| /// scratch register must not occur anywhere in the parallel-move |
| /// problem given to the resolver that produced this |
| /// `MoveVecWithScratch`. |
| pub fn with_scratch(self, scratch: Allocation) -> MoveVec<T> { |
| match self { |
| MoveVecWithScratch::NoScratch(moves) => moves, |
| MoveVecWithScratch::Scratch(mut moves) => { |
| for (src, dst, _) in &mut moves { |
| debug_assert!( |
| *src != scratch && *dst != scratch, |
| "Scratch register should not also be an actual source or dest of moves" |
| ); |
| debug_assert!( |
| !(src.is_none() && dst.is_none()), |
| "Move resolution should not have produced a scratch-to-scratch move" |
| ); |
| if src.is_none() { |
| *src = scratch; |
| } |
| if dst.is_none() { |
| *dst = scratch; |
| } |
| } |
| moves |
| } |
| } |
| } |
| |
| /// Unwrap without a scratch register. |
| pub fn without_scratch(self) -> Option<MoveVec<T>> { |
| match self { |
| MoveVecWithScratch::NoScratch(moves) => Some(moves), |
| MoveVecWithScratch::Scratch(..) => None, |
| } |
| } |
| |
| /// Do we need a scratch register? |
| pub fn needs_scratch(&self) -> bool { |
| match self { |
| MoveVecWithScratch::NoScratch(..) => false, |
| MoveVecWithScratch::Scratch(..) => true, |
| } |
| } |
| } |
| |
| /// Final stage of move resolution: finding or using scratch |
| /// registers, creating them if necessary by using stackslots, and |
| /// ensuring that the final list of moves contains no stack-to-stack |
| /// moves. |
| /// |
| /// The resolved list of moves may need one or two scratch registers, |
| /// and maybe a stackslot, to ensure these conditions. Our general |
| /// strategy is in two steps. |
| /// |
| /// First, we find a scratch register, so we only have to worry about |
| /// a list of moves, all with real locations as src and dest. If we're |
| /// lucky and there are any registers not allocated at this |
| /// program-point, we can use a real register. Otherwise, we use an |
| /// extra stackslot. This is fine, because at this step, |
| /// stack-to-stack moves are OK. |
| /// |
| /// 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 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, |
| { |
| /// Closure that finds us a PReg at the current location. |
| pub find_free_reg: GetReg, |
| /// Closure that gets us a stackslot, if needed. |
| pub get_stackslot: GetStackSlot, |
| /// Closure to determine whether an `Allocation` refers to a stack slot. |
| 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> |
| where |
| GetReg: FnMut() -> Option<Allocation>, |
| GetStackSlot: FnMut() -> Allocation, |
| IsStackAlloc: Fn(Allocation) -> bool, |
| { |
| 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() |
| }; |
| |
| // 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; |
| } |
| |
| // 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 mut result = smallvec![]; |
| for &(src, dst, data) in &moves { |
| // Do we have a stack-to-stack move? If so, resolve. |
| if self.is_stack_to_stack_move(src, dst) { |
| trace!("scratch resolver: stack to stack: {:?} -> {:?}", src, dst); |
| |
| // 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; |
| } |
| } |
| |
| // 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; |
| } |
| 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; |
| } |
| 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) |
| } |
| } |
