Skip to main content

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
35impl Backend {
36    /// The backend this processor runs: the best of [`Backend::available`],
37    /// found once.
38    #[must_use]
39    pub fn active() -> Self {
40        static ACTIVE: OnceLock<Backend> = OnceLock::new();
41        *ACTIVE.get_or_init(|| {
42            BEST_FIRST
43                .into_iter()
44                .find(|b| b.is_available())
45                .unwrap_or(Self::Portable)
46        })
47    }
48
49    /// The backends this processor can run, best first.
50    pub fn available() -> impl Iterator<Item = Self> {
51        BEST_FIRST.into_iter().filter(|b| b.is_available())
52    }
53
54    /// Whether this processor has the instructions of the backend.
55    #[must_use]
56    pub fn is_available(self) -> bool {
57        match self {
58            Self::Portable => true,
59            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
60            Self::Sse42 => std::arch::is_x86_feature_detected!("sse4.2"),
61            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
62            Self::Avx2 => std::arch::is_x86_feature_detected!("avx2"),
63            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
64            Self::Avx512 => {
65                std::arch::is_x86_feature_detected!("avx512f")
66                    && std::arch::is_x86_feature_detected!("avx2")
67            }
68            #[cfg(all(target_arch = "aarch64", target_endian = "little"))]
69            Self::Neon => std::arch::is_aarch64_feature_detected!("neon"),
70            _ => false,
71        }
72    }
73
74    /// Groups' numbers packed onto `out` ([`group_bytes`] of them each): `src`
75    /// holds [`GROUP`] numbers per group, each fitting `width` bits (`1..=32`).
76    ///
77    /// # Panics
78    ///
79    /// If the backend is not [available](Self::is_available), `width` is out
80    /// of range, or `out` is shorter than the groups' bytes.
81    pub fn pack_groups(self, src: &[[u32; GROUP]], width: u32, out: &mut [u8]) {
82        assert!(self.is_available(), "{self} is not available here");
83        assert!((1..=32).contains(&width), "a group is 1 to 32 bits wide");
84        let size = group_bytes(width);
85        assert!(out.len() >= src.len() * size, "a group's bytes");
86        let numbers = src.as_flattened();
87        let mut group = 0;
88        while group < src.len() {
89            let backend = self.for_groups(width, src.len() - group);
90            backend.pack_unchecked(&numbers[group * GROUP..], width, &mut out[group * size..]);
91            group += backend.groups_per_call();
92        }
93    }
94
95    /// Groups unpacked from `src` onto `out`.
96    ///
97    /// # Panics
98    ///
99    /// As [`pack_groups`](Self::pack_groups).
100    pub fn unpack_groups(self, src: &[u8], width: u32, out: &mut [[u32; GROUP]]) {
101        assert!(self.is_available(), "{self} is not available here");
102        assert!((1..=32).contains(&width), "a group is 1 to 32 bits wide");
103        let size = group_bytes(width);
104        assert!(src.len() >= out.len() * size, "a group's bytes");
105        let count = out.len();
106        let numbers = out.as_flattened_mut();
107        let mut group = 0;
108        while group < count {
109            let backend = self.for_groups(width, count - group);
110            backend.unpack_unchecked(&src[group * size..], width, &mut numbers[group * GROUP..]);
111            group += backend.groups_per_call();
112        }
113    }
114
115    /// The groups one call of the backend's kernels takes.
116    fn groups_per_call(self) -> usize {
117        match self {
118            Self::Avx512 => 2,
119            _ => 1,
120        }
121    }
122
123    /// The backend to use for the next call of `width`-bit groups when `left`
124    /// groups remain: this one, or a narrower one where it is faster. The
125    /// 512-bit kernels take two groups at a time, so a single group left goes
126    /// to AVX2; and at 32 bits a group is a copy, which AVX2 does faster than
127    /// shuffling two rows into one register and out again (measured).
128    fn for_groups(self, width: u32, left: usize) -> Self {
129        match self {
130            Self::Avx512 if left < self.groups_per_call() || width == u32::BITS => Self::Avx2,
131            other => other,
132        }
133    }
134
135    /// The kernels' work for a backend known to be available and arguments
136    /// known to fit: `groups_per_call` groups from the start of both.
137    pub(super) fn pack_unchecked(self, src: &[u32], width: u32, out: &mut [u8]) {
138        match self {
139            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
140            Self::Sse42 => super::sse42::pack(src, width, out),
141            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
142            Self::Avx2 => super::avx2::pack(src, width, out),
143            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
144            Self::Avx512 => super::avx512::pack(src, width, out),
145            #[cfg(all(target_arch = "aarch64", target_endian = "little"))]
146            Self::Neon => super::neon::pack(src, width, out),
147            _ => portable::pack(src, width, out),
148        }
149    }
150
151    /// The inverse of `pack_unchecked`.
152    pub(super) fn unpack_unchecked(self, src: &[u8], width: u32, out: &mut [u32]) {
153        match self {
154            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
155            Self::Sse42 => super::sse42::unpack(src, width, out),
156            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
157            Self::Avx2 => super::avx2::unpack(src, width, out),
158            #[cfg(all(target_arch = "x86_64", target_endian = "little"))]
159            Self::Avx512 => super::avx512::unpack(src, width, out),
160            #[cfg(all(target_arch = "aarch64", target_endian = "little"))]
161            Self::Neon => super::neon::unpack(src, width, out),
162            _ => portable::unpack(src, width, out),
163        }
164    }
165}
166
167impl fmt::Display for Backend {
168    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
169        f.write_str(match self {
170            Self::Portable => "portable",
171            Self::Sse42 => "SSE4.2",
172            Self::Avx2 => "AVX2",
173            Self::Avx512 => "AVX-512",
174            Self::Neon => "NEON",
175        })
176    }
177}