From 16ac4f65b3dc3437eb1caa5b8751ec939409312d Mon Sep 17 00:00:00 2001 From: Jakob Stoklund Olesen Date: Mon, 8 Jan 2018 13:29:53 -0800 Subject: [PATCH] Add support for textbook union-find to VirtRegs. The initial phase of computing virtual registers can now be implemented with a textbook union-find algorithm using a disjoint set forest complete with rank and path compression optimizations. The disjoint set forest is converted to virtual register value lists in a single linear scan implemented in finish_union_find(). This union-find algorithm will soon be used by the coalescer. --- lib/cretonne/src/regalloc/virtregs.rs | 335 +++++++++++++++++++++++++- 1 file changed, 333 insertions(+), 2 deletions(-) diff --git a/lib/cretonne/src/regalloc/virtregs.rs b/lib/cretonne/src/regalloc/virtregs.rs index b9388277e0..ce38dc3687 100644 --- a/lib/cretonne/src/regalloc/virtregs.rs +++ b/lib/cretonne/src/regalloc/virtregs.rs @@ -11,11 +11,15 @@ //! If any values in a virtual register are spilled, they will use the same stack slot. This avoids //! memory-to-memory copies when a spilled value is passed as an EBB argument. +use dbg::DisplayList; use entity::{EntityList, ListPool}; use entity::{PrimaryMap, EntityMap, Keys}; +use entity::EntityRef; use ir::Value; use packed_option::PackedOption; use ref_slice::ref_slice; +use std::cmp::Ordering; +use std::fmt; /// A virtual register reference. #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] @@ -35,26 +39,41 @@ pub struct VirtRegs { /// The primary table of virtual registers. vregs: PrimaryMap, + /// Allocated virtual register numbers that are no longer in use. + unused_vregs: Vec, + /// Each value belongs to at most one virtual register. value_vregs: EntityMap>, + + /// Table used during the union-find phase while `vregs` is empty. + union_find: EntityMap, + + /// Values that have been activated in the `union_find` table, but not yet added to any virtual + /// registers by the `finish_union_find()` function. + pending_values: Vec, } -#[allow(dead_code)] impl VirtRegs { /// Create a new virtual register collection. pub fn new() -> Self { Self { pool: ListPool::new(), vregs: PrimaryMap::new(), + unused_vregs: Vec::new(), value_vregs: EntityMap::new(), + union_find: EntityMap::new(), + pending_values: Vec::new(), } } /// Clear all virtual registers. pub fn clear(&mut self) { self.vregs.clear(); + self.unused_vregs.clear(); self.value_vregs.clear(); self.pool.clear(); + self.union_find.clear(); + self.pending_values.clear(); } /// Get the virtual register containing `value`, if any. @@ -63,7 +82,6 @@ impl VirtRegs { } /// Get the list of values in `vreg`. - /// The values are topologically ordered according dominance of their definition points. pub fn values(&self, vreg: VirtReg) -> &[Value] { self.vregs[vreg].as_slice(&self.pool) } @@ -94,6 +112,44 @@ impl VirtRegs { } } + /// Sort the values in `vreg` according to `compare`. + /// + /// If the ordering defined by `compare` is not total, value numbers are used as a last resort + /// tie-breaker. This makes it possible to use an unstable sorting algorithm which can be + /// faster because it doesn't allocate memory. + /// + /// Returns the slice of sorted values which `values(vreg)` will also return from now on. + pub fn sort_values(&mut self, vreg: VirtReg, mut compare: F) -> &[Value] + where + F: FnMut(Value, Value) -> Ordering, + { + let s = self.vregs[vreg].as_mut_slice(&mut self.pool); + s.sort_unstable_by(|&a, &b| compare(a, b).then(a.cmp(&b))); + s + } + + /// Remove a virtual register. + /// + /// The values in `vreg` become singletons, and the virtual register number may be reused in + /// the future. + pub fn remove(&mut self, vreg: VirtReg) { + // Start by reassigning all the values. + for &v in self.vregs[vreg].as_slice(&self.pool) { + let old = self.value_vregs[v].take(); + debug_assert_eq!(old, Some(vreg)); + } + + self.vregs[vreg].clear(&mut self.pool); + self.unused_vregs.push(vreg); + } + + /// Allocate a new empty virtual register. + fn alloc(&mut self) -> VirtReg { + self.unused_vregs.pop().unwrap_or_else(|| { + self.vregs.push(Default::default()) + }) + } + /// Unify `values` into a single virtual register. /// /// The values in the slice can be singletons or they can belong to a virtual register already. @@ -138,3 +194,278 @@ impl VirtRegs { vreg } } + +impl fmt::Display for VirtRegs { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + for vreg in self.all_virtregs() { + write!(f, "\n{} = {}", vreg, DisplayList(self.values(vreg)))?; + } + Ok(()) + } +} + +/// Expanded version of a union-find table entry. +enum UFEntry { + /// This value is a a set leader. The embedded number is the set's rank. + Rank(u32), + + /// This value belongs to the same set as the linked value. + Link(Value), +} + +/// The `union_find` table contains `i32` entries that are interpreted as follows: +/// +/// x = 0: The value belongs to its own singleton set. +/// x > 0: The value is the leader of a set with rank x. +/// x < 0: The value belongs to the same set as the value numbered !x. +/// +/// The rank of a set is an upper bound on the number of links that must be followed from a member +/// of the set to the set leader. +/// +/// A singleton set is the same as a set with rank 0. It contains only the leader value. +impl UFEntry { + /// Decode a table entry. + fn decode(x: i32) -> UFEntry { + if x < 0 { + UFEntry::Link(Value::new((!x) as usize)) + } else { + UFEntry::Rank(x as u32) + } + } + + /// Encode a link entry. + fn encode_link(v: Value) -> i32 { + !(v.index() as i32) + } +} + +/// Union-find algorithm for building virtual registers. +/// +/// Before values are added to virtual registers, it is possible to use a union-find algorithm to +/// construct virtual registers efficiently. This support implemented here is used as follows: +/// +/// 1. Repeatedly call the `union(a, b)` method to request that `a` and `b` are placed in the same +/// virtual register. +/// 2. When done, call `finish_union_find()` to construct the virtual register sets based on the +/// `union()` calls. +/// +/// The values that were passed to `union(a, b)` mist not belong to any existing virtual registers +/// by the time `finish_union_find()` is called. +/// +/// For more information on the algorithm implemented here, see Chapter 21 "Data Structures for +/// Disjoint Sets" of Cormen, Leiserson, Rivest, Stein, "Introduction to algorithms", 3rd Ed. +/// +/// The [Wikipedia entry on disjoint-set data +/// structures](https://en.wikipedia.org/wiki/Disjoint-set_data_structure) is also good. +impl VirtRegs { + /// Find the leader value and rank of the set containing `v`. + /// Compress the path if needed. + fn find(&mut self, val: Value) -> (Value, u32) { + match UFEntry::decode(self.union_find[val]) { + UFEntry::Rank(rank) => (val, rank), + UFEntry::Link(parent) => { + // TODO: This recursion would be more efficient as an iteration that pushes + // elements onto a SmallVector. + let found = self.find(parent); + // Compress the path if needed. + if found.0 != parent { + self.union_find[val] = UFEntry::encode_link(found.0); + } + found + } + } + } + + /// Union the two sets containing `a` and `b`. + /// + /// This ensures that `a` and `b` will belong to the same virtual register after calling + /// `finish_union_find()`. + pub fn union(&mut self, a: Value, b: Value) { + let (leader_a, rank_a) = self.find(a); + let (leader_b, rank_b) = self.find(b); + + if leader_a == leader_b { + return; + } + + // The first time we see a value, its rank will be 0. Add it to the list of pending values. + if rank_a == 0 { + debug_assert_eq!(a, leader_a); + self.pending_values.push(a); + } + if rank_b == 0 { + debug_assert_eq!(b, leader_b); + self.pending_values.push(b); + } + + // Merge into the set with the greater rank. This preserves the invariant that the rank is + // an upper bound on the number of links to the leader. + match rank_a.cmp(&rank_b) { + Ordering::Less => { + self.union_find[leader_a] = UFEntry::encode_link(leader_b); + } + Ordering::Greater => { + self.union_find[leader_b] = UFEntry::encode_link(leader_a); + } + Ordering::Equal => { + // When the two sets have the same rank, we arbitrarily pick the a-set to preserve. + // We need to increase the rank by one since the elements in the b-set are now one + // link further away from the leader. + self.union_find[leader_a] += 1; + self.union_find[leader_b] = UFEntry::encode_link(leader_a); + } + } + } + + /// Compute virtual registers based on previous calls to `union(a, b)`. + /// + /// This terminates the union-find algorithm, so the next time `union()` is called, it is for a + /// new independent batch of values. + /// + /// The values in each virtual register will be ordered according to when they were first + /// passed to `union()`, but backwards. It is expected that `sort_values()` will be used to + /// create a more sensible value order. + /// + /// The new virtual registers will be appended to `new_vregs`, if present. + pub fn finish_union_find(&mut self, mut new_vregs: Option<&mut Vec>) { + debug_assert_eq!( + self.pending_values.iter().find(|&&v| self.get(v).is_some()), + None, + "Values participating in union-find must not belong to existing virtual registers" + ); + + while let Some(val) = self.pending_values.pop() { + let (leader, _) = self.find(val); + + // Get the vreg for `leader`, or create it. + let vreg = self.get(leader).unwrap_or_else(|| { + // Allocate a vreg for `leader`, but leave it empty. + let vr = self.alloc(); + if let &mut Some(ref mut vec) = &mut new_vregs { + vec.push(vr); + } + self.value_vregs[leader] = vr.into(); + vr + }); + + // Push values in `pending_values` order, including when `v == leader`. + self.vregs[vreg].push(val, &mut self.pool); + self.value_vregs[val] = vreg.into(); + + // Clear the entry in the union-find table. The `find(val)` call may still look at this + // entry in a future iteration, but that it ok. It will return a rank 0 leader that has + // already been assigned to the correct virtual register. + self.union_find[val] = 0; + } + + // We do *not* call `union_find.clear()` table here because re-initializing the table for + // sparse use takes time linear in the number of values in the function. Instead we reset + // the entries that are known to be non-zero in the loop above. + } +} + +#[cfg(test)] +mod test { + use super::*; + use entity::EntityRef; + use ir::Value; + + #[test] + fn empty_union_find() { + let mut vregs = VirtRegs::new(); + vregs.finish_union_find(None); + assert_eq!(vregs.all_virtregs().count(), 0); + } + + #[test] + fn union_self() { + let mut vregs = VirtRegs::new(); + let v1 = Value::new(1); + vregs.union(v1, v1); + vregs.finish_union_find(None); + assert_eq!(vregs.get(v1), None); + assert_eq!(vregs.all_virtregs().count(), 0); + } + + #[test] + fn union_pair() { + let mut vregs = VirtRegs::new(); + let v1 = Value::new(1); + let v2 = Value::new(2); + vregs.union(v1, v2); + vregs.finish_union_find(None); + assert_eq!(vregs.congruence_class(&v1), &[v2, v1]); + assert_eq!(vregs.congruence_class(&v2), &[v2, v1]); + assert_eq!(vregs.all_virtregs().count(), 1); + } + + #[test] + fn union_pair_backwards() { + let mut vregs = VirtRegs::new(); + let v1 = Value::new(1); + let v2 = Value::new(2); + vregs.union(v2, v1); + vregs.finish_union_find(None); + assert_eq!(vregs.congruence_class(&v1), &[v1, v2]); + assert_eq!(vregs.congruence_class(&v2), &[v1, v2]); + assert_eq!(vregs.all_virtregs().count(), 1); + } + + #[test] + fn union_tree() { + let mut vregs = VirtRegs::new(); + let v1 = Value::new(1); + let v2 = Value::new(2); + let v3 = Value::new(3); + let v4 = Value::new(4); + + vregs.union(v2, v4); + vregs.union(v3, v1); + // Leaders: v2, v3 + vregs.union(v4, v1); + vregs.finish_union_find(None); + assert_eq!(vregs.congruence_class(&v1), &[v1, v3, v4, v2]); + assert_eq!(vregs.congruence_class(&v2), &[v1, v3, v4, v2]); + assert_eq!(vregs.congruence_class(&v3), &[v1, v3, v4, v2]); + assert_eq!(vregs.congruence_class(&v4), &[v1, v3, v4, v2]); + assert_eq!(vregs.all_virtregs().count(), 1); + } + + #[test] + fn union_two() { + let mut vregs = VirtRegs::new(); + let v1 = Value::new(1); + let v2 = Value::new(2); + let v3 = Value::new(3); + let v4 = Value::new(4); + + vregs.union(v2, v4); + vregs.union(v3, v1); + // Leaders: v2, v3 + vregs.finish_union_find(None); + assert_eq!(vregs.congruence_class(&v1), &[v1, v3]); + assert_eq!(vregs.congruence_class(&v2), &[v4, v2]); + assert_eq!(vregs.congruence_class(&v3), &[v1, v3]); + assert_eq!(vregs.congruence_class(&v4), &[v4, v2]); + assert_eq!(vregs.all_virtregs().count(), 2); + } + + #[test] + fn union_uneven() { + let mut vregs = VirtRegs::new(); + let v1 = Value::new(1); + let v2 = Value::new(2); + let v3 = Value::new(3); + let v4 = Value::new(4); + + vregs.union(v2, v4); // Rank 0-0 + vregs.union(v3, v2); // Rank 0-1 + vregs.union(v2, v1); // Rank 1-0 + vregs.finish_union_find(None); + assert_eq!(vregs.congruence_class(&v1), &[v1, v3, v4, v2]); + assert_eq!(vregs.congruence_class(&v2), &[v1, v3, v4, v2]); + assert_eq!(vregs.congruence_class(&v3), &[v1, v3, v4, v2]); + assert_eq!(vregs.congruence_class(&v4), &[v1, v3, v4, v2]); + assert_eq!(vregs.all_virtregs().count(), 1); + } +}