structured_zstd/

cpu_kernel.rs

//! CPU kernel dispatch — single detect+match at the dispatch site,
//! propagated through the inner pipeline as a generic parameter so
//! leaf hot-path code monomorphises against the chosen kernel.
//!
//! See issue #247 for the architecture rationale: per-subsystem
//! dispatch scatters the choice across HUF / FSE / SIMD-copy
//! independently and pays the cost N times per call. Lifting the
//! dispatch to the outermost feasible call site collapses it to one
//! detect there; the inner leaf-hot-path ops then route through
//! `K::method` calls on the chosen kernel zero-sized type.
//!
//! Current wiring (as of #247 Part 2): the only active dispatch site
//! is `decoding::literals_section_decoder::decompress_literals`,
//! which `match`es `detect_cpu_kernel()` and routes into per-K
//! `decompress_literals_*` `#[target_feature]` wrappers. The full
//! pipeline-wide propagation envisioned in the issue (FrameDecoder /
//! FrameCompressor entry, sequence executor, match copy) is
//! incremental; subsequent tiers extend the dispatch surface without
//! changing this trait or the kernel ZSTs.
//!
//! Structure code (block loop, FCS check, offset history, repeat
//! semantics) stays single-impl and only carries `K` as a phantom on
//! the outer function. Monomorphisation specialises ONLY the bodies
//! that actually differ per ISA — `mask_lower_bits`, `huf_burst`,
//! `copy_chunk`, etc.

#[cfg(feature = "std")]
use std::sync::OnceLock;

/// Trait covering the leaf hot-path operations whose bodies differ
/// per ISA. Implementations are ZSTs; the trait is `Copy` so it can
/// be `Default`-constructed at each call site without runtime cost.
///
/// New methods land here ONLY when their codegen genuinely differs
/// per kernel (BMI2 intrinsic vs scalar shift, AVX2 256-bit move vs
/// SSE2 128-bit move, etc.). Structure ops that have one canonical
/// implementation must NOT be on this trait — they stay on the
/// existing decoder / encoder types.
// Public (rather than `pub(crate)`) because `BitReaderReversed` is
// generic over `K: CpuKernel = ScalarKernel` and is re-exported via
// the `bench-internals`-gated `testing` module; under that feature
// the visibility of every type that appears in `BitReaderReversed`'s
// bounds (the trait + the default kernel) must match the type's own
// visibility, otherwise rustc rejects with `private_bounds` /
// `private_interfaces`. The trait surface stays narrow on stable
// crate users: nothing outside `bench-internals` constructs a
// non-Scalar kernel directly.
pub trait CpuKernel: Copy + 'static {
    /// Mask the low `n` bits of `value`, returning the remaining
    /// high bits zeroed. The FSE bitstream hot path fires this 3×
    /// per decoded sequence; on BMI2-capable hardware this maps to
    /// a single `_bzhi_u64` instruction, otherwise to a scalar
    /// `u64::MAX >> (64 - n)` shift + mask.
    ///
    /// Precondition: `n <= 64`. Behaviour for `n == 0` is "return 0";
    /// behaviour for `n > 64` is unspecified — callers MUST uphold
    /// the bound. The test-only `mask_lower_bits` helper in
    /// `bit_reader_reverse.rs` debug-asserts the bound for its
    /// unit tests, but production callers (FSE / HUF hot paths)
    /// derive `n` from `accuracy_log` / `max_num_bits` which the
    /// per-stream table builders pin to `n <= MAX_*_BITS` at
    /// construction time; no per-call wrapper assert runs.
    fn mask_lower_bits(value: u64, n: u8) -> u64;

    /// Split the low `n1 + n2 + n3` bits of `packed` into three fields, the
    /// highest first. The FSE sequence decoder reads its three state updates
    /// this way, once per sequence.
    ///
    /// The default is three [`Self::mask_lower_bits`]; a kernel whose hardware
    /// extracts them in one instruction overrides it. Every implementation
    /// returns the same three values, so which one ran is invisible to the
    /// stream being decoded.
    ///
    /// Precondition: `n1 + n2 + n3 <= 64`, as for `mask_lower_bits`.
    #[inline(always)]
    fn extract_triple(packed: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) {
        (
            Self::mask_lower_bits(packed.wrapping_shr(u32::from(n3) + u32::from(n2)), n1),
            Self::mask_lower_bits(packed.wrapping_shr(u32::from(n3)), n2),
            Self::mask_lower_bits(packed, n3),
        )
    }
}

/// Scalar fallback — portable, no SIMD or BMI2 intrinsics. Selected
/// when no x86 or aarch64 feature is detected at runtime.
#[derive(Copy, Clone, Default)]
pub struct ScalarKernel;

