Skip to main content

rucc_regalloc/
backtrack.rs

1//! Which register each value lives in, decided in the order the values are hardest to place, and
2//! undone when a value that would cost more to lose finds its register taken.
3//!
4//! Design: `spec/optimizer/39-register-allocation.md` section 39.7, and tamnd/rucc#1177.
5//!
6//! [`crate::assign`] is the `-O0` answer. It walks the line once, and when it runs out of
7//! registers the value that goes to the stack is the one whose range ends last. That is a guess
8//! about cost made from a fact about length, and it is wrong exactly where it matters: a value read
9//! in every turn of a loop and wanted again after the loop ends last, so it is the one that goes,
10//! and every turn of the loop pays a load for it.
11//!
12//! This answers the same question with the cost in it. Every value gets a weight, which is how
13//! often it is read or written, each time counted by how often the block it happens in runs, over
14//! how much of the function it is live across. A value with a high weight is one that would cost
15//! a lot to keep in memory for little register in return, and that is the one to keep.
16//!
17//! # The order values are placed in
18//!
19//! Longest first. A long value meets more of the others than a short one does, so it has the
20//! fewest registers to choose from, and giving it first choice is what leaves the short ones
21//! something to fit into. That is the order LLVM's greedy allocator takes them in, for the same
22//! reason.
23//!
24//! # Backtracking
25//!
26//! Going longest first means a long cold value takes a register before a short hot one has been
27//! looked at. When the hot one comes and finds nothing free, it asks what it would cost to take a
28//! register back. For each register the answer is the values in it that are in the way, and the
29//! register can be taken when every one of them weighs less than the value asking. Of the
30//! registers that can, the one taken is the one whose heaviest value in the way is lightest. The
31//! values that lose it go back in the queue and look for another register, and a value that has
32//! lost one [`ROUNDS`] times goes to the stack instead of looking again.
33//!
34//! That rule is also why it stops. A value only ever takes a register from values lighter than
35//! itself, and each value is put back a bounded number of times.
36//!
37//! # What it keeps from the linear scan
38//!
39//! Everything that says what a register may hold. The registers an instruction insists on, the
40//! values an instruction can only read from memory, the two address instructions and the hints
41//! are all read the way [`crate::assign`] reads them, from the same functions, so the two
42//! allocators cannot disagree about what the machine allows. They only disagree about who gets
43//! the register, and [`crate::check`] asks the same questions of either answer.
44//!
45//! A two address instruction is coalesced from both ends here. The linear scan only ever meets the
46//! answer after its source, since the source is written first. Here either can be placed first, so
47//! a source looks at where the answer that reuses it went as well as the other way round. That
48//! goes for the second source of an instruction that reads its sources either way round too: it
49//! can follow an answer placed before it by having the sources swapped, when the first source is
50//! wanted after the instruction and so can never be where the answer goes. A loaded value added to
51//! a base the loop reads again is that case.
52//!
53//! # When it gives up
54//!
55//! Every question of whether two values are both wanted is counted, and a function that asks more
56//! than [`BUDGET`] of them is handed to the linear scan instead. The answer is worse and it comes
57//! out in time, which is what the section asks of a pathological function. It is a count of work
58//! rather than of values because a function with many short values that never meet is cheap
59//! however many there are.
60//!
61//! # What the spill phase adds
62//!
63//! It runs twice when [`crate::pressure`] finds a point with more values live than registers. Once
64//! as above, where a value goes to the stack only when the queue reaches it and it can take no
65//! register back, and once with the values [`crate::spill`] picked sent to the stack before the
66//! queue starts. The first is better where the pressure is brief and eviction settles it in a few
67//! moves. The second is better where it is long, since the values that go are picked by weight
68//! across every point that is over rather than by which one the queue met last. Neither wins
69//! everywhere, so both are costed.
70//!
71//! It also gives up when it would lose. The linear scan runs as well, which is cheap next to this,
72//! and the answer kept is the one with the lower [`cost`]: the loads and stores of the values on
73//! the stack and the copies between the ends of each tie left apart, each counted by how often its
74//! block runs. Placing the long values first is right where registers are fought over in a loop,
75//! and it can be worse in a long straight run of arithmetic, where the order the linear scan walks
76//! in is also the order that lets each answer follow its source.
77
78use std::cmp::Reverse;
79use std::collections::{BinaryHeap, HashMap, HashSet};
80
81use rucc_mir::{Func, Inst, Reg};
82use rucc_target::{PhysReg, RegClass};
83
84use crate::assign::{self, Assignment, Blocks, Env, Place, Reuse, Want};
85use crate::live::{Area, Live, Range};
86use crate::order::Order;
87use crate::pressure::Pressure;
88use crate::spill;
89
90/// How many questions of whether two values are both wanted a function may ask before it is handed
91/// to the linear scan.
92pub const BUDGET: u64 = 50_000_000;
93
94/// How many times a value may lose its register and look for another before it goes to the stack.
95pub const ROUNDS: u32 = 4;
96
97/// One value, as the queue sees it.
98#[derive(Debug, Clone, Copy)]
99struct Value<'a> {
100    reg: Reg,
101    class: RegClass,
102    area: Area<'a>,
103    range: Range,
104    weight: u128,
105    size: u32,
106}
107
108/// Where every value goes, in the order that places the hard ones first and takes a register back
109/// when a heavier value wants it.
110///
111/// A function that asks more than [`BUDGET`] questions gets the linear scan's answer instead, and
112/// so does one where the linear scan's answer has the lower [`cost`].
113#[must_use]
114pub fn assign(func: &Func, order: &Order, live: &Live, env: &Env) -> Assignment {
115    let linear = assign::assign(func, order, live, env);
116    let mut best = cost(func, order, &linear);
117    let mut kept = linear;
118    let pressure = Pressure::of(func, order, live, env);
119    let spilled = spill::choose(func, live, &pressure);
120    // With nothing sent ahead the second try would be the first one again.
121    let tries: &[&[Reg]] = if spilled.is_empty() { &[&[]] } else { &[&[], &spilled] };
122    for &early in tries {
123        let Some(ours) = placed(func, order, live, env, BUDGET, early) else { break };
124        let spent = cost(func, order, &ours);
125        if spent <= best {
126            best = spent;
127            kept = ours;
128        }
129    }
130    kept
131}
132
133/// What an assignment is expected to cost a function, in instructions each counted by how often
134/// its block runs.
135///
136/// A value on the stack costs a load or a store every time it is read or written. A two address
137/// instruction whose answer is not where its source is costs the copy between them, and so does a
138/// block parameter that is not where the value passed to it is. Moves the machine needs whatever
139/// the assignment, like those into the registers a call insists on, are left out, since they are
140/// the same for any answer.
141///
142/// # Panics
143///
144/// Panics on a function with more values than a register number can name, as
145/// [`Reg::virtual_reg`] does.
146#[must_use]
147pub fn cost(func: &Func, order: &Order, assignment: &Assignment) -> u128 {
148    let costs = costs(func);
149    let mut total = 0;
150    for (reg, place) in assignment.placed() {
151        if !matches!(place, Place::Reg(_)) {
152            total += costs[index(reg)];
153        }
154    }
155    let mut weights = HashMap::new();
156    for block in func.blocks() {
157        let weight = u128::from(func[block].weight.raw().max(1));
158        for inst in func.insts(block) {
159            weights.insert(inst, weight);
160        }
161        for call in &func[block].succs {
162            for (&arg, param) in call.args.iter().zip(&func[call.block].params) {
163                if assignment.place(arg) != assignment.place(param.reg) {
164                    total += weight;
165                }
166            }
167        }
168    }
169    let commuted: HashSet<Inst> = assignment.commuted().iter().copied().collect();
170    for (number, reuse) in assign::reuses(func, order).iter().enumerate() {
171        let Some(reuse) = reuse else { continue };
172        let answer = Reg::virtual_reg(u32::try_from(number).expect("a register number"));
173        let tied = if commuted.contains(&reuse.inst) { reuse.second } else { Some(reuse.source) };
174        let Some(at) = assignment.place(answer) else { continue };
175        if tied.and_then(|tied| assignment.place(tied)) != Some(at) {
176            total += weights.get(&reuse.inst).copied().unwrap_or(1);
177        }
178    }
179    total
180}
181
182/// The same, with the budget said, or `None` for a function that went over it.
183///
184/// # Panics
185///
186/// Panics on a value in a class the environment hands out no registers from, as the linear scan
187/// does.
188#[must_use]
189pub fn within(
190    func: &Func,
191    order: &Order,
192    live: &Live,
193    env: &Env,
194    budget: u64,
195) -> Option<Assignment> {
196    placed(func, order, live, env, budget, &[])
197}
198
199fn placed(
200    func: &Func,
201    order: &Order,
202    live: &Live,
203    env: &Env,
204    budget: u64,
205    early: &[Reg],
206) -> Option<Assignment> {
207    let blocked = assign::blocked(func, order);
208    let forced = assign::forced(func);
209    let reuses = assign::reuses(func, order);
210    let hints = assign::hints(func);
211    let passed = assign::passed(func);
212    let received = received(func);
213    let reused = reused(&reuses);
214    let seconds = seconds(&reuses);
215    let costs = costs(func);
216
217    let count = func.vregs();
218    let mut values: Vec<Option<Value<'_>>> = vec![None; count];
219    let mut queue = BinaryHeap::new();
220    for (number, reuse) in reuses.iter().enumerate() {
221        let reg = Reg::virtual_reg(u32::try_from(number).expect("a register number"));
222        let (Some(mut area), Some(class)) = (live.area(reg), func.class_of(reg)) else {
223            continue;
224        };
225        if let Some(reuse) = reuse {
226            area = area.with(reuse.at);
227        }
228        let size = size(area);
229        let weight = costs[number] * 1024 / u128::from(size + 8);
230        values[number] = Some(Value { reg, class, area, range: area.hull(), weight, size });
231        queue.push((size, Reverse(number)));
232    }
233
234    let mut state = State {
235        live,
236        blocked: &blocked,
237        reuses: &reuses,
238        values: &values,
239        held: Vec::new(),
240        at: vec![None; count],
241        commuted: vec![None; count],
242        work: 0,
243        budget,
244    };
245    let mut assignment = Assignment::empty(count);
246    let mut lost = vec![0u32; count];
247    while let Some((_, Reverse(number))) = queue.pop() {
248        let Some(value) = values[number] else { continue };
249        if forced.contains(&value.reg) || early.contains(&value.reg) {
250            assignment.spill(value.reg, value.class);
251            continue;
252        }
253        assert!(
254            !env.order(value.class).is_empty(),
255            "a value in class {}, which the target hands out no registers from",
256            value.class.number()
257        );
258        let chosen = state
259            .coalesced(value, &reused[number])
260            .or_else(|| state.hinted(value, &hints[number], Want::Clear))
261            .or_else(|| {
262                let partners = passed[number].iter().chain(&received[number]);
263                let partners: Vec<PhysReg> =
264                    partners.filter_map(|&other| state.reg_of(other)).collect();
265                state.hinted(value, &partners, Want::Clear)
266            })
267            .or_else(|| state.swapped(value, &seconds[number]))
268            .or_else(|| {
269                let mut ties = reused[number].clone();
270                let source = reuses[number].and_then(|reuse| reuse.source.number());
271                ties.extend(source.and_then(|source| usize::try_from(source).ok()));
272                ties.retain(|&tie| state.at[tie].is_none() && values[tie].is_some());
273                let tied = ties.iter().flat_map(|&tie| hints[tie].iter());
274                let wanted: Vec<PhysReg> = tied.chain(env.order(value.class)).copied().collect();
275                state.together(value, &ties, &wanted)
276            })
277            .or_else(|| state.hinted(value, env.order(value.class), Want::Clear))
278            .or_else(|| state.hinted(value, env.order(value.class), Want::Allowed));
279        if state.work > state.budget {
280            return None;
281        }
282        if let Some(at) = chosen {
283            state.take(value, at);
284            continue;
285        }
286        match state.cheapest(value, env.order(value.class)) {
287            Some(at) => {
288                for other in state.evict(value, at) {
289                    lost[other] += 1;
290                    let Some(evicted) = values[other] else { continue };
291                    if lost[other] > ROUNDS {
292                        assignment.spill(evicted.reg, evicted.class);
293                    } else {
294                        queue.push((evicted.size, Reverse(other)));
295                    }
296                }
297                state.take(value, at);
298            }
299            None => assignment.spill(value.reg, value.class),
300        }
301        if state.work > state.budget {
302            return None;
303        }
304    }
305    state.settle(&reused, &passed, &received);
306    if state.work > state.budget {
307        return None;
308    }
309    for (number, at) in state.at.iter().enumerate() {
310        let Some(at) = *at else { continue };
311        let reg = Reg::virtual_reg(u32::try_from(number).expect("a register number"));
312        assignment.put(reg, Place::Reg(at));
313    }
314    for inst in state.commuted.iter().flatten() {
315        assignment.commute(*inst);
316    }
317    Some(assignment)
318}
319
320/// A value in a register, by number, with the stretch of the line it covers.
321type Held = (usize, Range);
322
323/// What the allocation knows while it runs.
324struct State<'a, 'v> {
325    live: &'a Live,
326    blocked: &'a Blocks,
327    reuses: &'a [Option<Reuse>],
328    values: &'v [Option<Value<'a>>],
329    /// Which values are in each register of each class, by number, each with the stretch of the
330    /// line it covers. Most values in a register are nowhere near the one being placed, and having
331    /// the stretch here tells so without reading the value itself from wherever it is in `values`.
332    held: Vec<((RegClass, PhysReg), Vec<Held>)>,
333    /// Which register each value is in now, if one.
334    at: Vec<Option<PhysReg>>,
335    /// The instruction whose sources were swapped for the answer written by each value, if its
336    /// register is the second source's.
337    commuted: Vec<Option<Inst>>,
338    work: u64,
339    budget: u64,
340}
341
342impl<'a> State<'a, '_> {
343    fn reg_of(&self, reg: Reg) -> Option<PhysReg> {
344        self.at.get(usize::try_from(reg.number()?).ok()?).copied().flatten()
345    }
346
347    fn slot(&mut self, class: RegClass, at: PhysReg) -> &mut Vec<Held> {
348        let found = self.held.iter().position(|(key, _)| *key == (class, at));
349        let index = found.unwrap_or_else(|| {
350            self.held.push(((class, at), Vec::new()));
351            self.held.len() - 1
352        });
353        &mut self.held[index].1
354    }
355
356    /// The values in `at` that are wanted while `value` is, leaving out the one a two address
357    /// instruction lets it share the register with.
358    fn clashes(&mut self, value: Value<'_>, at: PhysReg) -> Vec<usize> {
359        let Some(index) = self.held.iter().position(|(key, _)| *key == (value.class, at)) else {
360            return Vec::new();
361        };
362        let mut clashes = Vec::new();
363        for &(other, range) in &self.held[index].1 {
364            self.work += 1;
365            if !range.overlaps(value.range) {
366                continue;
367            }
368            let Some(held) = self.values[other] else { continue };
369            if !held.area.overlaps(value.area) {
370                continue;
371            }
372            if !self.shares(value.reg, held.reg) {
373                clashes.push(other);
374            }
375        }
376        clashes
377    }
378
379    /// Whether two values that are both wanted at one instruction may still be in one register,
380    /// which is when that instruction reads one for the last time and writes the other over it.
381    fn shares(&self, one: Reg, two: Reg) -> bool {
382        let reads = |answer: Reg, source: Reg| {
383            let Some(number) = answer.number().and_then(|n| usize::try_from(n).ok()) else {
384                return false;
385            };
386            let Some(reuse) = self.reuses[number] else { return false };
387            match self.commuted[number] {
388                Some(_) => reuse.second == Some(source),
389                None => reuse.source == source,
390            }
391        };
392        (reads(one, two) || reads(two, one)) && assign::apart(self.live, one, two)
393    }
394
395    fn free(&mut self, value: Value<'_>, at: PhysReg, want: Want) -> bool {
396        !self.blocked.insists(value.reg, value.class, value.area, value.range, at, want)
397            && self.clashes(value, at).is_empty()
398    }
399
400    /// The first register of `wanted` that is free for `value`.
401    fn hinted(&mut self, value: Value<'_>, wanted: &[PhysReg], want: Want) -> Option<PhysReg> {
402        wanted.iter().copied().find(|&at| self.work <= self.budget && self.free(value, at, want))
403    }
404
405    /// The first register of `order` free for `value` and for each value in `ties`, which are the
406    /// values a two address instruction ties it to that have no register yet. Taking one of these
407    /// leaves the tied value a register to follow it into when its turn comes. The registers the
408    /// tied values are hinted to come first in `order`, so that a parameter's answer waits for it in
409    /// the register the parameter arrives in.
410    fn together(&mut self, value: Value<'_>, ties: &[usize], order: &[PhysReg]) -> Option<PhysReg> {
411        if ties.is_empty() {
412            return None;
413        }
414        for &at in order {
415            if self.work > self.budget {
416                return None;
417            }
418            if !self.free(value, at, Want::Clear) {
419                continue;
420            }
421            let values = self.values;
422            let mut tied = ties.iter().filter_map(|&tie| values[tie]);
423            if tied.all(|other| self.free(other, at, Want::Allowed)) {
424                return Some(at);
425            }
426        }
427        None
428    }
429
430    /// The register of a value a two address instruction ties this one to, if it can have it.
431    ///
432    /// Asked of the answer, it is the register of the source, or of the second source if the
433    /// instruction reads its sources either way round. Asked of a source, it is the register of an
434    /// answer that reuses it and was placed first.
435    fn coalesced(&mut self, value: Value<'_>, answers: &[usize]) -> Option<PhysReg> {
436        let number = index(value.reg);
437        if let Some(reuse) = self.reuses[number] {
438            if let Some(at) = self.reg_of(reuse.source) {
439                if self.free(value, at, Want::Allowed) {
440                    return Some(at);
441                }
442            }
443            if let Some(second) = reuse.second {
444                if let Some(at) = self.reg_of(second) {
445                    self.commuted[number] = Some(reuse.inst);
446                    if self.free(value, at, Want::Allowed) {
447                        return Some(at);
448                    }
449                    self.commuted[number] = None;
450                }
451            }
452        }
453        for &answer in answers {
454            let Some(at) = self.at[answer] else { continue };
455            if self.commuted[answer].is_none() && self.free(value, at, Want::Allowed) {
456                return Some(at);
457            }
458        }
459        None
460    }
461
462    /// The register of an answer placed first that could be written over this value with its
463    /// sources swapped, and cannot be written over its first source, because that one is still
464    /// wanted after it. Asked after the hints, because the register an instruction wants the value
465    /// in saves a move as well, and following the answer would take the value away from it. The
466    /// value is one no two address instruction writes, such as a load.
467    fn swapped(&mut self, value: Value<'_>, seconds: &[usize]) -> Option<PhysReg> {
468        // A value that is itself the answer of a two address instruction is tied to its own source
469        // already, and following the answer it is read by would break that tie to save the same
470        // copy somewhere else.
471        if self.reuses[index(value.reg)].is_some() {
472            return None;
473        }
474        for &answer in seconds {
475            let (Some(at), Some(reuse)) = (self.at[answer], self.reuses[answer]) else { continue };
476            let Some(written) = self.values[answer] else { continue };
477            // An answer whose first source ends where it starts can still go over that one, which
478            // saves the same copy without swapping anything, and taking its register here would
479            // stop it moving there when the function is settled.
480            if self.commuted[answer].is_some()
481                || self.reg_of(reuse.source) == Some(at)
482                || assign::apart(self.live, written.reg, reuse.source)
483            {
484                continue;
485            }
486            self.commuted[answer] = Some(reuse.inst);
487            if self.free(value, at, Want::Allowed) {
488                return Some(at);
489            }
490            self.commuted[answer] = None;
491        }
492        None
493    }
494
495    /// The register that is cheapest to take back for `value`, if any is cheaper than sending
496    /// `value` to the stack.
497    fn cheapest(&mut self, value: Value<'_>, order: &[PhysReg]) -> Option<PhysReg> {
498        let mut best: Option<(u128, u128, PhysReg)> = None;
499        for &at in order {
500            if self.blocked.insists(
501                value.reg,
502                value.class,
503                value.area,
504                value.range,
505                at,
506                Want::Allowed,
507            ) {
508                continue;
509            }
510            let clashes = self.clashes(value, at);
511            let weights = clashes.iter().filter_map(|&other| self.values[other]).map(|v| v.weight);
512            let (heaviest, total) = weights.fold((0, 0), |(most, sum), w| (most.max(w), sum + w));
513            if heaviest >= value.weight {
514                continue;
515            }
516            if best.is_none_or(|(most, sum, _)| (heaviest, total) < (most, sum)) {
517                best = Some((heaviest, total, at));
518            }
519        }
520        best.map(|(_, _, at)| at)
521    }
522
523    /// Takes `at` back from the values in it that are in the way of `value`, and says which.
524    fn evict(&mut self, value: Value<'_>, at: PhysReg) -> Vec<usize> {
525        let clashes = self.clashes(value, at);
526        self.slot(value.class, at).retain(|(other, _)| !clashes.contains(other));
527        for &other in &clashes {
528            self.at[other] = None;
529            self.commuted[other] = None;
530        }
531        clashes
532    }
533
534    /// Moves each value into the register the most values tied to it are in, of those that are free
535    /// for it and hold more of them than where it is now.
536    ///
537    /// A tie is a two address instruction or a block parameter, and each one met is a copy the
538    /// rewrite does not have to write. Placing values in priority order can leave the two ends of a
539    /// tie apart though the register one of them is in stayed free for the other, because the other
540    /// was placed first and had nothing yet to follow. Every move meets more ties than it breaks, so
541    /// this ends.
542    fn settle(&mut self, reused: &[Vec<usize>], passed: &[Vec<Reg>], received: &[Vec<Reg>]) {
543        let mut moved = true;
544        while moved && self.work <= self.budget {
545            moved = false;
546            for number in 0..self.at.len() {
547                let (Some(value), Some(now)) = (self.values[number], self.at[number]) else {
548                    continue;
549                };
550                let reuse = self.reuses[number];
551                let source = reuse.and_then(|reuse| self.reg_of(reuse.source));
552                let second =
553                    reuse.and_then(|reuse| reuse.second).and_then(|second| self.reg_of(second));
554                let mut wanted: Vec<PhysReg> = source.into_iter().chain(second).collect();
555                for &answer in &reused[number] {
556                    if self.commuted[answer].is_none() {
557                        wanted.extend(self.at[answer]);
558                    }
559                }
560                let partners = passed[number].iter().chain(&received[number]);
561                wanted.extend(partners.filter_map(|&other| self.reg_of(other)));
562                let met = |at: PhysReg| wanted.iter().filter(|&&reg| reg == at).count();
563                let here = met(now);
564                let mut better: Vec<(usize, PhysReg)> = Vec::new();
565                for &at in &wanted {
566                    let count = met(at);
567                    if count > here && !better.contains(&(count, at)) {
568                        better.push((count, at));
569                    }
570                }
571                if better.is_empty() {
572                    continue;
573                }
574                better.sort_by_key(|&(count, _)| Reverse(count));
575                self.slot(value.class, now).retain(|&(other, _)| other != number);
576                self.at[number] = None;
577                let was = self.commuted[number];
578                let mut to = now;
579                for (_, at) in better {
580                    self.commuted[number] = match reuse {
581                        Some(reuse) if second == Some(at) && source != Some(at) => Some(reuse.inst),
582                        _ => None,
583                    };
584                    if self.free(value, at, Want::Allowed) {
585                        to = at;
586                        moved = true;
587                        break;
588                    }
589                }
590                if to == now {
591                    self.commuted[number] = was;
592                }
593                self.take(value, to);
594            }
595        }
596    }
597
598    fn take(&mut self, value: Value<'_>, at: PhysReg) {
599        let number = index(value.reg);
600        self.at[number] = Some(at);
601        self.slot(value.class, at).push((number, value.range));
602    }
603}
604
605/// How much of the line a value is live over.
606pub(crate) fn size(area: Area<'_>) -> u32 {
607    area.pieces().map(|piece| piece.end - piece.start + 1).sum()
608}
609
610/// How often each value is read or written, each time counted by how often its block runs.
611///
612/// A block that says nothing about how often it runs counts as running once, so a function with no
613/// weights on it is counted by how many times each value is named.
614pub(crate) fn costs(func: &Func) -> Vec<u128> {
615    let mut costs = vec![0u128; func.vregs()];
616    let mut add = |reg: Reg, weight: u128| {
617        let number = reg.number().and_then(|number| usize::try_from(number).ok());
618        if let Some(cost) = number.and_then(|number| costs.get_mut(number)) {
619            *cost += weight;
620        }
621    };
622    for block in func.blocks() {
623        let weight = u128::from(func[block].weight.raw().max(1));
624        for param in &func[block].params {
625            add(param.reg, weight);
626        }
627        for inst in func.insts(block) {
628            for operand in &func[func[inst].operands] {
629                add(operand.reg, weight);
630            }
631        }
632        for call in &func[block].succs {
633            for &arg in &call.args {
634                add(arg, weight);
635            }
636        }
637    }
638    costs
639}
640
641/// For each parameter of a block, the values the edges into it pass it.
642fn received(func: &Func) -> Vec<Vec<Reg>> {
643    let mut received = vec![Vec::new(); func.vregs()];
644    for block in func.blocks() {
645        for call in &func[block].succs {
646            for (&arg, param) in call.args.iter().zip(&func[call.block].params) {
647                let number = param.reg.number().and_then(|number| usize::try_from(number).ok());
648                let Some(from) = number.and_then(|number| received.get_mut(number)) else {
649                    continue;
650                };
651                if !from.contains(&arg) {
652                    from.push(arg);
653                }
654            }
655        }
656    }
657    received
658}
659
660/// For each value, the answers of two address instructions that reuse it as their first source.
661fn reused(reuses: &[Option<Reuse>]) -> Vec<Vec<usize>> {
662    let mut reused = vec![Vec::new(); reuses.len()];
663    for (answer, reuse) in reuses.iter().enumerate() {
664        let Some(reuse) = reuse else { continue };
665        let number = reuse.source.number().and_then(|number| usize::try_from(number).ok());
666        if let Some(answers) = number.and_then(|number| reused.get_mut(number)) {
667            answers.push(answer);
668        }
669    }
670    reused
671}
672
673/// The answers that could be written over each value as the second source of an instruction that
674/// reads its sources either way round.
675fn seconds(reuses: &[Option<Reuse>]) -> Vec<Vec<usize>> {
676    let mut seconds = vec![Vec::new(); reuses.len()];
677    for (answer, reuse) in reuses.iter().enumerate() {
678        let Some(second) = reuse.and_then(|reuse| reuse.second) else { continue };
679        let number = second.number().and_then(|number| usize::try_from(number).ok());
680        if let Some(answers) = number.and_then(|number| seconds.get_mut(number)) {
681            answers.push(answer);
682        }
683    }
684    seconds
685}
686
687fn index(reg: Reg) -> usize {
688    usize::try_from(reg.number().expect("a virtual register")).expect("a register number")
689}
690
691#[cfg(test)]
692mod tests {
693    use rucc_base::Interner;
694    use rucc_mir::{BlockCall, Constraint, Flags, Opcode, Operand, Param};
695    use rucc_target::x86_64::{GPR, REGS, SYSV};
696
697    use super::*;
698    use crate::check;
699
700    fn env() -> Env {
701        let (order, scratch) = SYSV.int_order.split_at(SYSV.int_order.len() - 3);
702        Env::new().with(GPR, order, scratch)
703    }
704
705    fn narrow(count: usize) -> Env {
706        Env::new().with(GPR, &SYSV.int_order[..count], &SYSV.int_order[count..count + 1])
707    }
708
709    fn named(place: Option<Place>) -> String {
710        match place {
711            Some(Place::Reg(reg)) => REGS.name(GPR, reg).expect("a register").to_string(),
712            Some(Place::Slot(_)) => "slot".to_string(),
713            None => "nowhere".to_string(),
714        }
715    }
716
717    /// Where every value went, after asking the checker whether the machine could run it.
718    fn places(func: &mut Func, env: &Env) -> Vec<String> {
719        let order = Order::of(func);
720        let live = Live::of(func, &order);
721        let assignment = within(func, &order, &live, env, BUDGET).expect("inside the budget");
722        // The sources of an instruction that commutes are swapped before the check, as `run_with`
723        // swaps them.
724        for &inst in assignment.commuted() {
725            let list = func[inst].operands;
726            func[list].swap(1, 2);
727        }
728        let problems = check::check(func, &order, &live, &assignment);
729        assert!(problems.is_empty(), "{}", check::report(&problems));
730        (0..func.vregs())
731            .map(|number| {
732                let reg = Reg::virtual_reg(u32::try_from(number).expect("a register number"));
733                named(assignment.place(reg))
734            })
735            .collect()
736    }
737
738    fn linear(func: &Func, env: &Env) -> Vec<String> {
739        let order = Order::of(func);
740        let live = Live::of(func, &order);
741        let assignment = assign::assign(func, &order, &live, env);
742        (0..func.vregs())
743            .map(|number| {
744                let reg = Reg::virtual_reg(u32::try_from(number).expect("a register number"));
745                named(assignment.place(reg))
746            })
747            .collect()
748    }
749
750    #[test]
751    fn a_value_the_spill_phase_picked_goes_to_the_stack_and_the_rest_fit() {
752        let mut names = Interner::new();
753        let mut func = Func::new(names.intern("f"));
754        let opcode = Opcode::new(names.intern("x64.nop"));
755        let block = func.create_block();
756        let busy = func.new_vreg(GPR);
757        let once = func.new_vreg(GPR);
758        let other = func.new_vreg(GPR);
759        func.build(block, opcode).def(busy, GPR).finish();
760        func.build(block, opcode).def(once, GPR).finish();
761        func.build(block, opcode).def(other, GPR).finish();
762        for _ in 0..3 {
763            func.build(block, opcode).uses(busy, GPR).uses(other, GPR).finish();
764        }
765        func.build(block, opcode).uses(once, GPR).uses(busy, GPR).uses(other, GPR).finish();
766
767        let env = narrow(2);
768        let order = Order::of(&func);
769        let live = Live::of(&func, &order);
770        let pressure = Pressure::of(&func, &order, &live, &env);
771        let early = spill::choose(&func, &live, &pressure);
772        assert_eq!(early, [once]);
773        let assignment = placed(&func, &order, &live, &env, BUDGET, &early).expect("in budget");
774        let problems = check::check(&func, &order, &live, &assignment);
775        assert!(problems.is_empty(), "{}", check::report(&problems));
776        assert_eq!(named(assignment.place(once)), "slot");
777        assert_eq!(assignment.spilled(), 1);
778    }
779
780    #[test]
781    fn two_values_that_are_never_both_wanted_share_a_register() {
782        let mut names = Interner::new();
783        let mut func = Func::new(names.intern("f"));
784        let opcode = Opcode::new(names.intern("x64.nop"));
785        let block = func.create_block();
786        let first = func.new_vreg(GPR);
787        let second = func.new_vreg(GPR);
788        func.build(block, opcode).def(first, GPR).finish();
789        func.build(block, opcode).uses(first, GPR).finish();
790        func.build(block, opcode).def(second, GPR).finish();
791        func.build(block, opcode).uses(second, GPR).finish();
792
793        assert_eq!(places(&mut func, &env()), ["rax", "rax"]);
794    }
795
796    #[test]
797    fn the_value_read_most_often_keeps_its_register_though_it_is_wanted_longest() {
798        let mut names = Interner::new();
799        let mut func = Func::new(names.intern("f"));
800        let opcode = Opcode::new(names.intern("x64.nop"));
801        let block = func.create_block();
802        let busy = func.new_vreg(GPR);
803        let once = func.new_vreg(GPR);
804        func.build(block, opcode).def(busy, GPR).finish();
805        func.build(block, opcode).def(once, GPR).finish();
806        for _ in 0..6 {
807            func.build(block, opcode).uses(busy, GPR).finish();
808        }
809        func.build(block, opcode).uses(once, GPR).finish();
810        func.build(block, opcode).uses(busy, GPR).finish();
811
812        // With one register the linear scan sends the value that ends last to the stack, which is
813        // the one read seven times. Weighing them keeps that one and sends the one read once.
814        assert_eq!(linear(&func, &narrow(1)), ["slot", "rax"]);
815        assert_eq!(places(&mut func, &narrow(1)), ["rax", "slot"]);
816    }
817
818    #[test]
819    fn a_value_read_in_a_loop_keeps_its_register_over_one_read_outside_it() {
820        let mut names = Interner::new();
821        let mut func = Func::new(names.intern("f"));
822        let opcode = Opcode::new(names.intern("x64.nop"));
823        let entry = func.create_block();
824        let body = func.create_block();
825        let out = func.create_block();
826        let step = func.new_vreg(GPR);
827        let cold = func.new_vreg(GPR);
828        func.build(entry, opcode).def(cold, GPR).finish();
829        func.build(entry, opcode).def(step, GPR).finish();
830        *func.succs_mut(entry) = vec![BlockCall::to(body)];
831        func.build(body, opcode).uses(step, GPR).finish();
832        *func.succs_mut(body) = vec![BlockCall::to(body), BlockCall::to(out)];
833        func.set_weight(body, rucc_mir::Weight::parts(100 * rucc_mir::Weight::SCALE));
834        func.build(out, opcode).uses(cold, GPR).finish();
835        func.build(out, opcode).uses(step, GPR).finish();
836
837        // Read once in the loop and once after it, against a value read once after it: the count
838        // of reads is the same and it is how often the loop runs that decides.
839        assert_eq!(places(&mut func, &narrow(1)), ["rax", "slot"]);
840    }
841
842    #[test]
843    fn a_long_value_gives_its_register_back_to_a_short_busy_one() {
844        let mut names = Interner::new();
845        let mut func = Func::new(names.intern("f"));
846        let opcode = Opcode::new(names.intern("x64.nop"));
847        let block = func.create_block();
848        let long = func.new_vreg(GPR);
849        let short = func.new_vreg(GPR);
850        func.build(block, opcode).def(long, GPR).finish();
851        for _ in 0..4 {
852            func.build(block, opcode).finish();
853        }
854        func.build(block, opcode).def(short, GPR).finish();
855        for _ in 0..4 {
856            func.build(block, opcode).uses(short, GPR).finish();
857        }
858        func.build(block, opcode).uses(long, GPR).finish();
859
860        // The long one is placed first, since it is the harder to place, and then loses the one
861        // register to the short one, which weighs more, and has nowhere else to go.
862        assert_eq!(places(&mut func, &narrow(1)), ["slot", "rax"]);
863    }
864
865    #[test]
866    fn the_answer_of_a_two_address_instruction_goes_where_its_source_ends() {
867        let mut names = Interner::new();
868        let mut func = Func::new(names.intern("f"));
869        let opcode = Opcode::new(names.intern("x64.nop"));
870        let block = func.create_block();
871        let left = func.new_vreg(GPR);
872        let right = func.new_vreg(GPR);
873        let sum = func.new_vreg(GPR);
874        func.build(block, opcode).def(left, GPR).finish();
875        func.build(block, opcode).def(right, GPR).finish();
876        func.build(block, opcode)
877            .operand(Operand::write(sum, GPR).with(Constraint::Reuse(1)))
878            .uses(left, GPR)
879            .uses(right, GPR)
880            .finish();
881        func.build(block, opcode).uses(right, GPR).finish();
882        func.build(block, opcode).uses(sum, GPR).finish();
883
884        let places = places(&mut func, &env());
885        assert_eq!(places[2], places[0]);
886        assert_ne!(places[1], places[0]);
887    }
888
889    #[test]
890    fn a_source_placed_after_its_answer_goes_where_the_answer_is() {
891        let mut names = Interner::new();
892        let mut func = Func::new(names.intern("f"));
893        let opcode = Opcode::new(names.intern("x64.nop"));
894        let block = func.create_block();
895        let left = func.new_vreg(GPR);
896        let sum = func.new_vreg(GPR);
897        func.build(block, opcode).def(left, GPR).finish();
898        func.build(block, opcode)
899            .operand(Operand::write(sum, GPR).with(Constraint::Reuse(1)))
900            .uses(left, GPR)
901            .finish();
902        for _ in 0..6 {
903            func.build(block, opcode).uses(sum, GPR).finish();
904        }
905
906        // The answer is the longer of the two, so it is placed first, and the source finds it.
907        let places = places(&mut func, &env());
908        assert_eq!(places[0], places[1]);
909    }
910
911    #[test]
912    fn an_answer_that_commutes_stays_where_the_loop_passes_it() {
913        let mut names = Interner::new();
914        let mut func = Func::new(names.intern("f"));
915        let opcode = Opcode::new(names.intern("x64.nop"));
916        let entry = func.create_block();
917        let head = func.create_block();
918        let out = func.create_block();
919        let seed = func.new_vreg(GPR);
920        let total = func.new_vreg(GPR);
921        let term = func.new_vreg(GPR);
922        let next = func.new_vreg(GPR);
923        func.build(entry, opcode).def(seed, GPR).finish();
924        *func.succs_mut(entry) = vec![BlockCall::with(head, vec![seed])];
925        func.params_mut(head).push(Param { reg: total, class: GPR });
926        func.build(head, opcode).def(term, GPR).finish();
927        func.build(head, opcode)
928            .flags(Flags::COMMUTES)
929            .operand(Operand::write(next, GPR).with(Constraint::Reuse(1)))
930            .uses(total, GPR)
931            .uses(term, GPR)
932            .finish();
933        *func.succs_mut(head) = vec![BlockCall::with(head, vec![next]), BlockCall::to(out)];
934
935        let places = places(&mut func, &env());
936        assert_eq!(places[3], places[1]);
937    }
938
939    #[test]
940    fn an_answer_whose_first_source_lives_on_goes_where_the_second_one_ends() {
941        let mut names = Interner::new();
942        let mut func = Func::new(names.intern("f"));
943        let opcode = Opcode::new(names.intern("x64.nop"));
944        let block = func.create_block();
945        let base = func.new_vreg(GPR);
946        let entry = func.new_vreg(GPR);
947        let target = func.new_vreg(GPR);
948        func.build(block, opcode).def(base, GPR).finish();
949        func.build(block, opcode).def(entry, GPR).finish();
950        func.build(block, opcode)
951            .flags(Flags::COMMUTES)
952            .operand(Operand::write(target, GPR).with(Constraint::Reuse(1)))
953            .uses(base, GPR)
954            .uses(entry, GPR)
955            .finish();
956        func.build(block, opcode).uses(target, GPR).finish();
957        func.build(block, opcode).uses(base, GPR).finish();
958
959        let places = places(&mut func, &env());
960        assert_eq!(places[2], places[1]);
961        assert_ne!(places[2], places[0]);
962    }
963
964    #[test]
965    fn an_offset_added_to_a_base_the_loop_reads_again_goes_where_the_answer_went() {
966        let mut names = Interner::new();
967        let mut func = Func::new(names.intern("f"));
968        let opcode = Opcode::new(names.intern("x64.nop"));
969        let entry = func.create_block();
970        let head = func.create_block();
971        let out = func.create_block();
972        let base = func.new_vreg(GPR);
973        let at = func.new_vreg(GPR);
974        let offset = func.new_vreg(GPR);
975        let target = func.new_vreg(GPR);
976        let later = func.new_vreg(GPR);
977        func.build(entry, opcode).def(base, GPR).finish();
978        *func.succs_mut(entry) = vec![BlockCall::to(head)];
979        func.build(head, opcode).def(at, GPR).finish();
980        func.build(head, opcode).def(offset, GPR).uses(base, GPR).uses(at, GPR).finish();
981        func.build(head, opcode)
982            .flags(Flags::COMMUTES)
983            .operand(Operand::write(target, GPR).with(Constraint::Reuse(1)))
984            .uses(base, GPR)
985            .uses(offset, GPR)
986            .finish();
987        func.build(head, opcode).def(later, GPR).finish();
988        func.build(head, opcode).uses(target, GPR).finish();
989        for _ in 0..4 {
990            func.build(head, opcode).finish();
991        }
992        func.build(head, opcode).uses(later, GPR).finish();
993        *func.succs_mut(head) = vec![BlockCall::to(head), BlockCall::to(out)];
994
995        // The answer is placed before the offset, which is wanted over less of the line, and the
996        // base the loop reads again has the only register the answer could have followed. The
997        // offset is placed last and finds the register `later` has after the add free before it,
998        // which is where it went before it looked at the answer, and the answer could not then be
999        // moved to it. tamnd/rucc#2064.
1000        let places = places(&mut func, &env());
1001        assert_eq!(places[3], places[2]);
1002        assert_ne!(places[3], places[0]);
1003    }
1004
1005    #[test]
1006    fn a_register_an_instruction_insists_on_is_left_to_the_value_it_names() {
1007        let mut names = Interner::new();
1008        let mut func = Func::new(names.intern("f"));
1009        let opcode = Opcode::new(names.intern("x64.nop"));
1010        let block = func.create_block();
1011        let kept = func.new_vreg(GPR);
1012        let passed = func.new_vreg(GPR);
1013        func.build(block, opcode).def(kept, GPR).finish();
1014        func.build(block, opcode).def(passed, GPR).finish();
1015        func.build(block, opcode)
1016            .operand(Operand::read(passed, GPR).with(Constraint::Fixed(rucc_target::x86_64::RAX)))
1017            .finish();
1018        func.build(block, opcode).uses(kept, GPR).finish();
1019
1020        let places = places(&mut func, &env());
1021        assert_eq!(places[1], "rax");
1022        assert_ne!(places[0], "rax");
1023    }
1024
1025    #[test]
1026    fn a_value_only_memory_can_hold_goes_to_the_stack() {
1027        let mut names = Interner::new();
1028        let mut func = Func::new(names.intern("f"));
1029        let opcode = Opcode::new(names.intern("x64.nop"));
1030        let block = func.create_block();
1031        let value = func.new_vreg(GPR);
1032        func.build(block, opcode).def(value, GPR).finish();
1033        func.build(block, opcode)
1034            .operand(Operand::read(value, GPR).with(Constraint::Stack))
1035            .finish();
1036
1037        assert_eq!(places(&mut func, &env()), ["slot"]);
1038    }
1039
1040    #[test]
1041    fn a_function_over_the_budget_gets_the_linear_scan_answer() {
1042        let mut names = Interner::new();
1043        let mut func = Func::new(names.intern("f"));
1044        let opcode = Opcode::new(names.intern("x64.nop"));
1045        let block = func.create_block();
1046        let busy = func.new_vreg(GPR);
1047        let once = func.new_vreg(GPR);
1048        func.build(block, opcode).def(busy, GPR).finish();
1049        func.build(block, opcode).def(once, GPR).finish();
1050        for _ in 0..6 {
1051            func.build(block, opcode).uses(busy, GPR).finish();
1052        }
1053        func.build(block, opcode).uses(once, GPR).finish();
1054        func.build(block, opcode).uses(busy, GPR).finish();
1055
1056        let order = Order::of(&func);
1057        let live = Live::of(&func, &order);
1058        assert!(within(&func, &order, &live, &narrow(1), 0).is_none());
1059        let fallen = linear(&func, &narrow(1));
1060        assert_eq!(fallen, ["slot", "rax"]);
1061    }
1062
1063    #[test]
1064    fn the_cheaper_of_the_two_answers_is_the_one_kept() {
1065        let mut names = Interner::new();
1066        let mut func = Func::new(names.intern("f"));
1067        let opcode = Opcode::new(names.intern("x64.nop"));
1068        let block = func.create_block();
1069        let busy = func.new_vreg(GPR);
1070        let once = func.new_vreg(GPR);
1071        func.build(block, opcode).def(busy, GPR).finish();
1072        func.build(block, opcode).def(once, GPR).finish();
1073        for _ in 0..6 {
1074            func.build(block, opcode).uses(busy, GPR).finish();
1075        }
1076        func.build(block, opcode).uses(once, GPR).finish();
1077        func.build(block, opcode).uses(busy, GPR).finish();
1078
1079        let order = Order::of(&func);
1080        let live = Live::of(&func, &order);
1081        let env = narrow(1);
1082        let linear = assign::assign(&func, &order, &live, &env);
1083        let ours = within(&func, &order, &live, &env, BUDGET).expect("inside the budget");
1084        // The linear scan puts `busy` on the stack, which is its write and its seven reads. This
1085        // puts `once` there, which is one write and one read.
1086        let once_through = u128::from(func[block].weight.raw());
1087        assert_eq!(cost(&func, &order, &linear), 8 * once_through);
1088        assert_eq!(cost(&func, &order, &ours), 2 * once_through);
1089        let kept = assign(&func, &order, &live, &env);
1090        assert_eq!(kept.place(busy), ours.place(busy));
1091        assert_eq!(kept.place(once), ours.place(once));
1092    }
1093
1094    #[test]
1095    fn a_copy_left_between_a_two_address_answer_and_its_source_is_counted() {
1096        let mut names = Interner::new();
1097        let mut func = Func::new(names.intern("f"));
1098        let opcode = Opcode::new(names.intern("x64.nop"));
1099        let block = func.create_block();
1100        let left = func.new_vreg(GPR);
1101        let right = func.new_vreg(GPR);
1102        let sum = func.new_vreg(GPR);
1103        func.build(block, opcode).def(left, GPR).finish();
1104        func.build(block, opcode).def(right, GPR).finish();
1105        func.build(block, opcode)
1106            .operand(Operand::write(sum, GPR).with(Constraint::Reuse(1)))
1107            .uses(left, GPR)
1108            .uses(right, GPR)
1109            .finish();
1110        func.build(block, opcode).uses(sum, GPR).finish();
1111
1112        let order = Order::of(&func);
1113        let live = Live::of(&func, &order);
1114        let chosen = assign(&func, &order, &live, &env());
1115        assert_eq!(cost(&func, &order, &chosen), 0);
1116        let mut apart = chosen.clone();
1117        apart.put(sum, chosen.place(right).expect("a place for right"));
1118        assert_eq!(cost(&func, &order, &apart), u128::from(func[block].weight.raw()));
1119    }
1120
1121    #[test]
1122    fn many_values_in_few_registers_come_out_as_something_the_machine_can_run() {
1123        let mut names = Interner::new();
1124        let mut func = Func::new(names.intern("f"));
1125        let opcode = Opcode::new(names.intern("x64.nop"));
1126        let block = func.create_block();
1127        let regs: Vec<Reg> = (0..24).map(|_| func.new_vreg(GPR)).collect();
1128        for &reg in &regs {
1129            func.build(block, opcode).def(reg, GPR).finish();
1130        }
1131        for (at, &reg) in regs.iter().enumerate().rev() {
1132            for _ in 0..(at % 5) {
1133                func.build(block, opcode).uses(reg, GPR).finish();
1134            }
1135            func.build(block, opcode).uses(reg, GPR).finish();
1136        }
1137
1138        let places = places(&mut func, &narrow(4));
1139        assert_eq!(places.iter().filter(|place| *place != "slot").count(), 4);
1140    }
1141}