regalloc/
reg_maps.rs

1use crate::{RealReg, RegUsageMapper, VirtualReg};
2use smallvec::SmallVec;
3use std::mem;
4
5/// This data structure holds the mappings needed to map an instruction's uses, mods and defs from
6/// virtual to real registers.
7///
8/// It remembers the sets of mappings (of a virtual register to a real register) over time, based
9/// on precise virtual ranges and their allocations.
10///
11/// This is the right implementation to use when a register allocation algorithm keeps track of
12/// precise virtual ranges, and maintains them over time.
13#[derive(Debug)]
14pub struct VrangeRegUsageMapper {
15    /// Dense vector-map indexed by virtual register number. This is consulted
16    /// directly for use-queries and augmented with the overlay for def-queries.
17    slots: Vec<RealReg>,
18
19    /// Overlay for def-queries. This is a set of updates that occurs "during"
20    /// the instruction in question, and will be applied to the slots array
21    /// once we are done processing this instruction (in preparation for
22    /// the next one).
23    overlay: SmallVec<[(VirtualReg, RealReg); 16]>,
24}
25
26impl VrangeRegUsageMapper {
27    /// Allocate a reg-usage mapper with the given predicted vreg capacity.
28    pub(crate) fn new(vreg_capacity: usize) -> VrangeRegUsageMapper {
29        VrangeRegUsageMapper {
30            slots: Vec::with_capacity(vreg_capacity),
31            overlay: SmallVec::new(),
32        }
33    }
34
35    /// Is the overlay past the sorted-size threshold?
36    fn is_overlay_large_enough_to_sort(&self) -> bool {
37        // Use the SmallVec spill-to-heap threshold as a threshold for "large
38        // enough to sort"; this has the effect of amortizing the cost of
39        // sorting along with the cost of copying out to heap memory, and also
40        // ensures that when we access heap (more likely to miss in cache), we
41        // do it with O(log N) accesses instead of O(N).
42        self.overlay.spilled()
43    }
44
45    /// Update the overlay.
46    pub(crate) fn set_overlay(&mut self, vreg: VirtualReg, rreg: Option<RealReg>) {
47        let rreg = rreg.unwrap_or(RealReg::invalid());
48        self.overlay.push((vreg, rreg));
49    }
50
51    /// Finish updates to the overlay, sorting if necessary.
52    pub(crate) fn finish_overlay(&mut self) {
53        if self.overlay.len() == 0 || !self.is_overlay_large_enough_to_sort() {
54            return;
55        }
56
57        // Sort stably, so that later updates continue to come after earlier
58        // ones.
59        self.overlay.sort_by_key(|pair| pair.0);
60        // Remove duplicates by collapsing runs of same-vreg pairs down to
61        // the last one.
62        let mut last_vreg = self.overlay[0].0;
63        let mut out = 0;
64        for i in 1..self.overlay.len() {
65            let this_vreg = self.overlay[i].0;
66            if this_vreg != last_vreg {
67                out += 1;
68            }
69            if i != out {
70                self.overlay[out] = self.overlay[i];
71            }
72            last_vreg = this_vreg;
73        }
74        let new_len = out + 1;
75        self.overlay.truncate(new_len);
76    }
77
78    /// Merge the overlay into the main map.
79    pub(crate) fn merge_overlay(&mut self) {
80        // Take the SmallVec and swap with empty to allow `&mut self` method
81        // call below.
82        let mappings = mem::replace(&mut self.overlay, SmallVec::new());
83        for (vreg, rreg) in mappings.into_iter() {
84            self.set_direct_internal(vreg, rreg);
85        }
86    }
87
88    /// Make a direct update to the mapping. Only usable when the overlay
89    /// is empty.
90    pub(crate) fn set_direct(&mut self, vreg: VirtualReg, rreg: Option<RealReg>) {
91        debug_assert!(self.overlay.is_empty());
92        let rreg = rreg.unwrap_or(RealReg::invalid());
93        self.set_direct_internal(vreg, rreg);
94    }
95
96    fn set_direct_internal(&mut self, vreg: VirtualReg, rreg: RealReg) {
97        let idx = vreg.get_index();
98        if idx >= self.slots.len() {
99            self.slots.resize(idx + 1, RealReg::invalid());
100        }
101        self.slots[idx] = rreg;
102    }
103
104    /// Perform a lookup directly in the main map. Returns `None` for
105    /// not-present.
106    fn lookup_direct(&self, vreg: VirtualReg) -> Option<RealReg> {
107        let idx = vreg.get_index();
108        if idx >= self.slots.len() {
109            None
110        } else {
111            Some(self.slots[idx])
112        }
113    }
114
115    /// Perform a lookup in the overlay. Returns `None` for not-present. No
116    /// fallback to main map (that happens in callers). Returns `Some` even
117    /// if mapped to `RealReg::invalid()`, because this is a tombstone
118    /// (represents deletion) in the overlay.
119    fn lookup_overlay(&self, vreg: VirtualReg) -> Option<RealReg> {
120        if self.is_overlay_large_enough_to_sort() {
121            // Do a binary search; we are guaranteed to have at most one
122            // matching because duplicates were collapsed after sorting.
123            if let Ok(idx) = self.overlay.binary_search_by_key(&vreg, |pair| pair.0) {
124                return Some(self.overlay[idx].1);
125            }
126        } else {
127            // Search in reverse order to find later updates first.
128            for &(this_vreg, this_rreg) in self.overlay.iter().rev() {
129                if this_vreg == vreg {
130                    return Some(this_rreg);
131                }
132            }
133        }
134        None
135    }
136
137    /// Sanity check: check that all slots are empty. Typically for use at the
138    /// end of processing as a debug-assert.
139    pub(crate) fn is_empty(&self) -> bool {
140        self.overlay.iter().all(|pair| pair.1.is_invalid())
141            && self.slots.iter().all(|rreg| rreg.is_invalid())
142    }
143}
144
145impl RegUsageMapper for VrangeRegUsageMapper {
146    /// Return the `RealReg` if mapped, or `None`, for `vreg` occuring as a use
147    /// on the current instruction.
148    fn get_use(&self, vreg: VirtualReg) -> Option<RealReg> {
149        self.lookup_direct(vreg)
150            // Convert Some(RealReg::invalid()) to None.
151            .and_then(|reg| reg.maybe_valid())
152    }
153
154    /// Return the `RealReg` if mapped, or `None`, for `vreg` occuring as a def
155    /// on the current instruction.
156    fn get_def(&self, vreg: VirtualReg) -> Option<RealReg> {
157        self.lookup_overlay(vreg)
158            .or_else(|| self.lookup_direct(vreg))
159            // Convert Some(RealReg::invalid()) to None.
160            .and_then(|reg| reg.maybe_valid())
161    }
162
163    /// Return the `RealReg` if mapped, or `None`, for a `vreg` occuring as a
164    /// mod on the current instruction.
165    fn get_mod(&self, vreg: VirtualReg) -> Option<RealReg> {
166        let result = self.get_use(vreg);
167        debug_assert_eq!(result, self.get_def(vreg));
168        result
169    }
170}
171
172#[cfg(test)]
173mod test {
174    use super::*;
175    use crate::{Reg, RegClass, VirtualReg};
176
177    fn vreg(idx: u32) -> VirtualReg {
178        Reg::new_virtual(RegClass::I64, idx).to_virtual_reg()
179    }
180    fn rreg(idx: u8) -> RealReg {
181        Reg::new_real(RegClass::I64, /* enc = */ 0, /* index = */ idx).to_real_reg()
182    }
183
184    #[test]
185    fn test_reg_use_mapper() {
186        let mut mapper = VrangeRegUsageMapper::new(/* estimated vregs = */ 16);
187        assert_eq!(None, mapper.get_use(vreg(0)));
188        assert_eq!(None, mapper.get_def(vreg(0)));
189        assert_eq!(None, mapper.get_mod(vreg(0)));
190
191        mapper.set_direct(vreg(0), Some(rreg(1)));
192        mapper.set_direct(vreg(1), Some(rreg(2)));
193
194        assert_eq!(Some(rreg(1)), mapper.get_use(vreg(0)));
195        assert_eq!(Some(rreg(1)), mapper.get_def(vreg(0)));
196        assert_eq!(Some(rreg(1)), mapper.get_mod(vreg(0)));
197        assert_eq!(Some(rreg(2)), mapper.get_use(vreg(1)));
198        assert_eq!(Some(rreg(2)), mapper.get_def(vreg(1)));
199        assert_eq!(Some(rreg(2)), mapper.get_mod(vreg(1)));
200
201        mapper.set_overlay(vreg(0), Some(rreg(3)));
202        mapper.set_overlay(vreg(2), Some(rreg(4)));
203        mapper.finish_overlay();
204
205        assert_eq!(Some(rreg(1)), mapper.get_use(vreg(0)));
206        assert_eq!(Some(rreg(3)), mapper.get_def(vreg(0)));
207        // vreg 0 not valid for mod (use and def differ).
208        assert_eq!(Some(rreg(2)), mapper.get_use(vreg(1)));
209        assert_eq!(Some(rreg(2)), mapper.get_def(vreg(1)));
210        assert_eq!(Some(rreg(2)), mapper.get_mod(vreg(1)));
211        assert_eq!(None, mapper.get_use(vreg(2)));
212        assert_eq!(Some(rreg(4)), mapper.get_def(vreg(2)));
213        // vreg 2 not valid for mod (use and def differ).
214
215        mapper.merge_overlay();
216        assert_eq!(Some(rreg(3)), mapper.get_use(vreg(0)));
217        assert_eq!(Some(rreg(2)), mapper.get_use(vreg(1)));
218        assert_eq!(Some(rreg(4)), mapper.get_use(vreg(2)));
219        assert_eq!(None, mapper.get_use(vreg(3)));
220
221        // Check tombstoning behavior.
222        mapper.set_overlay(vreg(0), None);
223        mapper.finish_overlay();
224        assert_eq!(Some(rreg(3)), mapper.get_use(vreg(0)));
225        assert_eq!(None, mapper.get_def(vreg(0)));
226        mapper.merge_overlay();
227
228        // Check large (sorted) overlay mode.
229        for i in (2..50).rev() {
230            mapper.set_overlay(vreg(i), Some(rreg((i + 100) as u8)));
231        }
232        mapper.finish_overlay();
233        assert_eq!(None, mapper.get_use(vreg(0)));
234        assert_eq!(Some(rreg(2)), mapper.get_use(vreg(1)));
235        assert_eq!(Some(rreg(4)), mapper.get_use(vreg(2)));
236        for i in 2..50 {
237            assert_eq!(Some(rreg((i + 100) as u8)), mapper.get_def(vreg(i)));
238        }
239        mapper.merge_overlay();
240
241        for i in (0..100).rev() {
242            mapper.set_overlay(vreg(i), None);
243        }
244        mapper.finish_overlay();
245        for i in 0..100 {
246            assert_eq!(None, mapper.get_def(vreg(i)));
247        }
248        assert_eq!(false, mapper.is_empty());
249        mapper.merge_overlay();
250        assert_eq!(true, mapper.is_empty());
251
252        // Check multiple-update behavior in small mode.
253        mapper.set_overlay(vreg(1), Some(rreg(1)));
254        mapper.set_overlay(vreg(1), Some(rreg(2)));
255        mapper.finish_overlay();
256        assert_eq!(Some(rreg(2)), mapper.get_def(vreg(1)));
257        mapper.merge_overlay();
258        assert_eq!(Some(rreg(2)), mapper.get_use(vreg(1)));
259
260        mapper.set_overlay(vreg(1), Some(rreg(2)));
261        mapper.set_overlay(vreg(1), None);
262        mapper.finish_overlay();
263        assert_eq!(None, mapper.get_def(vreg(1)));
264        mapper.merge_overlay();
265        assert_eq!(None, mapper.get_use(vreg(1)));
266
267        // Check multiple-update behavior in sorted mode.
268        for i in 0..100 {
269            mapper.set_overlay(vreg(2), Some(rreg(i)));
270        }
271        for i in 0..100 {
272            mapper.set_overlay(vreg(2), Some(rreg(2 * i)));
273        }
274        mapper.finish_overlay();
275        assert_eq!(Some(rreg(198)), mapper.get_def(vreg(2)));
276        mapper.merge_overlay();
277        assert_eq!(Some(rreg(198)), mapper.get_use(vreg(2)));
278
279        for i in 0..100 {
280            mapper.set_overlay(vreg(2), Some(rreg(i)));
281        }
282        for _ in 0..100 {
283            mapper.set_overlay(vreg(2), None);
284        }
285        mapper.finish_overlay();
286        assert_eq!(None, mapper.get_def(vreg(50)));
287        mapper.merge_overlay();
288        assert_eq!(None, mapper.get_use(vreg(50)));
289    }
290}
291
292/// This implementation of RegUsageMapper relies on explicit mentions of vregs in instructions. The
293/// caller must keep them, and for each instruction:
294///
295/// - clear the previous mappings, using `clear()`,
296/// - feed the mappings from vregs to rregs for uses and defs, with `set_use`/`set_def`,
297/// - then call the `Function::map_regs` function with this structure.
298///
299/// This avoids a lot of resizes, and makes it possible for algorithms that don't have precise live
300/// ranges to fill in vreg -> rreg mappings.
301#[derive(Debug)]
302pub struct MentionRegUsageMapper {
303    /// Sparse vector-map indexed by virtual register number. This is consulted for use-queries.
304    uses: SmallVec<[(VirtualReg, RealReg); 8]>,
305
306    /// Sparse vector-map indexed by virtual register number. This is consulted for def-queries.
307    defs: SmallVec<[(VirtualReg, RealReg); 8]>,
308}
309
310impl MentionRegUsageMapper {
311    pub(crate) fn new() -> Self {
312        Self {
313            uses: SmallVec::new(),
314            defs: SmallVec::new(),
315        }
316    }
317    pub(crate) fn clear(&mut self) {
318        self.uses.clear();
319        self.defs.clear();
320    }
321    pub(crate) fn lookup_use(&self, vreg: VirtualReg) -> Option<RealReg> {
322        self.uses.iter().find(|&pair| pair.0 == vreg).map(|x| x.1)
323    }
324    pub(crate) fn lookup_def(&self, vreg: VirtualReg) -> Option<RealReg> {
325        self.defs.iter().find(|&pair| pair.0 == vreg).map(|x| x.1)
326    }
327    pub(crate) fn set_use(&mut self, vreg: VirtualReg, rreg: RealReg) {
328        self.uses.push((vreg, rreg));
329    }
330    pub(crate) fn set_def(&mut self, vreg: VirtualReg, rreg: RealReg) {
331        self.defs.push((vreg, rreg));
332    }
333}
334
335impl RegUsageMapper for MentionRegUsageMapper {
336    fn get_use(&self, vreg: VirtualReg) -> Option<RealReg> {
337        return self.lookup_use(vreg);
338    }
339    fn get_def(&self, vreg: VirtualReg) -> Option<RealReg> {
340        return self.lookup_def(vreg);
341    }
342    fn get_mod(&self, vreg: VirtualReg) -> Option<RealReg> {
343        let result = self.lookup_use(vreg);
344        debug_assert_eq!(result, self.lookup_def(vreg));
345        return result;
346    }
347}
regalloc/reg_maps.rs

regalloc/
reg_maps.rs