/// `BIT_MASK[n]` is the low `n` bits set for `n` in `0..=64`, and all bits for
/// anything past that (a width the formats cannot ask for).
///
/// A table rather than `u64::MAX >> (64 - n)`: the shift form needs a guard for
/// `n == 0`, since a 64-bit shift is undefined, and that guard is a branch or a
/// `cmov` on every field of every sequence. Indexed by a `u8`, and sized for
/// every `u8`, so the load carries no bounds check either. The widths in use
/// are small, so the hot part is the first few cache lines of it.
pub(crate) const BIT_MASK: [u64; 256] = {
    let mut table = [u64::MAX; 256];
    let mut i: usize = 0;
    while i < 64 {
        table[i] = (1u64 << i) - 1;
        i += 1;
    }
    table
};

impl CpuKernel for ScalarKernel {
    #[inline(always)]
    fn mask_lower_bits(value: u64, n: u8) -> u64 {
        value & BIT_MASK[n as usize]
    }
}

// The SSE2 tier exists in `CpuKernelTag` (it carries the 128-bit copy-chunk
// choice for the unified copy dispatch) but needs no `CpuKernel` ZST yet: the
// only trait method, `mask_lower_bits`, has no SSE2-specific form (SSE2 has no
// bit-extract), so the Sse2 tag routes through the scalar bodies for the
// FSE/HUF paths. A dedicated `Sse2Kernel` lands when `copy_chunk` moves onto
// the trait.

/// BMI2-only kernel: `bzhi` for mask_lower_bits. Selected when the CPU has
/// BMI2 but not the AVX2 SIMD width to upgrade to the Avx2 kernel. Treated as
/// a stepping stone between Sse2 and Avx2 on hardware that has BMI2 but not
/// AVX2 (rare in practice but matches upstream zstd's gating). Present on
/// 32-bit x86 as well as x86_64: the instruction is there, only its width
/// differs, and without this tier a 32-bit build would decode on the scalar
/// bodies whatever the CPU offers.
#[cfg(all(
    any(target_arch = "x86", target_arch = "x86_64"),
    feature = "kernel-bmi2"
))]
#[derive(Copy, Clone, Default)]
pub(crate) struct Bmi2Kernel;

#[cfg(all(
    any(target_arch = "x86", target_arch = "x86_64"),
    feature = "kernel-bmi2"
))]
impl CpuKernel for Bmi2Kernel {
    #[inline(always)]
    fn mask_lower_bits(value: u64, n: u8) -> u64 {
        // SAFETY: this kernel ZST is only reachable via the
        // `match detect_cpu_kernel() { CpuKernelTag::Bmi2 => ... }`
        // dispatch arms at decoder entry sites, all of which fire only
        // after `detect_cpu_kernel` confirmed BMI2 is available on the
        // running CPU.
        unsafe { mask_lower_bits_bmi2_impl(value, n) }
    }
}

/// x86_64 AVX2 + BMI2 kernel (x86-64-v3 baseline). The common modern
/// x86 case — most CPUs released since 2013 (Haswell) have AVX2+BMI2.
/// Uses `_bzhi_u64` for mask ops; future trait methods will use AVX2
/// 256-bit moves for `copy_chunk` and pext for HUF burst.
#[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))]
#[derive(Copy, Clone, Default)]
pub(crate) struct Avx2Kernel;

#[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))]
impl CpuKernel for Avx2Kernel {
    #[inline(always)]
    fn mask_lower_bits(value: u64, n: u8) -> u64 {
        // SAFETY: Avx2Kernel is selected only after runtime detect
        // confirmed both AVX2 and BMI2 — `_bzhi_u64` is callable.
        unsafe { mask_lower_bits_bmi2_impl(value, n) }
    }
}

/// x86_64 AVX-512 VBMI2 + AVX2 + BMI2 kernel. Selected when the CPU
/// has the AVX-512 VBMI2 family available — VBMI2 unlocks a faster
/// HUF burst inner loop (VPSHUFB-based table lookup); BMI2 mask_lower
/// bits stays identical to Avx2 kernel.
#[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))]
#[derive(Copy, Clone, Default)]
pub(crate) struct Vbmi2Kernel;

#[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))]
impl CpuKernel for Vbmi2Kernel {
    #[inline(always)]
    fn mask_lower_bits(value: u64, n: u8) -> u64 {
        // SAFETY: same precondition as Avx2Kernel — BMI2 confirmed
        // at runtime before this kernel is instantiated.
        unsafe { mask_lower_bits_bmi2_impl(value, n) }
    }
}

