Skip to main content

structured_zstd/
cpu_kernel.rs

1//! CPU kernel dispatch — single detect+match at the dispatch site,
2//! propagated through the inner pipeline as a generic parameter so
3//! leaf hot-path code monomorphises against the chosen kernel.
4//!
5//! See issue #247 for the architecture rationale: per-subsystem
6//! dispatch scatters the choice across HUF / FSE / SIMD-copy
7//! independently and pays the cost N times per call. Lifting the
8//! dispatch to the outermost feasible call site collapses it to one
9//! detect there; the inner leaf-hot-path ops then route through
10//! `K::method` calls on the chosen kernel zero-sized type.
11//!
12//! Current wiring (as of #247 Part 2): the only active dispatch site
13//! is `decoding::literals_section_decoder::decompress_literals`,
14//! which `match`es `detect_cpu_kernel()` and routes into per-K
15//! `decompress_literals_*` `#[target_feature]` wrappers. The full
16//! pipeline-wide propagation envisioned in the issue (FrameDecoder /
17//! FrameCompressor entry, sequence executor, match copy) is
18//! incremental; subsequent tiers extend the dispatch surface without
19//! changing this trait or the kernel ZSTs.
20//!
21//! Structure code (block loop, FCS check, offset history, repeat
22//! semantics) stays single-impl and only carries `K` as a phantom on
23//! the outer function. Monomorphisation specialises ONLY the bodies
24//! that actually differ per ISA — `mask_lower_bits`, `huf_burst`,
25//! `copy_chunk`, etc.
26
27#[cfg(feature = "std")]
28use std::sync::OnceLock;
29
30/// An instruction-set level the codec's kernels may be limited to, from the
31/// portable scalar code up. The x86 levels (`Sse2` through `Avx512`) and the
32/// aarch64 ones (`Neon`, `Sve`) are separate ladders; `Scalar` sits under
33/// both.
34///
35/// Used with [`set_cpu_ceiling`] to compare kernel tiers on one machine.
36#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
37#[non_exhaustive]
38pub enum CpuLevel {
39    /// Portable code only.
40    Scalar,
41    /// x86 SSE2.
42    Sse2,
43    /// x86 SSE4.2.
44    Sse42,
45    /// x86 BMI2, with everything below it.
46    Bmi2,
47    /// x86 AVX2 (and BMI2), the x86-64-v3 kernels.
48    Avx2,
49    /// x86 AVX-512 (the VBMI2 kernels, when built with `kernel-vbmi2`).
50    Avx512,
51    /// aarch64 NEON.
52    Neon,
53    /// aarch64 SVE and SVE2.
54    Sve,
55}
56
57impl CpuLevel {
58    /// Every level, lowest first within each ladder.
59    pub const ALL: [CpuLevel; 8] = [
60        CpuLevel::Scalar,
61        CpuLevel::Sse2,
62        CpuLevel::Sse42,
63        CpuLevel::Bmi2,
64        CpuLevel::Avx2,
65        CpuLevel::Avx512,
66        CpuLevel::Neon,
67        CpuLevel::Sve,
68    ];
69
70    /// The level's name as [`FromStr`](core::str::FromStr) reads it.
71    ///
72    /// # Examples
73    /// ```
74    /// use structured_zstd::CpuLevel;
75    ///
76    /// assert_eq!(CpuLevel::Sse42.name(), "sse4.2");
77    /// assert_eq!("avx2".parse::<CpuLevel>(), Ok(CpuLevel::Avx2));
78    /// ```
79    pub const fn name(self) -> &'static str {
80        match self {
81            CpuLevel::Scalar => "scalar",
82            CpuLevel::Sse2 => "sse2",
83            CpuLevel::Sse42 => "sse4.2",
84            CpuLevel::Bmi2 => "bmi2",
85            CpuLevel::Avx2 => "avx2",
86            CpuLevel::Avx512 => "avx512",
87            CpuLevel::Neon => "neon",
88            CpuLevel::Sve => "sve",
89        }
90    }
91
92    /// Which ladder the level is on, and its rung: 0 for `Scalar` on either.
93    const fn ladder(self) -> (Ladder, u8) {
94        match self {
95            CpuLevel::Scalar => (Ladder::Any, 0),
96            CpuLevel::Sse2 => (Ladder::X86, 1),
97            CpuLevel::Sse42 => (Ladder::X86, 2),
98            CpuLevel::Bmi2 => (Ladder::X86, 3),
99            CpuLevel::Avx2 => (Ladder::X86, 4),
100            CpuLevel::Avx512 => (Ladder::X86, 5),
101            CpuLevel::Neon => (Ladder::Arm, 1),
102            CpuLevel::Sve => (Ladder::Arm, 2),
103        }
104    }
105
106    /// Whether the level exists on the architecture this build targets.
107    const fn native(self) -> bool {
108        match self.ladder().0 {
109            Ladder::Any => true,
110            Ladder::X86 => cfg!(any(target_arch = "x86", target_arch = "x86_64")),
111            Ladder::Arm => cfg!(target_arch = "aarch64"),
112        }
113    }
114}
115
116impl core::fmt::Display for CpuLevel {
117    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
118        f.write_str(self.name())
119    }
120}
121
122/// A name that is not a [`CpuLevel`].
123#[derive(Copy, Clone, Debug, PartialEq, Eq)]
124pub struct UnknownCpuLevel;
125
126impl core::fmt::Display for UnknownCpuLevel {
127    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
128        f.write_str("unknown CPU level; expected one of ")?;
129        for (index, level) in CpuLevel::ALL.iter().enumerate() {
130            if index > 0 {
131                f.write_str(", ")?;
132            }
133            f.write_str(level.name())?;
134        }
135        Ok(())
136    }
137}
138
139#[cfg(feature = "std")]
140impl std::error::Error for UnknownCpuLevel {}
141
142impl core::str::FromStr for CpuLevel {
143    type Err = UnknownCpuLevel;
144
145    /// The names [`CpuLevel::name`] gives, case-insensitive, plus the spellings
146    /// `sse42`, `avx-512` and `avx512f`.
147    fn from_str(text: &str) -> Result<Self, Self::Err> {
148        let lower = |expected: &str| text.eq_ignore_ascii_case(expected);
149        CpuLevel::ALL
150            .into_iter()
151            .find(|level| lower(level.name()))
152            .or_else(|| {
153                if lower("sse42") {
154                    Some(CpuLevel::Sse42)
155                } else if lower("avx-512") || lower("avx512f") {
156                    Some(CpuLevel::Avx512)
157                } else {
158                    None
159                }
160            })
161            .ok_or(UnknownCpuLevel)
162    }
163}
164
165#[derive(Copy, Clone, PartialEq, Eq)]
166enum Ladder {
167    Any,
168    X86,
169    Arm,
170}
171
172/// Why [`set_cpu_ceiling`] could not set the ceiling.
173#[derive(Copy, Clone, Debug, PartialEq, Eq)]
174#[non_exhaustive]
175pub enum CpuCeilingError {
176    /// A kernel was already chosen, or a ceiling already set: the tiers in
177    /// use cannot change for the rest of the process.
178    AlreadyResolved,
179    /// The level belongs to another architecture than this build's.
180    OtherArchitecture(CpuLevel),
181    /// This target has no atomics to hold a ceiling in.
182    Unsupported,
183}
184
185impl core::fmt::Display for CpuCeilingError {
186    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
187        match self {
188            CpuCeilingError::AlreadyResolved => {
189                f.write_str("the CPU kernels were already chosen; set the ceiling before any work")
190            }
191            CpuCeilingError::OtherArchitecture(level) => {
192                write!(f, "{level} is not an instruction set of this architecture")
193            }
194            CpuCeilingError::Unsupported => f.write_str("this target cannot hold a CPU ceiling"),
195        }
196    }
197}
198
199#[cfg(feature = "std")]
200impl std::error::Error for CpuCeilingError {}
201
202/// The ceiling, frozen at its first read: 0 is not yet read or set,
203/// [`CEILING_NONE`] is frozen without a limit, anything else is a level's
204/// index in [`CpuLevel::ALL`] plus one.
205#[cfg(target_has_atomic = "8")]
206static CEILING: core::sync::atomic::AtomicU8 = core::sync::atomic::AtomicU8::new(0);
207#[cfg(target_has_atomic = "8")]
208const CEILING_NONE: u8 = u8::MAX;
209
210/// Limit every kernel the codec picks at run time to `level` and below, for
211/// the rest of the process. A kernel is still used only where the CPU
212/// supports it: the ceiling never raises a choice.
213///
214/// Only the run-time choices are limited. A build compiled for a wider
215/// baseline (`-C target-cpu=...`) uses those instructions throughout, below
216/// any ceiling.
217///
218/// # Errors
219///
220/// [`CpuCeilingError::AlreadyResolved`] once any kernel has been chosen or a
221/// ceiling set. Building a codec context can already choose kernels (a
222/// decoder does, at construction), and so does querying one
223/// ([`active_cpu_kernel_name`],
224/// [`cpu_ceiling`]): set the ceiling before creating any context, not merely
225/// before the first compression or decompression.
226/// [`CpuCeilingError::OtherArchitecture`] for a level of another architecture.
227///
228/// # Examples
229/// ```standalone_crate
230/// use structured_zstd::{CpuLevel, set_cpu_ceiling};
231///
232/// // Before any compression or decompression in the process.
233/// set_cpu_ceiling(CpuLevel::Scalar).unwrap();
234/// assert_eq!(structured_zstd::cpu_ceiling(), Some(CpuLevel::Scalar));
235/// ```
236pub fn set_cpu_ceiling(level: CpuLevel) -> Result<(), CpuCeilingError> {
237    if !level.native() {
238        return Err(CpuCeilingError::OtherArchitecture(level));
239    }
240    #[cfg(target_has_atomic = "8")]
241    {
242        use core::sync::atomic::Ordering;
243        let index = CpuLevel::ALL
244            .iter()
245            .position(|candidate| *candidate == level)
246            .expect("every level is in ALL") as u8;
247        CEILING
248            .compare_exchange(0, index + 1, Ordering::AcqRel, Ordering::Acquire)
249            .map(|_| ())
250            .map_err(|_| CpuCeilingError::AlreadyResolved)
251    }
252    #[cfg(not(target_has_atomic = "8"))]
253    {
254        Err(CpuCeilingError::Unsupported)
255    }
256}
257
258/// The ceiling in force, freezing it (as none) if nothing set one yet.
259pub fn cpu_ceiling() -> Option<CpuLevel> {
260    #[cfg(target_has_atomic = "8")]
261    {
262        use core::sync::atomic::Ordering;
263        // Once frozen the value never changes, so a plain load answers every
264        // later call; only the first read pays the locked compare-exchange.
265        let mut stored = CEILING.load(Ordering::Acquire);
266        if stored == 0 {
267            stored = match CEILING.compare_exchange(
268                0,
269                CEILING_NONE,
270                Ordering::AcqRel,
271                Ordering::Acquire,
272            ) {
273                Ok(_) => CEILING_NONE,
274                Err(stored) => stored,
275            };
276        }
277        (stored != CEILING_NONE).then(|| CpuLevel::ALL[usize::from(stored - 1)])
278    }
279    #[cfg(not(target_has_atomic = "8"))]
280    {
281        None
282    }
283}
284
285/// Whether a kernel needing `needed` may run under `ceiling`.
286// Only a build that compiles a run-time kernel choice asks; wasm, or a target
287// with its kernel features off, compiles none.
288#[cfg_attr(
289    not(any(
290        test,
291        target_arch = "x86_64",
292        all(target_arch = "x86", feature = "kernel-sse"),
293        all(target_arch = "aarch64", feature = "kernel-neon"),
294    )),
295    allow(dead_code)
296)]
297const fn allowed_under(ceiling: Option<CpuLevel>, needed: CpuLevel) -> bool {
298    let Some(ceiling) = ceiling else {
299        return true;
300    };
301    let (needed_ladder, needed_rung) = needed.ladder();
302    let (ceiling_ladder, ceiling_rung) = ceiling.ladder();
303    needed_rung == 0 || (needed_ladder as u8 == ceiling_ladder as u8 && needed_rung <= ceiling_rung)
304}
305
306/// Whether a kernel needing `needed` may be chosen. Every run-time kernel
307/// selection asks this beside its CPU-feature probe.
308#[cfg_attr(
309    not(any(
310        target_arch = "x86_64",
311        all(target_arch = "x86", feature = "kernel-sse"),
312        all(target_arch = "aarch64", feature = "kernel-neon"),
313    )),
314    allow(dead_code)
315)]
316#[inline]
317pub(crate) fn cpu_allows(needed: CpuLevel) -> bool {
318    allowed_under(cpu_ceiling(), needed)
319}
320
321/// Trait covering the leaf hot-path operations whose bodies differ
322/// per ISA. Implementations are ZSTs; the trait is `Copy` so it can
323/// be `Default`-constructed at each call site without runtime cost.
324///
325/// New methods land here ONLY when their codegen genuinely differs
326/// per kernel (BMI2 intrinsic vs scalar shift, AVX2 256-bit move vs
327/// SSE2 128-bit move, etc.). Structure ops that have one canonical
328/// implementation must NOT be on this trait — they stay on the
329/// existing decoder / encoder types.
330// Public (rather than `pub(crate)`) because `BitReaderReversed` is
331// generic over `K: CpuKernel = ScalarKernel` and is re-exported via
332// the `bench-internals`-gated `testing` module; under that feature
333// the visibility of every type that appears in `BitReaderReversed`'s
334// bounds (the trait + the default kernel) must match the type's own
335// visibility, otherwise rustc rejects with `private_bounds` /
336// `private_interfaces`. The trait surface stays narrow on stable
337// crate users: nothing outside `bench-internals` constructs a
338// non-Scalar kernel directly.
339pub trait CpuKernel: Copy + 'static {
340    /// Mask the low `n` bits of `value`, returning the remaining
341    /// high bits zeroed. The FSE bitstream hot path fires this 3×
342    /// per decoded sequence; on BMI2-capable hardware this maps to
343    /// a single `_bzhi_u64` instruction, otherwise to a scalar
344    /// `u64::MAX >> (64 - n)` shift + mask.
345    ///
346    /// Precondition: `n <= 64`. Behaviour for `n == 0` is "return 0";
347    /// behaviour for `n > 64` is unspecified — callers MUST uphold
348    /// the bound. The test-only `mask_lower_bits` helper in
349    /// `bit_reader_reverse.rs` debug-asserts the bound for its
350    /// unit tests, but production callers (FSE / HUF hot paths)
351    /// derive `n` from `accuracy_log` / `max_num_bits` which the
352    /// per-stream table builders pin to `n <= MAX_*_BITS` at
353    /// construction time; no per-call wrapper assert runs.
354    fn mask_lower_bits(value: u64, n: u8) -> u64;
355
356    /// Split the low `n1 + n2 + n3` bits of `packed` into three fields, the
357    /// highest first. The FSE sequence decoder reads its three state updates
358    /// this way, once per sequence.
359    ///
360    /// The default is three [`Self::mask_lower_bits`]; a kernel whose hardware
361    /// extracts them in one instruction overrides it. Every implementation
362    /// returns the same three values, so which one ran is invisible to the
363    /// stream being decoded.
364    ///
365    /// Precondition: `n1 + n2 + n3 <= 64`, as for `mask_lower_bits`.
366    #[inline(always)]
367    fn extract_triple(packed: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) {
368        (
369            Self::mask_lower_bits(packed.wrapping_shr(u32::from(n3) + u32::from(n2)), n1),
370            Self::mask_lower_bits(packed.wrapping_shr(u32::from(n3)), n2),
371            Self::mask_lower_bits(packed, n3),
372        )
373    }
374
375    /// Width in bytes of [`Self::copy_chunks`]'s stores. The default is
376    /// portable code, a machine word (`simd128` on wasm, which has no run-time
377    /// tier); each SIMD tier overrides it with its own vector.
378    const COPY_CHUNK: usize = crate::decoding::simd_copy::PORTABLE_COPY_CHUNK;
379
380    /// Copy `len` bytes, a multiple of [`Self::COPY_CHUNK`], in whole chunks:
381    /// the wildcopy body of the decoder's buffer copies.
382    ///
383    /// # Safety
384    /// `src` readable and `dst` writable for `len` bytes; the regions do not
385    /// overlap; the running CPU supports this kernel's tier.
386    #[inline(always)]
387    unsafe fn copy_chunks(src: *const u8, dst: *mut u8, len: usize) {
388        unsafe { crate::decoding::simd_copy::copy_chunks_portable(src, dst, len) }
389    }
390
391    /// Width of [`Self::copy_step`], the chunk a wide tier steps down to when
392    /// the slack left does not fit its own; equal to [`Self::COPY_CHUNK`]
393    /// where there is nothing narrower to step to.
394    const STEP_CHUNK: usize = Self::COPY_CHUNK;
395
396    /// Copy `len` bytes, a multiple of [`Self::STEP_CHUNK`], in whole chunks.
397    ///
398    /// # Safety
399    /// As [`Self::copy_chunks`].
400    #[inline(always)]
401    unsafe fn copy_step(src: *const u8, dst: *mut u8, len: usize) {
402        unsafe { Self::copy_chunks(src, dst, len) }
403    }
404
405    /// Copy exactly 16 bytes in one transfer, for a short copy whose buffers
406    /// have the room to overshoot to 16.
407    ///
408    /// # Safety
409    /// `src` readable and `dst` writable for 16 bytes; the regions do not
410    /// overlap; the running CPU supports this kernel's tier.
411    #[inline(always)]
412    unsafe fn copy16(src: *const u8, dst: *mut u8) {
413        unsafe { crate::decoding::simd_copy::copy16_portable(src, dst) }
414    }
415}
416
417/// Copies at the build's baseline vector (SSE2 where the target guarantees it,
418/// NEON, `simd128`) and masks as the scalar tier does: for the copies that run
419/// outside any tier's dispatch, a raw block's or a literals-only block's, which
420/// every CPU the build runs on can take at that width.
421#[derive(Copy, Clone, Default)]
422pub(crate) struct BaselineKernel;
423
424impl CpuKernel for BaselineKernel {
425    #[inline(always)]
426    fn mask_lower_bits(value: u64, n: u8) -> u64 {
427        ScalarKernel::mask_lower_bits(value, n)
428    }
429
430    const COPY_CHUNK: usize = crate::decoding::simd_copy::BASELINE_COPY_CHUNK;
431
432    #[inline(always)]
433    unsafe fn copy_chunks(src: *const u8, dst: *mut u8, len: usize) {
434        // SAFETY: the baseline runs on every CPU the build does.
435        unsafe { crate::decoding::simd_copy::copy_chunks_baseline(src, dst, len) }
436    }
437
438    #[inline(always)]
439    unsafe fn copy16(src: *const u8, dst: *mut u8) {
440        // SAFETY: as above.
441        unsafe { crate::decoding::simd_copy::copy16_baseline(src, dst) }
442    }
443}
444
445/// The SSE2 copies every x86 SIMD tier shares: 16-byte chunks. A macro, so
446/// each kernel's methods expand them in place.
447#[cfg(all(
448    any(target_arch = "x86", target_arch = "x86_64"),
449    feature = "kernel-sse"
450))]
451macro_rules! sse2_copies {
452    () => {
453        const COPY_CHUNK: usize = 16;
454
455        #[inline(always)]
456        unsafe fn copy_chunks(src: *const u8, dst: *mut u8, len: usize) {
457            // SAFETY: every tier using this was selected with SSE2 confirmed.
458            unsafe { crate::decoding::simd_copy::copy_sse2(src, dst, len) }
459        }
460
461        #[inline(always)]
462        unsafe fn copy16(src: *const u8, dst: *mut u8) {
463            // SAFETY: as above.
464            unsafe { crate::decoding::simd_copy::copy_sse2(src, dst, 16) }
465        }
466    };
467}
468
469/// x86 SSE2 kernel: scalar masks (SSE2 has no bit extract) and 16-byte SSE2
470/// copies. Chosen at run time, so a 32-bit build whose baseline lacks SSE2
471/// still copies in vectors on a CPU that has it.
472#[cfg(all(
473    any(target_arch = "x86", target_arch = "x86_64"),
474    feature = "kernel-sse"
475))]
476#[derive(Copy, Clone, Default)]
477pub(crate) struct Sse2Kernel;
478
479#[cfg(all(
480    any(target_arch = "x86", target_arch = "x86_64"),
481    feature = "kernel-sse"
482))]
483impl CpuKernel for Sse2Kernel {
484    #[inline(always)]
485    fn mask_lower_bits(value: u64, n: u8) -> u64 {
486        ScalarKernel::mask_lower_bits(value, n)
487    }
488
489    sse2_copies!();
490}
491
492/// Scalar fallback — portable, no SIMD or BMI2 intrinsics. Selected
493/// when no x86 or aarch64 feature is detected at runtime.
494#[derive(Copy, Clone, Default)]
495pub struct ScalarKernel;
496
497/// `BIT_MASK[n]` is the low `n` bits set for `n` in `0..=64`, and all bits for
498/// anything past that (a width the formats cannot ask for).
499///
500/// A table rather than `u64::MAX >> (64 - n)`: the shift form needs a guard for
501/// `n == 0`, since a 64-bit shift is undefined, and that guard is a branch or a
502/// `cmov` on every field of every sequence. Indexed by a `u8`, and sized for
503/// every `u8`, so the load carries no bounds check either. The widths in use
504/// are small, so the hot part is the first few cache lines of it.
505///
506/// A `static`, as upstream zstd's `static const unsigned BIT_mask[]`
507/// (`lib/common/bitstream.h`) is one object. A `const` array is a value
508/// copied at each use; a debug (`-O0`) build copied all 2 KiB of it onto the
509/// stack at every inlined load, 43 KiB of the scalar sequence decoder's frame.
510pub(crate) static BIT_MASK: [u64; 256] = {
511    let mut table = [u64::MAX; 256];
512    let mut i: usize = 0;
513    while i < 64 {
514        table[i] = (1u64 << i) - 1;
515        i += 1;
516    }
517    table
518};
519
520impl CpuKernel for ScalarKernel {
521    #[inline(always)]
522    fn mask_lower_bits(value: u64, n: u8) -> u64 {
523        value & BIT_MASK[n as usize]
524    }
525}
526
527/// BMI2-only kernel: `bzhi` for mask_lower_bits. Selected when the CPU has
528/// BMI2 but not the AVX2 SIMD width to upgrade to the Avx2 kernel. Treated as
529/// a stepping stone between Sse2 and Avx2 on hardware that has BMI2 but not
530/// AVX2 (rare in practice but matches upstream zstd's gating). Present on
531/// 32-bit x86 as well as x86_64: the instruction is there, only its width
532/// differs, and without this tier a 32-bit build would decode on the scalar
533/// bodies whatever the CPU offers.
534#[cfg(all(
535    any(target_arch = "x86", target_arch = "x86_64"),
536    feature = "kernel-bmi2"
537))]
538#[derive(Copy, Clone, Default)]
539pub(crate) struct Bmi2Kernel;
540
541#[cfg(all(
542    any(target_arch = "x86", target_arch = "x86_64"),
543    feature = "kernel-bmi2"
544))]
545impl CpuKernel for Bmi2Kernel {
546    #[inline(always)]
547    fn mask_lower_bits(value: u64, n: u8) -> u64 {
548        // SAFETY: this kernel ZST is only reachable via the
549        // `match detect_cpu_kernel() { CpuKernelTag::Bmi2 => ... }`
550        // dispatch arms at decoder entry sites, all of which fire only
551        // after `detect_cpu_kernel` confirmed BMI2 is available on the
552        // running CPU.
553        unsafe { mask_lower_bits_bmi2_impl(value, n) }
554    }
555
556    // Every BMI2 CPU has SSE2.
557    sse2_copies!();
558}
559
560/// x86 AVX2 + BMI2 kernel (x86-64-v3 baseline). The common modern
561/// x86 case — most CPUs released since 2013 (Haswell) have AVX2+BMI2.
562/// Uses `bzhi` for mask ops and 256-bit moves for the buffer copies. Present
563/// on 32-bit x86 too, where the kernel's masks are the table ones and its
564/// copies are what it adds; the tier still decodes literals on the BMI2 path,
565/// so, like every tier, it requires the ones below it.
566#[cfg(all(
567    any(target_arch = "x86", target_arch = "x86_64"),
568    feature = "kernel-avx2"
569))]
570#[derive(Copy, Clone, Default)]
571pub(crate) struct Avx2Kernel;
572
573#[cfg(all(
574    any(target_arch = "x86", target_arch = "x86_64"),
575    feature = "kernel-avx2"
576))]
577impl CpuKernel for Avx2Kernel {
578    #[inline(always)]
579    fn mask_lower_bits(value: u64, n: u8) -> u64 {
580        // SAFETY: Avx2Kernel is selected only after runtime detect
581        // confirmed both AVX2 and BMI2 — `_bzhi_u64` is callable.
582        #[cfg(target_arch = "x86_64")]
583        unsafe {
584            mask_lower_bits_bmi2_impl(value, n)
585        }
586        // 32-bit x86 has only a 32-bit `bzhi`: two of them and the branches
587        // choosing between them cost more than the table lookup in the
588        // sequence loop, which calls this three times a sequence (decoding
589        // z000033 on i686 measured 12.6% slower at level 1 with them).
590        #[cfg(target_arch = "x86")]
591        {
592            ScalarKernel::mask_lower_bits(value, n)
593        }
594    }
595
596    // 32 bytes, the width the buffers' trailing slack is sized for.
597    const COPY_CHUNK: usize = 32;
598
599    #[inline(always)]
600    unsafe fn copy_chunks(src: *const u8, dst: *mut u8, len: usize) {
601        // SAFETY: this kernel is selected only after detect confirmed AVX2.
602        unsafe { crate::decoding::simd_copy::copy_avx2(src, dst, len) }
603    }
604
605    const STEP_CHUNK: usize = 16;
606
607    #[inline(always)]
608    unsafe fn copy_step(src: *const u8, dst: *mut u8, len: usize) {
609        // SAFETY: an AVX2 CPU has SSE2.
610        unsafe { crate::decoding::simd_copy::copy_sse2(src, dst, len) }
611    }
612
613    #[inline(always)]
614    unsafe fn copy16(src: *const u8, dst: *mut u8) {
615        // SAFETY: as above.
616        unsafe { crate::decoding::simd_copy::copy_sse2(src, dst, 16) }
617    }
618}
619
620/// x86_64 AVX-512 VBMI2 + AVX2 + BMI2 kernel. Selected when the CPU
621/// has the AVX-512 VBMI2 family available — VBMI2 unlocks a faster
622/// HUF burst inner loop (VPSHUFB-based table lookup); BMI2 mask_lower
623/// bits stays identical to Avx2 kernel.
624#[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))]
625#[derive(Copy, Clone, Default)]
626pub(crate) struct Vbmi2Kernel;
627
628#[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))]
629impl CpuKernel for Vbmi2Kernel {
630    #[inline(always)]
631    fn mask_lower_bits(value: u64, n: u8) -> u64 {
632        // SAFETY: same precondition as Avx2Kernel — BMI2 confirmed
633        // at runtime before this kernel is instantiated.
634        unsafe { mask_lower_bits_bmi2_impl(value, n) }
635    }
636
637    // The AVX2 width: a 64-byte store would need twice the buffers' trailing
638    // slack to fire on the short copies that dominate.
639    const COPY_CHUNK: usize = 32;
640
641    #[inline(always)]
642    unsafe fn copy_chunks(src: *const u8, dst: *mut u8, len: usize) {
643        // SAFETY: the VBMI2 tier is selected only with AVX2 confirmed.
644        unsafe { crate::decoding::simd_copy::copy_avx2(src, dst, len) }
645    }
646
647    const STEP_CHUNK: usize = 16;
648
649    #[inline(always)]
650    unsafe fn copy_step(src: *const u8, dst: *mut u8, len: usize) {
651        // SAFETY: an AVX2 CPU has SSE2.
652        unsafe { crate::decoding::simd_copy::copy_sse2(src, dst, len) }
653    }
654
655    #[inline(always)]
656    unsafe fn copy16(src: *const u8, dst: *mut u8) {
657        // SAFETY: as above.
658        unsafe { crate::decoding::simd_copy::copy_sse2(src, dst, 16) }
659    }
660}
661
662/// aarch64 NEON kernel: scalar masks and 16-byte NEON copies. Used on all
663/// aarch64 hardware that exposes NEON (effectively universal on the
664/// supported targets).
665#[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))]
666#[derive(Copy, Clone, Default)]
667pub(crate) struct NeonKernel;
668
669#[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))]
670impl CpuKernel for NeonKernel {
671    #[inline(always)]
672    fn mask_lower_bits(value: u64, n: u8) -> u64 {
673        // aarch64 has no BMI2 equivalent that improves on the scalar
674        // shift-and-mask sequence for this op; the codegen is
675        // identical to the Scalar kernel here.
676        ScalarKernel::mask_lower_bits(value, n)
677    }
678
679    const COPY_CHUNK: usize = BaselineKernel::COPY_CHUNK;
680
681    #[inline(always)]
682    unsafe fn copy_chunks(src: *const u8, dst: *mut u8, len: usize) {
683        // SAFETY: NEON is the aarch64 baseline.
684        unsafe { BaselineKernel::copy_chunks(src, dst, len) }
685    }
686
687    #[inline(always)]
688    unsafe fn copy16(src: *const u8, dst: *mut u8) {
689        // SAFETY: as above.
690        unsafe { BaselineKernel::copy16(src, dst) }
691    }
692}
693
694/// aarch64 SVE kernel. Variable-vector-length SVE extends NEON for
695/// HUF burst / SIMD copy on Graviton3 / Apple M-series with SVE
696/// support. Masks and copies as the NEON kernel's.
697#[cfg(all(target_arch = "aarch64", feature = "kernel-sve"))]
698#[allow(dead_code)]
699#[derive(Copy, Clone, Default)]
700pub(crate) struct SveKernel;
701
702#[cfg(all(target_arch = "aarch64", feature = "kernel-sve"))]
703impl CpuKernel for SveKernel {
704    #[inline(always)]
705    fn mask_lower_bits(value: u64, n: u8) -> u64 {
706        ScalarKernel::mask_lower_bits(value, n)
707    }
708
709    const COPY_CHUNK: usize = BaselineKernel::COPY_CHUNK;
710
711    #[inline(always)]
712    unsafe fn copy_chunks(src: *const u8, dst: *mut u8, len: usize) {
713        // SAFETY: NEON is the aarch64 baseline.
714        unsafe { BaselineKernel::copy_chunks(src, dst, len) }
715    }
716
717    #[inline(always)]
718    unsafe fn copy16(src: *const u8, dst: *mut u8) {
719        // SAFETY: as above.
720        unsafe { BaselineKernel::copy16(src, dst) }
721    }
722}
723
724/// Single `#[target_feature(enable = "bmi2")]` wrapper around the
725/// `_bzhi_u64` intrinsic. Lifted to a free function so each kernel
726/// impl that needs the BMI2 path (Bmi2 / Avx2 / Vbmi2) calls the
727/// same shared body. With `#[inline]` LLVM inlines the call into
728/// any caller that itself has BMI2 in scope; outside that scope the
729/// target_feature boundary is preserved.
730#[cfg(all(
731    any(target_arch = "x86", target_arch = "x86_64"),
732    feature = "kernel-bmi2"
733))]
734#[target_feature(enable = "bmi2")]
735#[inline]
736unsafe fn mask_lower_bits_bmi2_impl(value: u64, n: u8) -> u64 {
737    // The intrinsic call is permitted directly inside a function
738    // already annotated `#[target_feature(enable = "bmi2")]` — no
739    // `unsafe { ... }` block needed (the function-level `unsafe`
740    // already covers it). SAFETY: caller selected a kernel whose
741    // CpuKernelTag was resolved after `is_x86_feature_detected!("bmi2")`
742    // returned true, so the BMI2 instruction set is available.
743    #[cfg(target_arch = "x86_64")]
744    {
745        core::arch::x86_64::_bzhi_u64(value, n as u32)
746    }
747    // 32-bit x86 has `bzhi` on 32-bit registers only. Widths up to 32 take one
748    // instruction on the low half; wider ones keep the low 32 bits whole and
749    // apply it to the high half, which is what a 64-bit `bzhi` does in one go.
750    #[cfg(target_arch = "x86")]
751    {
752        use core::arch::x86::_bzhi_u32;
753        if n >= 64 {
754            return value;
755        }
756        if n <= 32 {
757            return u64::from(_bzhi_u32(value as u32, u32::from(n)));
758        }
759        let high = _bzhi_u32((value >> 32) as u32, u32::from(n) - 32);
760        (value & u64::from(u32::MAX)) | (u64::from(high) << 32)
761    }
762}
763
764/// Pure boolean-input variant of the x86 kernel-tag selection. Both the
765/// `std` runtime-detect path and the `no_std` compile-time-cfg path
766/// route through this helper so the precedence rules stay in one place
767/// (and are unit-testable without runtime CPUID).
768///
769/// The VBMI2 tier requires every AVX-512 sub-feature it touches AND the
770/// AVX2 baseline — VBMI2 kernels mix VBMI2-only intrinsics with AVX2
771/// 256-bit moves, so the dispatch must be conditioned on `has_avx2` too.
772/// Likewise the Avx2 tier requires both AVX2 and BMI2.
773#[cfg(target_arch = "x86_64")]
774#[inline(always)]
775// Params go unused when the matching `kernel_*` feature is disabled (the
776// rung that consumes them is `#[cfg]`-ed out); they are still passed by the
777// detect callers. Silence the conditional unused-variable warning rather
778// than thread per-feature `_`-prefixes through the signature.
779#[allow(unused_variables)]
780const fn select_x86_kernel(
781    has_avx512vbmi2: bool,
782    has_avx512f: bool,
783    has_avx512vl: bool,
784    has_avx512bw: bool,
785    has_bmi2: bool,
786    has_avx2: bool,
787    has_sse2: bool,
788) -> CpuKernelTag {
789    #[cfg(feature = "kernel-vbmi2")]
790    if has_avx512vbmi2 && has_avx512f && has_avx512vl && has_avx512bw && has_bmi2 && has_avx2 {
791        return CpuKernelTag::Vbmi2;
792    }
793    #[cfg(feature = "kernel-avx2")]
794    if has_avx2 && has_bmi2 {
795        return CpuKernelTag::Avx2;
796    }
797    #[cfg(feature = "kernel-bmi2")]
798    if has_bmi2 {
799        return CpuKernelTag::Bmi2;
800    }
801    #[cfg(feature = "kernel-sse")]
802    if has_sse2 {
803        return CpuKernelTag::Sse2;
804    }
805    CpuKernelTag::Scalar
806}
807
808/// Cached runtime-detected kernel tag. The actual `CpuKernel` impl
809/// (`ScalarKernel` / `Bmi2Kernel` / `Avx2Kernel` / `Vbmi2Kernel` /
810/// `NeonKernel` / `SveKernel`) is constructed at the dispatch site —
811/// currently only `decoding::literals_section_decoder::decompress_literals`
812/// — via a `match` on this tag that branches into the per-K
813/// `target_feature`-wrapped specialisation. Pipeline-wide dispatch
814/// (FrameDecoder / FrameCompressor entry, sequence executor, match
815/// copy) lands incrementally in follow-up tiers.
816#[derive(Copy, Clone, Debug, Eq, PartialEq)]
817pub(crate) enum CpuKernelTag {
818    Scalar,
819    /// Chosen at run time on 32-bit x86 too, where a build's baseline may
820    /// lack SSE2 while the CPU has it.
821    #[cfg(all(
822        any(target_arch = "x86", target_arch = "x86_64"),
823        feature = "kernel-sse"
824    ))]
825    Sse2,
826    /// Reachable on 32-bit x86 as well: `bzhi` is there, and without the tier
827    /// such a build would decode on the scalar bodies whatever the CPU offers.
828    #[cfg(all(
829        any(target_arch = "x86", target_arch = "x86_64"),
830        feature = "kernel-bmi2"
831    ))]
832    Bmi2,
833    /// On 32-bit x86 it runs the portable sequence walk with this kernel's
834    /// 32-byte copies.
835    #[cfg(all(
836        any(target_arch = "x86", target_arch = "x86_64"),
837        feature = "kernel-avx2"
838    ))]
839    Avx2,
840    #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))]
841    Vbmi2,
842    #[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))]
843    Neon,
844    // Both constructors of `Sve` need a reachable feature: runtime
845    // detection via `std::arch::is_aarch64_feature_detected!` (so
846    // `feature = "std"`) or compile-time `target_feature = "sve"` in
847    // RUSTFLAGS. Without either, the variant is unreachable and a
848    // `match` arm referencing it warns as dead.
849    #[cfg(all(
850        target_arch = "aarch64",
851        feature = "kernel-sve",
852        any(feature = "std", target_feature = "sve"),
853    ))]
854    Sve,
855}
856
857/// Detect once and cache the best available CPU kernel for the
858/// current process. Subsequent calls return the cached tag without
859/// re-running CPU-feature detection. Std-only — no-std targets use
860/// the compile-time variant below that resolves at build time.
861#[cfg(feature = "std")]
862pub(crate) fn detect_cpu_kernel() -> CpuKernelTag {
863    static CACHED: OnceLock<CpuKernelTag> = OnceLock::new();
864    *CACHED.get_or_init(detect_cpu_kernel_uncached)
865}
866
867#[cfg(feature = "std")]
868fn detect_cpu_kernel_uncached() -> CpuKernelTag {
869    #[cfg(target_arch = "x86_64")]
870    {
871        use std::arch::is_x86_feature_detected;
872        // Gate each probe on its tier feature: `cfg!(...)` const-folds, so the
873        // `&&` short-circuits away the runtime `is_x86_feature_detected!` call
874        // (and its CPUID/cache traffic) for tiers the build disabled — the
875        // matching `select_x86_kernel` rung is `#[cfg]`-ed out anyway.
876        let avx512 = cfg!(feature = "kernel-vbmi2") && cpu_allows(CpuLevel::Avx512);
877        return select_x86_kernel(
878            avx512 && is_x86_feature_detected!("avx512vbmi2"),
879            avx512 && is_x86_feature_detected!("avx512f"),
880            avx512 && is_x86_feature_detected!("avx512vl"),
881            avx512 && is_x86_feature_detected!("avx512bw"),
882            cfg!(feature = "kernel-bmi2")
883                && cpu_allows(CpuLevel::Bmi2)
884                && is_x86_feature_detected!("bmi2"),
885            cfg!(feature = "kernel-avx2")
886                && cpu_allows(CpuLevel::Avx2)
887                && is_x86_feature_detected!("avx2"),
888            cfg!(feature = "kernel-sse")
889                && cpu_allows(CpuLevel::Sse2)
890                && is_x86_feature_detected!("sse2"),
891        );
892    }
893    // 32-bit x86 carries the SSE2, BMI2 and AVX2 tiers: SSE2 copies, `bzhi`
894    // for the entropy tables and the AVX2 buffer copies. The VBMI2 bodies are
895    // x86_64-only.
896    #[cfg(target_arch = "x86")]
897    {
898        #[cfg(feature = "kernel-sse")]
899        {
900            use std::arch::is_x86_feature_detected;
901            #[cfg(feature = "kernel-bmi2")]
902            {
903                let bmi2 = cpu_allows(CpuLevel::Bmi2) && is_x86_feature_detected!("bmi2");
904                // The AVX2 tier sits above BMI2, as on x86_64: its literals take
905                // the BMI2 body, so an AVX2 CPU without BMI2 (in practice a
906                // virtual machine that hides it) runs the SSE2 tier. Its copies
907                // stay 16 bytes wide; a tier of AVX2 copies without BMI2 would be
908                // one more sequence decoder for a CPU no vendor ships.
909                #[cfg(feature = "kernel-avx2")]
910                if bmi2 && cpu_allows(CpuLevel::Avx2) && is_x86_feature_detected!("avx2") {
911                    return CpuKernelTag::Avx2;
912                }
913                if bmi2 {
914                    return CpuKernelTag::Bmi2;
915                }
916            }
917            if cpu_allows(CpuLevel::Sse2) && is_x86_feature_detected!("sse2") {
918                return CpuKernelTag::Sse2;
919            }
920        }
921        return CpuKernelTag::Scalar;
922    }
923    #[cfg(target_arch = "aarch64")]
924    {
925        #[cfg(any(feature = "kernel-sve", feature = "kernel-neon"))]
926        use std::arch::is_aarch64_feature_detected;
927        #[cfg(feature = "kernel-sve")]
928        if cpu_allows(CpuLevel::Sve) && is_aarch64_feature_detected!("sve") {
929            return CpuKernelTag::Sve;
930        }
931        #[cfg(feature = "kernel-neon")]
932        if cpu_allows(CpuLevel::Neon) && is_aarch64_feature_detected!("neon") {
933            return CpuKernelTag::Neon;
934        }
935        return CpuKernelTag::Scalar;
936    }
937    #[allow(unreachable_code)]
938    CpuKernelTag::Scalar
939}
940
941/// no-std variant: rely on compile-time `target_feature` flags
942/// instead of runtime detection. Resolves to the most-capable kernel
943/// that the build target supports.
944#[cfg(not(feature = "std"))]
945pub(crate) fn detect_cpu_kernel() -> CpuKernelTag {
946    #[cfg(target_arch = "x86_64")]
947    {
948        // Route through the same const-fn precedence helper as the
949        // `feature = "std"` path. `cfg!(target_feature = ...)`
950        // returns a compile-time bool that constant-folds through
951        // `select_x86_kernel`, so the runtime call has the same
952        // codegen as the previous hand-written #[cfg] chain.
953        let avx512 = cpu_allows(CpuLevel::Avx512);
954        return select_x86_kernel(
955            avx512 && cfg!(target_feature = "avx512vbmi2"),
956            avx512 && cfg!(target_feature = "avx512f"),
957            avx512 && cfg!(target_feature = "avx512vl"),
958            avx512 && cfg!(target_feature = "avx512bw"),
959            cpu_allows(CpuLevel::Bmi2) && cfg!(target_feature = "bmi2"),
960            cpu_allows(CpuLevel::Avx2) && cfg!(target_feature = "avx2"),
961            cpu_allows(CpuLevel::Sse2) && cfg!(target_feature = "sse2"),
962        );
963    }
964    // `cfg!` rather than `#[cfg]`, as `select_x86_kernel` does: the tests fold
965    // at compile time, and the tiers stay constructed in every build.
966    #[cfg(all(target_arch = "x86", feature = "kernel-bmi2"))]
967    {
968        let bmi2 = cpu_allows(CpuLevel::Bmi2) && cfg!(target_feature = "bmi2");
969        #[cfg(feature = "kernel-avx2")]
970        if bmi2 && cpu_allows(CpuLevel::Avx2) && cfg!(target_feature = "avx2") {
971            return CpuKernelTag::Avx2;
972        }
973        if bmi2 {
974            return CpuKernelTag::Bmi2;
975        }
976    }
977    #[cfg(all(target_arch = "x86", feature = "kernel-sse"))]
978    if cpu_allows(CpuLevel::Sse2) && cfg!(target_feature = "sse2") {
979        return CpuKernelTag::Sse2;
980    }
981    #[cfg(target_arch = "aarch64")]
982    {
983        #[cfg(all(feature = "kernel-sve", target_feature = "sve"))]
984        if cpu_allows(CpuLevel::Sve) {
985            return CpuKernelTag::Sve;
986        }
987        #[cfg(all(feature = "kernel-neon", target_feature = "neon"))]
988        if cpu_allows(CpuLevel::Neon) {
989            return CpuKernelTag::Neon;
990        }
991    }
992    #[allow(unreachable_code)]
993    CpuKernelTag::Scalar
994}
995
996impl CpuKernelTag {
997    /// Stable lowercase diagnostic name for this tier (used by
998    /// [`active_cpu_kernel_name`] and the bench/dashboard reporting). Pure
999    /// mapping over the tag, so every arm is exercisable in tests regardless
1000    /// of which tier the running CPU actually resolves to.
1001    pub(crate) fn name(self) -> &'static str {
1002        match self {
1003            CpuKernelTag::Scalar => "scalar",
1004            #[cfg(all(
1005                any(target_arch = "x86", target_arch = "x86_64"),
1006                feature = "kernel-sse"
1007            ))]
1008            CpuKernelTag::Sse2 => "sse2",
1009            #[cfg(all(
1010                any(target_arch = "x86", target_arch = "x86_64"),
1011                feature = "kernel-bmi2"
1012            ))]
1013            CpuKernelTag::Bmi2 => "bmi2",
1014            #[cfg(all(
1015                any(target_arch = "x86", target_arch = "x86_64"),
1016                feature = "kernel-avx2"
1017            ))]
1018            CpuKernelTag::Avx2 => "avx2",
1019            #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))]
1020            CpuKernelTag::Vbmi2 => "vbmi2",
1021            #[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))]
1022            CpuKernelTag::Neon => "neon",
1023            #[cfg(all(
1024                target_arch = "aarch64",
1025                feature = "kernel-sve",
1026                any(feature = "std", target_feature = "sve"),
1027            ))]
1028            CpuKernelTag::Sve => "sve",
1029        }
1030    }
1031}
1032
1033/// Name of the CPU kernel tier this process selected for the entropy /
1034/// sequence hot paths: decode (literals + FSE sequence decode) and encode
1035/// (entropy) share this dispatch (see #247). Returned as a stable lowercase
1036/// string for diagnostics and benchmark/dashboard reporting; the value is
1037/// what the runtime CPU-feature detection (or compile-time `target_feature`
1038/// on `no_std`) actually resolves to on this machine, so a dashboard can
1039/// attribute a measurement to the kernel that produced it.
1040pub fn active_cpu_kernel_name() -> &'static str {
1041    detect_cpu_kernel().name()
1042}
1043
1044#[cfg(test)]
1045mod tests;