Skip to main content

iridia_exec/
mask.rs

1//! Sets of lanes: which invocations of a run are active. A pixel run is a quad or a few invocations in lockstep, its
2//! set an integer's bits; a thread group of up to 1024 invocations needs [`Wide`].
3
4/// A set of the lanes of a run.
5pub trait Mask: Copy + PartialEq + std::fmt::Debug {
6    /// The empty set.
7    fn none() -> Self;
8    /// Lanes `0..n`.
9    fn first(n: usize) -> Self;
10    /// Whether `lane` is in the set.
11    fn has(self, lane: usize) -> bool;
12    /// Add `lane`; a lane past the width of the set changes nothing.
13    fn insert(&mut self, lane: usize);
14    /// Remove `lane`; a lane past the width of the set changes nothing.
15    fn remove(&mut self, lane: usize);
16    /// The lanes in both sets.
17    fn and(self, other: Self) -> Self;
18    /// The lanes in either set.
19    fn or(self, other: Self) -> Self;
20    /// The lanes of `self` that are not in `other`.
21    fn minus(self, other: Self) -> Self;
22    /// Whether the set has no lane.
23    fn is_empty(self) -> bool;
24}
25
26/// The shift count of a lane: a lane past `u32` is past every mask.
27fn shift(lane: usize) -> u32 {
28    u32::try_from(lane).unwrap_or(u32::MAX)
29}
30
31/// A set of up to as many lanes as the integer has bits. A lane past the width is in no set, and inserting it or removing
32/// it changes none.
33macro_rules! bits {
34    ($($t:ty),*) => {$(
35        impl Mask for $t {
36            fn none() -> Self {
37                0
38            }
39            fn first(n: usize) -> Self {
40                (1 as $t).checked_shl(shift(n)).map_or(<$t>::MAX, |one| one.wrapping_sub(1))
41            }
42            fn has(self, lane: usize) -> bool {
43                self.checked_shr(shift(lane)).is_some_and(|rest| rest & 1 != 0)
44            }
45            fn insert(&mut self, lane: usize) {
46                if let Some(bit) = (1 as $t).checked_shl(shift(lane)) {
47                    *self |= bit;
48                }
49            }
50            fn remove(&mut self, lane: usize) {
51                if let Some(bit) = (1 as $t).checked_shl(shift(lane)) {
52                    *self &= !bit;
53                }
54            }
55            fn and(self, other: Self) -> Self {
56                self & other
57            }
58            fn or(self, other: Self) -> Self {
59                self | other
60            }
61            fn minus(self, other: Self) -> Self {
62                self & !other
63            }
64            fn is_empty(self) -> bool {
65                self == 0
66            }
67        }
68    )*};
69}
70bits!(u8, u16, u32);
71
72/// Up to 1024 lanes. A lane past them is in no set, and inserting it or removing it changes none.
73#[derive(Clone, Copy, Debug, PartialEq, Eq)]
74pub struct Wide([u64; 16]);
75
76impl Wide {
77    /// The word of `lane` and the bit in it.
78    fn place(lane: usize) -> (usize, u64) {
79        (lane / 64, 1u64.checked_shl(shift(lane % 64)).unwrap_or(0))
80    }
81    /// Each word of `self` combined with the word of `other` at the same position.
82    fn zip(self, other: Self, combine: impl Fn(u64, u64) -> u64) -> Self {
83        let mut words = self.0;
84        for (word, other) in words.iter_mut().zip(other.0) {
85            *word = combine(*word, other);
86        }
87        Wide(words)
88    }
89}
90
91impl Mask for Wide {
92    fn none() -> Self {
93        Wide([0; 16])
94    }
95    fn first(n: usize) -> Self {
96        let mut w = [0u64; 16];
97        for (i, word) in w.iter_mut().enumerate() {
98            // the lanes of this word that are among the first n
99            let taken = n.saturating_sub(i.saturating_mul(64));
100            *word = 1u64
101                .checked_shl(shift(taken))
102                .map_or(u64::MAX, |one| one.wrapping_sub(1));
103        }
104        Wide(w)
105    }
106    fn has(self, lane: usize) -> bool {
107        let (word, bit) = Self::place(lane);
108        self.0.get(word).is_some_and(|w| w & bit != 0)
109    }
110    fn insert(&mut self, lane: usize) {
111        let (word, bit) = Self::place(lane);
112        if let Some(w) = self.0.get_mut(word) {
113            *w |= bit;
114        }
115    }
116    fn remove(&mut self, lane: usize) {
117        let (word, bit) = Self::place(lane);
118        if let Some(w) = self.0.get_mut(word) {
119            *w &= !bit;
120        }
121    }
122    fn and(self, other: Self) -> Self {
123        self.zip(other, |a, b| a & b)
124    }
125    fn or(self, other: Self) -> Self {
126        self.zip(other, |a, b| a | b)
127    }
128    fn minus(self, other: Self) -> Self {
129        self.zip(other, |a, b| a & !b)
130    }
131    fn is_empty(self) -> bool {
132        self.0.iter().all(|w| *w == 0)
133    }
134}