/// aarch64 NEON baseline kernel. Used on all aarch64 hardware that
/// exposes NEON (effectively universal on the supported targets).
///
/// `#[allow(dead_code)]`: scaffolding for the future aarch64 dispatch
/// arm in `decompress_literals` / `decode_and_execute_sequences`.
/// The struct + trait impl land first so the dispatch wiring can be
/// added incrementally without churning the CpuKernel surface; until
/// the dispatch arm uses it the type is reachable only as a phantom.
#[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))]
#[allow(dead_code)]
#[derive(Copy, Clone, Default)]
pub(crate) struct NeonKernel;

#[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))]
impl CpuKernel for NeonKernel {
    #[inline(always)]
    fn mask_lower_bits(value: u64, n: u8) -> u64 {
        // aarch64 has no BMI2 equivalent that improves on the scalar
        // shift-and-mask sequence for this op; the codegen is
        // identical to the Scalar kernel here. Other trait methods
        // (huf_burst, copy_chunk) will diverge once they land.
        ScalarKernel::mask_lower_bits(value, n)
    }
}

/// aarch64 SVE kernel. Variable-vector-length SVE extends NEON for
/// HUF burst / SIMD copy on Graviton3 / Apple M-series with SVE
/// support. Mask op identical to NEON / Scalar.
///
/// `#[allow(dead_code)]`: same scaffolding rationale as `NeonKernel`.
#[cfg(all(target_arch = "aarch64", feature = "kernel-sve"))]
#[allow(dead_code)]
#[derive(Copy, Clone, Default)]
pub(crate) struct SveKernel;

#[cfg(all(target_arch = "aarch64", feature = "kernel-sve"))]
impl CpuKernel for SveKernel {
    #[inline(always)]
    fn mask_lower_bits(value: u64, n: u8) -> u64 {
        ScalarKernel::mask_lower_bits(value, n)
    }
}

/// Single `#[target_feature(enable = "bmi2")]` wrapper around the
/// `_bzhi_u64` intrinsic. Lifted to a free function so each kernel
/// impl that needs the BMI2 path (Bmi2 / Avx2 / Vbmi2) calls the
/// same shared body. With `#[inline]` LLVM inlines the call into
/// any caller that itself has BMI2 in scope; outside that scope the
/// target_feature boundary is preserved.
#[cfg(all(
    any(target_arch = "x86", target_arch = "x86_64"),
    feature = "kernel-bmi2"
))]
#[target_feature(enable = "bmi2")]
#[inline]
unsafe fn mask_lower_bits_bmi2_impl(value: u64, n: u8) -> u64 {
    // The intrinsic call is permitted directly inside a function
    // already annotated `#[target_feature(enable = "bmi2")]` — no
    // `unsafe { ... }` block needed (the function-level `unsafe`
    // already covers it). SAFETY: caller selected a kernel whose
    // CpuKernelTag was resolved after `is_x86_feature_detected!("bmi2")`
    // returned true, so the BMI2 instruction set is available.
    #[cfg(target_arch = "x86_64")]
    {
        core::arch::x86_64::_bzhi_u64(value, n as u32)
    }
    // 32-bit x86 has `bzhi` on 32-bit registers only. Widths up to 32 take one
    // instruction on the low half; wider ones keep the low 32 bits whole and
    // apply it to the high half, which is what a 64-bit `bzhi` does in one go.
    #[cfg(target_arch = "x86")]
    {
        use core::arch::x86::_bzhi_u32;
        if n >= 64 {
            return value;
        }
        if n <= 32 {
            return u64::from(_bzhi_u32(value as u32, u32::from(n)));
        }
        let high = _bzhi_u32((value >> 32) as u32, u32::from(n) - 32);
        (value & u64::from(u32::MAX)) | (u64::from(high) << 32)
    }
}

