1 use crate::{RealReg, RegUsageMapper, VirtualReg};
2 use smallvec::SmallVec;
3 use 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)]
14 pub 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 
26 impl VrangeRegUsageMapper {
27     /// Allocate a reg-usage mapper with the given predicted vreg capacity.
new(vreg_capacity: usize) -> VrangeRegUsageMapper28     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?
is_overlay_large_enough_to_sort(&self) -> bool36     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.
set_overlay(&mut self, vreg: VirtualReg, rreg: Option<RealReg>)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.
finish_overlay(&mut self)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.
merge_overlay(&mut self)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.
set_direct(&mut self, vreg: VirtualReg, rreg: Option<RealReg>)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 
set_direct_internal(&mut self, vreg: VirtualReg, rreg: RealReg)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.
lookup_direct(&self, vreg: VirtualReg) -> Option<RealReg>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.
lookup_overlay(&self, vreg: VirtualReg) -> Option<RealReg>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.
is_empty(&self) -> bool139     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 
145 impl RegUsageMapper for VrangeRegUsageMapper {
146     /// Return the `RealReg` if mapped, or `None`, for `vreg` occuring as a use
147     /// on the current instruction.
get_use(&self, vreg: VirtualReg) -> Option<RealReg>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.
get_def(&self, vreg: VirtualReg) -> Option<RealReg>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.
get_mod(&self, vreg: VirtualReg) -> Option<RealReg>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)]
173 mod test {
174     use super::*;
175     use crate::{Reg, RegClass, VirtualReg};
176 
vreg(idx: u32) -> VirtualReg177     fn vreg(idx: u32) -> VirtualReg {
178         Reg::new_virtual(RegClass::I64, idx).to_virtual_reg()
179     }
rreg(idx: u8) -> RealReg180     fn rreg(idx: u8) -> RealReg {
181         Reg::new_real(RegClass::I64, /* enc = */ 0, /* index = */ idx).to_real_reg()
182     }
183 
184     #[test]
test_reg_use_mapper()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)]
302 pub 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 
310 impl MentionRegUsageMapper {
new() -> Self311     pub(crate) fn new() -> Self {
312         Self {
313             uses: SmallVec::new(),
314             defs: SmallVec::new(),
315         }
316     }
clear(&mut self)317     pub(crate) fn clear(&mut self) {
318         self.uses.clear();
319         self.defs.clear();
320     }
lookup_use(&self, vreg: VirtualReg) -> Option<RealReg>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     }
lookup_def(&self, vreg: VirtualReg) -> Option<RealReg>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     }
set_use(&mut self, vreg: VirtualReg, rreg: RealReg)327     pub(crate) fn set_use(&mut self, vreg: VirtualReg, rreg: RealReg) {
328         self.uses.push((vreg, rreg));
329     }
set_def(&mut self, vreg: VirtualReg, rreg: RealReg)330     pub(crate) fn set_def(&mut self, vreg: VirtualReg, rreg: RealReg) {
331         self.defs.push((vreg, rreg));
332     }
333 }
334 
335 impl RegUsageMapper for MentionRegUsageMapper {
get_use(&self, vreg: VirtualReg) -> Option<RealReg>336     fn get_use(&self, vreg: VirtualReg) -> Option<RealReg> {
337         return self.lookup_use(vreg);
338     }
get_def(&self, vreg: VirtualReg) -> Option<RealReg>339     fn get_def(&self, vreg: VirtualReg) -> Option<RealReg> {
340         return self.lookup_def(vreg);
341     }
get_mod(&self, vreg: VirtualReg) -> Option<RealReg>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 }
348