celox_backend_common/regalloc/
parallel_copy.rs

1//! Target-independent scheduling of edge-local parallel copies.
2
3use std::collections::{BTreeMap, BTreeSet, VecDeque};
4use std::fmt;
5
6/// Writable physical location in a parallel assignment.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
8pub enum CopyDestination<R> {
9    Register(R),
10    Stack(i32),
11}
12
13/// Readable physical location or constant in a parallel assignment.
14#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
15pub enum CopySource<R> {
16    Register(R),
17    Stack(i32),
18    Immediate(u64),
19}
20
21/// One simultaneous assignment row.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct ParallelCopy<R> {
24    pub destination: CopyDestination<R>,
25    pub source: CopySource<R>,
26}
27
28impl<R: PartialEq> ParallelCopy<R> {
29    pub fn is_identity(&self) -> bool {
30        matches!(
31            (&self.destination, &self.source),
32            (
33                CopyDestination::Register(destination),
34                CopySource::Register(source)
35            ) if destination == source
36        ) || matches!(
37            (&self.destination, &self.source),
38            (CopyDestination::Stack(destination), CopySource::Stack(source))
39                if destination == source
40        )
41    }
42}
43
44/// One sequential operation implementing a parallel assignment.
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub enum CopyOperation<R> {
47    Move {
48        destination: CopyDestination<R>,
49        source: CopySource<R>,
50    },
51    SwapRegisters {
52        left: R,
53        right: R,
54    },
55    SaveTemporary(CopyDestination<R>),
56    RestoreTemporary(CopyDestination<R>),
57}
58
59/// Work counters retained for allocator diagnostics and regression tests.
60#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
61pub struct CopyResolutionWork {
62    pub effective_copies: usize,
63    pub direct_moves: usize,
64    pub register_swaps: usize,
65    pub cycle_breaks: usize,
66    pub temporary_cycle_breaks: usize,
67    pub ready_queue_pops: usize,
68    pub dependency_releases: usize,
69}
70
71/// Dependency-ordered realization of one parallel assignment.
72#[derive(Debug, Clone, PartialEq, Eq)]
73pub struct CopyResolution<R> {
74    pub operations: Vec<CopyOperation<R>>,
75    pub work: CopyResolutionWork,
76}
77
78/// Malformed copy rows or an inconsistent resolver state.
79#[derive(Debug, Clone, PartialEq, Eq)]
80pub struct CopyResolutionError<R> {
81    pub rule: &'static str,
82    pub destination: Option<CopyDestination<R>>,
83    pub message: String,
84}
85
86impl<R> CopyResolutionError<R> {
87    fn new(rule: &'static str, message: impl Into<String>) -> Self {
88        Self {
89            rule,
90            destination: None,
91            message: message.into(),
92        }
93    }
94
95    fn at(mut self, destination: CopyDestination<R>) -> Self {
96        self.destination = Some(destination);
97        self
98    }
99}
100
101impl<R: fmt::Debug> fmt::Display for CopyResolutionError<R> {
102    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
103        write!(formatter, "parallel copy [{}]", self.rule)?;
104        if let Some(destination) = &self.destination {
105            write!(formatter, " at {destination:?}")?;
106        }
107        write!(formatter, ": {}", self.message)
108    }
109}
110
111impl<R: fmt::Debug> std::error::Error for CopyResolutionError<R> {}
112
113#[derive(Debug, Clone, Copy, PartialEq, Eq)]
114enum PendingSource<R> {
115    Value(CopySource<R>),
116    Temporary,
117}
118
119#[derive(Debug, Clone, Copy)]
120struct PendingCopy<R> {
121    destination: CopyDestination<R>,
122    source: PendingSource<R>,
123    pending: bool,
124}
125
126fn source_as_destination<R: Copy>(source: CopySource<R>) -> Option<CopyDestination<R>> {
127    match source {
128        CopySource::Register(register) => Some(CopyDestination::Register(register)),
129        CopySource::Stack(slot) => Some(CopyDestination::Stack(slot)),
130        CopySource::Immediate(_) => None,
131    }
132}
133
134fn register_cycle<R>(
135    copies: &[PendingCopy<R>],
136    destination_index: &BTreeMap<CopyDestination<R>, usize>,
137    start: usize,
138) -> Option<(Vec<usize>, Vec<(R, R)>)>
139where
140    R: Copy + Ord,
141{
142    let CopyDestination::Register(start_register) = copies.get(start)?.destination else {
143        return None;
144    };
145    let mut current = start_register;
146    let mut members = Vec::new();
147    let mut swaps = Vec::new();
148    let mut visited = BTreeSet::new();
149    loop {
150        if !visited.insert(current) {
151            return None;
152        }
153        let destination = CopyDestination::Register(current);
154        let &index = destination_index.get(&destination)?;
155        let copy = copies.get(index)?;
156        if !copy.pending {
157            return None;
158        }
159        let PendingSource::Value(CopySource::Register(source)) = copy.source else {
160            return None;
161        };
162        members.push(index);
163        if source == start_register {
164            return (members.len() >= 2).then_some((members, swaps));
165        }
166        swaps.push((current, source));
167        current = source;
168    }
169}
170
171/// Schedule copies without clobbering any still-needed source.
172///
173/// Acyclic rows become direct moves, two-register cycles become one swap, and
174/// longer or stack cycles use one target-provided temporary location.
175pub fn resolve_parallel_copies<R>(
176    rows: &[ParallelCopy<R>],
177) -> Result<CopyResolution<R>, CopyResolutionError<R>>
178where
179    R: Copy + Ord + fmt::Debug,
180{
181    let mut destination_owner = BTreeSet::new();
182    for row in rows {
183        if !destination_owner.insert(row.destination) {
184            return Err(CopyResolutionError::new(
185                "PARALLEL_COPY.NON_UNIQUE_DESTINATION",
186                "parallel assignment writes one destination more than once",
187            )
188            .at(row.destination));
189        }
190    }
191
192    let mut copies = rows
193        .iter()
194        .filter(|row| !row.is_identity())
195        .map(|row| PendingCopy {
196            destination: row.destination,
197            source: PendingSource::Value(row.source),
198            pending: true,
199        })
200        .collect::<Vec<_>>();
201    let mut work = CopyResolutionWork {
202        effective_copies: copies.len(),
203        ..CopyResolutionWork::default()
204    };
205    if copies.is_empty() {
206        return Ok(CopyResolution {
207            operations: Vec::new(),
208            work,
209        });
210    }
211
212    let destination_index = copies
213        .iter()
214        .enumerate()
215        .map(|(index, copy)| (copy.destination, index))
216        .collect::<BTreeMap<_, _>>();
217    let mut readers = BTreeMap::<CopyDestination<R>, BTreeSet<usize>>::new();
218    for (index, copy) in copies.iter().enumerate() {
219        let PendingSource::Value(source) = copy.source else {
220            continue;
221        };
222        if let Some(location) = source_as_destination(source) {
223            readers.entry(location).or_default().insert(index);
224        }
225    }
226
227    let mut ready = VecDeque::new();
228    let mut queued = vec![false; copies.len()];
229    for (index, copy) in copies.iter().enumerate() {
230        if !readers.contains_key(&copy.destination) {
231            ready.push_back(index);
232            queued[index] = true;
233        }
234    }
235
236    let mut operations = Vec::with_capacity(copies.len());
237    let mut remaining = copies.len();
238    let mut temporary_live = false;
239    let mut cycle_search_start = 0usize;
240    while remaining != 0 {
241        while let Some(index) = ready.pop_front() {
242            queued[index] = false;
243            if !copies[index].pending {
244                continue;
245            }
246            work.ready_queue_pops += 1;
247            match copies[index].source {
248                PendingSource::Value(source) => {
249                    operations.push(CopyOperation::Move {
250                        destination: copies[index].destination,
251                        source,
252                    });
253                    work.direct_moves += 1;
254                    if let Some(location) = source_as_destination(source) {
255                        let released = readers.get_mut(&location).is_some_and(|location_readers| {
256                            location_readers.remove(&index);
257                            location_readers.is_empty()
258                        });
259                        if released {
260                            readers.remove(&location);
261                            work.dependency_releases += 1;
262                            if let Some(&writer) = destination_index.get(&location)
263                                && copies[writer].pending
264                                && !queued[writer]
265                            {
266                                ready.push_back(writer);
267                                queued[writer] = true;
268                            }
269                        }
270                    }
271                }
272                PendingSource::Temporary => {
273                    if !temporary_live {
274                        return Err(CopyResolutionError::new(
275                            "PARALLEL_COPY.TEMPORARY_STATE",
276                            "resolver attempted to restore an inactive temporary",
277                        )
278                        .at(copies[index].destination));
279                    }
280                    operations.push(CopyOperation::RestoreTemporary(copies[index].destination));
281                    temporary_live = false;
282                }
283            }
284            copies[index].pending = false;
285            remaining -= 1;
286        }
287
288        if remaining == 0 {
289            break;
290        }
291        if temporary_live {
292            return Err(CopyResolutionError::new(
293                "PARALLEL_COPY.TEMPORARY_STATE",
294                "resolver stalled while a temporary was live",
295            ));
296        }
297        let Some(cycle) = copies
298            .iter()
299            .enumerate()
300            .skip(cycle_search_start)
301            .find_map(|(index, copy)| copy.pending.then_some(index))
302        else {
303            return Err(CopyResolutionError::new(
304                "PARALLEL_COPY.RESOLVER_STATE",
305                "nonzero pending count has no pending row",
306            ));
307        };
308        cycle_search_start = cycle + 1;
309
310        if let Some((members, swaps)) = register_cycle(&copies, &destination_index, cycle)
311            && members.len() == 2
312        {
313            for (left, right) in swaps.iter().copied() {
314                operations.push(CopyOperation::SwapRegisters { left, right });
315            }
316            for &member in &members {
317                readers.remove(&copies[member].destination);
318                copies[member].pending = false;
319                queued[member] = false;
320            }
321            remaining -= members.len();
322            work.register_swaps += swaps.len();
323            work.cycle_breaks += 1;
324            work.dependency_releases += members.len();
325            continue;
326        }
327
328        let saved = copies[cycle].destination;
329        let saved_readers = readers.remove(&saved).unwrap_or_default();
330        if saved_readers.len() != 1 {
331            return Err(CopyResolutionError::new(
332                "PARALLEL_COPY.CYCLE_SHAPE",
333                format!("stalled cycle location has {} readers", saved_readers.len()),
334            )
335            .at(saved));
336        }
337        let reader = *saved_readers
338            .iter()
339            .next()
340            .expect("one cycle reader was checked above");
341        copies[reader].source = PendingSource::Temporary;
342        operations.push(CopyOperation::SaveTemporary(saved));
343        work.cycle_breaks += 1;
344        work.temporary_cycle_breaks += 1;
345        work.dependency_releases += 1;
346        temporary_live = true;
347        if !queued[cycle] {
348            ready.push_back(cycle);
349            queued[cycle] = true;
350        }
351    }
352
353    if temporary_live {
354        return Err(CopyResolutionError::new(
355            "PARALLEL_COPY.TEMPORARY_STATE",
356            "resolver left its temporary live",
357        ));
358    }
359    Ok(CopyResolution { operations, work })
360}
361
362#[cfg(test)]
363mod tests {
364    use super::*;
365
366    fn register_copy(destination: u8, source: u8) -> ParallelCopy<u8> {
367        ParallelCopy {
368            destination: CopyDestination::Register(destination),
369            source: CopySource::Register(source),
370        }
371    }
372
373    #[test]
374    fn orders_an_acyclic_chain_backwards() {
375        let result = resolve_parallel_copies(&[register_copy(0, 1), register_copy(2, 0)]).unwrap();
376        assert_eq!(
377            result.operations,
378            vec![
379                CopyOperation::Move {
380                    destination: CopyDestination::Register(2),
381                    source: CopySource::Register(0),
382                },
383                CopyOperation::Move {
384                    destination: CopyDestination::Register(0),
385                    source: CopySource::Register(1),
386                },
387            ]
388        );
389    }
390
391    #[test]
392    fn swaps_a_two_register_cycle() {
393        let result = resolve_parallel_copies(&[register_copy(0, 1), register_copy(1, 0)]).unwrap();
394        assert_eq!(result.work.register_swaps, 1);
395        assert!(matches!(
396            result.operations.as_slice(),
397            [CopyOperation::SwapRegisters { .. }]
398        ));
399    }
400
401    #[test]
402    fn breaks_a_long_cycle_with_one_temporary() {
403        let result = resolve_parallel_copies(&[
404            register_copy(0, 1),
405            register_copy(1, 2),
406            register_copy(2, 0),
407        ])
408        .unwrap();
409        assert_eq!(result.work.temporary_cycle_breaks, 1);
410        assert!(matches!(
411            result.operations.first(),
412            Some(CopyOperation::SaveTemporary(_))
413        ));
414        assert!(matches!(
415            result.operations.last(),
416            Some(CopyOperation::RestoreTemporary(_))
417        ));
418    }
419
420    #[test]
421    fn rejects_duplicate_destinations_even_when_one_row_is_identity() {
422        let error =
423            resolve_parallel_copies(&[register_copy(0, 0), register_copy(0, 1)]).unwrap_err();
424        assert_eq!(error.rule, "PARALLEL_COPY.NON_UNIQUE_DESTINATION");
425    }
426}
celox_backend_common/regalloc/parallel_copy.rs

celox_backend_common/regalloc/
parallel_copy.rs