/// Pure boolean-input variant of the x86 kernel-tag selection. Both the
/// `std` runtime-detect path and the `no_std` compile-time-cfg path
/// route through this helper so the precedence rules stay in one place
/// (and are unit-testable without runtime CPUID).
///
/// The VBMI2 tier requires every AVX-512 sub-feature it touches AND the
/// AVX2 baseline — VBMI2 kernels mix VBMI2-only intrinsics with AVX2
/// 256-bit moves, so the dispatch must be conditioned on `has_avx2` too.
/// Likewise the Avx2 tier requires both AVX2 and BMI2.
#[cfg(target_arch = "x86_64")]
#[inline(always)]
// Params go unused when the matching `kernel_*` feature is disabled (the
// rung that consumes them is `#[cfg]`-ed out); they are still passed by the
// detect callers. Silence the conditional unused-variable warning rather
// than thread per-feature `_`-prefixes through the signature.
#[allow(unused_variables)]
const fn select_x86_kernel(
    has_avx512vbmi2: bool,
    has_avx512f: bool,
    has_avx512vl: bool,
    has_avx512bw: bool,
    has_bmi2: bool,
    has_avx2: bool,
    has_sse2: bool,
) -> CpuKernelTag {
    #[cfg(feature = "kernel-vbmi2")]
    if has_avx512vbmi2 && has_avx512f && has_avx512vl && has_avx512bw && has_bmi2 && has_avx2 {
        return CpuKernelTag::Vbmi2;
    }
    #[cfg(feature = "kernel-avx2")]
    if has_avx2 && has_bmi2 {
        return CpuKernelTag::Avx2;
    }
    #[cfg(feature = "kernel-bmi2")]
    if has_bmi2 {
        return CpuKernelTag::Bmi2;
    }
    #[cfg(feature = "kernel-sse")]
    if has_sse2 {
        return CpuKernelTag::Sse2;
    }
    CpuKernelTag::Scalar
}

/// Cached runtime-detected kernel tag. The actual `CpuKernel` impl
/// (`ScalarKernel` / `Bmi2Kernel` / `Avx2Kernel` / `Vbmi2Kernel` /
/// `NeonKernel` / `SveKernel`) is constructed at the dispatch site —
/// currently only `decoding::literals_section_decoder::decompress_literals`
/// — via a `match` on this tag that branches into the per-K
/// `target_feature`-wrapped specialisation. Pipeline-wide dispatch
/// (FrameDecoder / FrameCompressor entry, sequence executor, match
/// copy) lands incrementally in follow-up tiers.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub(crate) enum CpuKernelTag {
    Scalar,
    #[cfg(all(target_arch = "x86_64", feature = "kernel-sse"))]
    Sse2,
    /// Reachable on 32-bit x86 as well: `bzhi` is there, and without the tier
    /// such a build would decode on the scalar bodies whatever the CPU offers.
    #[cfg(all(
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-bmi2"
    ))]
    Bmi2,
    #[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))]
    Avx2,
    #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))]
    Vbmi2,
    #[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))]
    Neon,
    // Both constructors of `Sve` need a reachable feature: runtime
    // detection via `std::arch::is_aarch64_feature_detected!` (so
    // `feature = "std"`) or compile-time `target_feature = "sve"` in
    // RUSTFLAGS. Without either, the variant is unreachable and a
    // `match` arm referencing it warns as dead.
    #[cfg(all(
        target_arch = "aarch64",
        feature = "kernel-sve",
        any(feature = "std", target_feature = "sve"),
    ))]
    Sve,
}

/// Detect once and cache the best available CPU kernel for the
/// current process. Subsequent calls return the cached tag without
/// re-running CPU-feature detection. Std-only — no-std targets use
/// the compile-time variant below that resolves at build time.
#[cfg(feature = "std")]
pub(crate) fn detect_cpu_kernel() -> CpuKernelTag {
    static CACHED: OnceLock<CpuKernelTag> = OnceLock::new();
    *CACHED.get_or_init(detect_cpu_kernel_uncached)
}

#[cfg(feature = "std")]
fn detect_cpu_kernel_uncached() -> CpuKernelTag {
    #[cfg(target_arch = "x86_64")]
    {
        use std::arch::is_x86_feature_detected;
        // Gate each probe on its tier feature: `cfg!(...)` const-folds, so the
        // `&&` short-circuits away the runtime `is_x86_feature_detected!` call
        // (and its CPUID/cache traffic) for tiers the build disabled — the
        // matching `select_x86_kernel` rung is `#[cfg]`-ed out anyway.
        return select_x86_kernel(
            cfg!(feature = "kernel-vbmi2") && is_x86_feature_detected!("avx512vbmi2"),
            cfg!(feature = "kernel-vbmi2") && is_x86_feature_detected!("avx512f"),
            cfg!(feature = "kernel-vbmi2") && is_x86_feature_detected!("avx512vl"),
            cfg!(feature = "kernel-vbmi2") && is_x86_feature_detected!("avx512bw"),
            cfg!(feature = "kernel-bmi2") && is_x86_feature_detected!("bmi2"),
            cfg!(feature = "kernel-avx2") && is_x86_feature_detected!("avx2"),
            cfg!(feature = "kernel-sse") && is_x86_feature_detected!("sse2"),
        );
    }
    // 32-bit x86 carries only the BMI2 tier: the wider tiers' kernels and
    // their `target_feature` bodies are x86_64-only, so there is nothing
    // above `bzhi` to select here.
    #[cfg(target_arch = "x86")]
    {
        #[cfg(feature = "kernel-bmi2")]
        {
            use std::arch::is_x86_feature_detected;
            if is_x86_feature_detected!("bmi2") {
                return CpuKernelTag::Bmi2;
            }
        }
        return CpuKernelTag::Scalar;
    }
    #[cfg(target_arch = "aarch64")]
    {
        #[cfg(any(feature = "kernel-sve", feature = "kernel-neon"))]
        use std::arch::is_aarch64_feature_detected;
        #[cfg(feature = "kernel-sve")]
        if is_aarch64_feature_detected!("sve") {
            return CpuKernelTag::Sve;
        }
        #[cfg(feature = "kernel-neon")]
        if is_aarch64_feature_detected!("neon") {
            return CpuKernelTag::Neon;
        }
        return CpuKernelTag::Scalar;
    }
    #[allow(unreachable_code)]
    CpuKernelTag::Scalar
}

