substrate_data/wire/frame/lanes/
backend.rs

1//! The instruction sets a group can be packed with, and which one this
2//! processor uses.
3
4use std::fmt;
5
6use super::{GROUP, group_bytes, portable};
7use crate::sync::OnceLock;
8
9/// An instruction set that packs and unpacks groups. Every backend writes and
10/// reads the same bytes; they differ only in speed.
11#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
12#[non_exhaustive]
13pub enum Backend {
14    /// Plain Rust, any processor.
15    Portable,
16    /// x86-64 SSE4.2: a row as two 128-bit vectors.
17    Sse42,
18    /// x86-64 AVX2: a row as one 256-bit vector.
19    Avx2,
20    /// x86-64 AVX-512: two rows of a group in one 512-bit vector.
21    Avx512,
22    /// `AArch64` NEON: a row as two 128-bit vectors.
23    Neon,
24}
25
26/// Every backend, best first.
27const BEST_FIRST: [Backend; 5] = [
28    Backend::Avx512,
29    Backend::Avx2,
30    Backend::Neon,
31    Backend::Sse42,
32    Backend::Portable,
33];
34
35/// The most 32-bit groups AVX-512 unpacks faster than AVX2: a copy of this many
36/// groups, in and out, still fits a first-level data cache (48 KiB).
37const COPY_AVX512_GROUPS: usize = 24;
38
39impl Backend {
40    /// The backend this processor runs: the best of [`Backend::available`],
41    /// found once.
42    #[must_use]
43    pub fn active() -> Self {
44        static ACTIVE: OnceLock<Backend> = OnceLock::new();
45        *ACTIVE.get_or_init(|| {
46            BEST_FIRST
47                .into_iter()
48                .find(|b| b.is_available())
49                .unwrap_or(Self::Portable)
50        })
51    }
52
53    /// The backends this processor can run, best first.
54    pub fn available() -> impl Iterator<Item = Self> {
55        BEST_FIRST.into_iter().filter(|b| b.is_available())
56    }
57
58    /// Whether this processor has the instructions of the backend.
59    #[must_use]
60    pub fn is_available(self) -> bool {
61        match self {
62            Self::Portable => true,
63            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
64            Self::Sse42 => std::arch::is_x86_feature_detected!("sse4.2"),
65            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
66            Self::Avx2 => std::arch::is_x86_feature_detected!("avx2"),
67            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
68            Self::Avx512 => {
69                std::arch::is_x86_feature_detected!("avx512f")
70                    && std::arch::is_x86_feature_detected!("avx2")
71            }
72            #[cfg(all(target_arch = "aarch64", target_endian = "little"))]
73            Self::Neon => std::arch::is_aarch64_feature_detected!("neon"),
74            _ => false,
75        }
76    }
77
78    /// Groups' numbers packed onto `out` ([`group_bytes`] of them each): `src`
79    /// holds [`GROUP`] numbers per group, each fitting `width` bits (`1..=32`).
80    ///
81    /// # Panics
82    ///
83    /// If the backend is not [available](Self::is_available), `width` is out
84    /// of range, or `out` is shorter than the groups' bytes.
85    pub fn pack_groups(self, src: &[[u32; GROUP]], width: u32, out: &mut [u8]) {
86        assert!(self.is_available(), "{self} is not available here");
87        assert!((1..=32).contains(&width), "a group is 1 to 32 bits wide");
88        let size = group_bytes(width);
89        assert!(out.len() >= src.len() * size, "a group's bytes");
90        let numbers = src.as_flattened();
91        let mut group = 0;
92        while group < src.len() {
93            let backend = self.for_pack(width, src.len() - group);
94            backend.pack_unchecked(&numbers[group * GROUP..], width, &mut out[group * size..]);
95            group += backend.groups_per_call();
96        }
97    }
98
99    /// Groups unpacked from `src` onto `out`.
100    ///
101    /// # Panics
102    ///
103    /// As [`pack_groups`](Self::pack_groups).
104    pub fn unpack_groups(self, src: &[u8], width: u32, out: &mut [[u32; GROUP]]) {
105        assert!(self.is_available(), "{self} is not available here");
106        assert!((1..=32).contains(&width), "a group is 1 to 32 bits wide");
107        let size = group_bytes(width);
108        assert!(src.len() >= out.len() * size, "a group's bytes");
109        let count = out.len();
110        let numbers = out.as_flattened_mut();
111        let mut group = 0;
112        while group < count {
113            let backend = self.for_unpack(width, count - group);
114            backend.unpack_unchecked(&src[group * size..], width, &mut numbers[group * GROUP..]);
115            group += backend.groups_per_call();
116        }
117    }
118
119    /// The groups one call of the backend's kernels takes.
120    fn groups_per_call(self) -> usize {
121        match self {
122            Self::Avx512 => 2,
123            _ => 1,
124        }
125    }
126
127    /// The backend to use for the next call of packing `width`-bit groups when
128    /// `left` groups remain: this one, or a narrower one where it is faster
129    /// (measured, `docs/performance.md`). The 512-bit kernels take two groups
130    /// at a time, so a single group left goes to AVX2; and at 32 bits a group
131    /// is a copy, which AVX2 does 1.2 to 1.5 times faster than shuffling two
132    /// rows into one register and out again.
133    fn for_pack(self, width: u32, left: usize) -> Self {
134        match self {
135            Self::Avx512 if left < self.groups_per_call() || width == u32::BITS => Self::Avx2,
136            other => other,
137        }
138    }
139
140    /// As [`for_pack`](Self::for_pack), for unpacking. At 32 bits AVX-512 is
141    /// 1.6 to 2.2 times faster than AVX2 while the groups fit the first-level
142    /// cache (up to [`COPY_AVX512_GROUPS`]), and no better beyond it, where
143    /// the copy is bound by memory and the winner depends on how the buffers
144    /// fall; AVX2 is the one that is never much worse.
145    fn for_unpack(self, width: u32, left: usize) -> Self {
146        match self {
147            Self::Avx512
148                if left < self.groups_per_call()
149                    || (width == u32::BITS && left > COPY_AVX512_GROUPS) =>
150            {
151                Self::Avx2
152            }
153            other => other,
154        }
155    }
156
157    /// The kernels' work for a backend known to be available and arguments
158    /// known to fit: `groups_per_call` groups from the start of both.
159    pub(super) fn pack_unchecked(self, src: &[u32], width: u32, out: &mut [u8]) {
160        match self {
161            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
162            Self::Sse42 => super::sse42::pack(src, width, out),
163            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
164            Self::Avx2 => super::avx2::pack(src, width, out),
165            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
166            Self::Avx512 => super::avx512::pack(src, width, out),
167            #[cfg(all(target_arch = "aarch64", target_endian = "little"))]
168            Self::Neon => super::neon::pack(src, width, out),
169            _ => portable::pack(src, width, out),
170        }
171    }
172
173    /// The inverse of `pack_unchecked`.
174    pub(super) fn unpack_unchecked(self, src: &[u8], width: u32, out: &mut [u32]) {
175        match self {
176            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
177            Self::Sse42 => super::sse42::unpack(src, width, out),
178            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
179            Self::Avx2 => super::avx2::unpack(src, width, out),
180            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
181            Self::Avx512 => super::avx512::unpack(src, width, out),
182            #[cfg(all(target_arch = "aarch64", target_endian = "little"))]
183            Self::Neon => super::neon::unpack(src, width, out),
184            _ => portable::unpack(src, width, out),
185        }
186    }
187}
188
189impl fmt::Display for Backend {
190    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
191        f.write_str(match self {
192            Self::Portable => "portable",
193            Self::Sse42 => "SSE4.2",
194            Self::Avx2 => "AVX2",
195            Self::Avx512 => "AVX-512",
196            Self::Neon => "NEON",
197        })
198    }
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204
205    #[test]
206    fn the_512_bit_kernels_hand_a_lone_group_to_avx2() {
207        assert_eq!(Backend::Avx512.for_pack(7, 1), Backend::Avx2);
208        assert_eq!(Backend::Avx512.for_unpack(7, 1), Backend::Avx2);
209        assert_eq!(Backend::Avx512.for_pack(7, 2), Backend::Avx512);
210        assert_eq!(Backend::Avx512.for_unpack(7, 2), Backend::Avx512);
211    }
212
213    #[test]
214    fn a_full_width_copy_is_packed_by_avx2_and_unpacked_by_avx512_while_it_fits_l1() {
215        assert_eq!(Backend::Avx512.for_pack(32, 8), Backend::Avx2);
216        assert_eq!(
217            Backend::Avx512.for_unpack(32, COPY_AVX512_GROUPS),
218            Backend::Avx512
219        );
220        assert_eq!(
221            Backend::Avx512.for_unpack(32, COPY_AVX512_GROUPS + 1),
222            Backend::Avx2
223        );
224    }
225
226    #[test]
227    fn narrower_backends_are_never_redirected() {
228        for backend in [Backend::Portable, Backend::Avx2] {
229            assert_eq!(backend.for_pack(32, 1), backend);
230            assert_eq!(backend.for_unpack(32, 100), backend);
231        }
232    }
233}
substrate_data/wire/frame/lanes/backend.rs

substrate_data/wire/frame/lanes/
backend.rs