/// no-std variant: rely on compile-time `target_feature` flags
/// instead of runtime detection. Resolves to the most-capable kernel
/// that the build target supports.
#[cfg(not(feature = "std"))]
pub(crate) fn detect_cpu_kernel() -> CpuKernelTag {
    #[cfg(target_arch = "x86_64")]
    {
        // Route through the same const-fn precedence helper as the
        // `feature = "std"` path. `cfg!(target_feature = ...)`
        // returns a compile-time bool that constant-folds through
        // `select_x86_kernel`, so the runtime call has the same
        // codegen as the previous hand-written #[cfg] chain.
        return select_x86_kernel(
            cfg!(target_feature = "avx512vbmi2"),
            cfg!(target_feature = "avx512f"),
            cfg!(target_feature = "avx512vl"),
            cfg!(target_feature = "avx512bw"),
            cfg!(target_feature = "bmi2"),
            cfg!(target_feature = "avx2"),
            cfg!(target_feature = "sse2"),
        );
    }
    #[cfg(target_arch = "x86")]
    {
        #[cfg(all(feature = "kernel-bmi2", target_feature = "bmi2"))]
        {
            return CpuKernelTag::Bmi2;
        }
    }
    #[cfg(target_arch = "aarch64")]
    {
        #[cfg(all(feature = "kernel-sve", target_feature = "sve"))]
        {
            return CpuKernelTag::Sve;
        }
        #[cfg(all(feature = "kernel-neon", target_feature = "neon"))]
        {
            return CpuKernelTag::Neon;
        }
    }
    #[allow(unreachable_code)]
    CpuKernelTag::Scalar
}

impl CpuKernelTag {
    /// Stable lowercase diagnostic name for this tier (used by
    /// [`active_cpu_kernel_name`] and the bench/dashboard reporting). Pure
    /// mapping over the tag, so every arm is exercisable in tests regardless
    /// of which tier the running CPU actually resolves to.
    pub(crate) fn name(self) -> &'static str {
        match self {
            CpuKernelTag::Scalar => "scalar",
            #[cfg(all(target_arch = "x86_64", feature = "kernel-sse"))]
            CpuKernelTag::Sse2 => "sse2",
            #[cfg(all(
                any(target_arch = "x86", target_arch = "x86_64"),
                feature = "kernel-bmi2"
            ))]
            CpuKernelTag::Bmi2 => "bmi2",
            #[cfg(all(target_arch = "x86_64", feature = "kernel-avx2"))]
            CpuKernelTag::Avx2 => "avx2",
            #[cfg(all(target_arch = "x86_64", feature = "kernel-vbmi2"))]
            CpuKernelTag::Vbmi2 => "vbmi2",
            #[cfg(all(target_arch = "aarch64", feature = "kernel-neon"))]
            CpuKernelTag::Neon => "neon",
            #[cfg(all(
                target_arch = "aarch64",
                feature = "kernel-sve",
                any(feature = "std", target_feature = "sve"),
            ))]
            CpuKernelTag::Sve => "sve",
        }
    }
}

/// Name of the CPU kernel tier this process selected for the entropy /
/// sequence hot paths: decode (literals + FSE sequence decode) and encode
/// (entropy) share this dispatch (see #247). Returned as a stable lowercase
/// string for diagnostics and benchmark/dashboard reporting; the value is
/// what the runtime CPU-feature detection (or compile-time `target_feature`
/// on `no_std`) actually resolves to on this machine, so a dashboard can
/// attribute a measurement to the kernel that produced it.
pub fn active_cpu_kernel_name() -> &'static str {
    detect_cpu_kernel().name()
}

#[cfg(test)]
mod tests;