perf(decoding): hold streamed frames to a window plus a block, decode unsized frames into the caller's slice

zstd/src/bit_io/bit_reader_reverse.rs

@@ -1,18 +1,12 @@
 use crate::cpu_kernel::{CpuKernel, ScalarKernel};
 use core::convert::TryInto;
 use core::marker::PhantomData;
-#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))]
-use std::sync::OnceLock;
 
 /// Pre-computed mask table: `BIT_MASK[n]` equals the lower `n` bits set,
-/// i.e. `(1u64 << n) - 1` for `n` in `0..=64`.
-///
-/// `mask_lower_bits` no longer reads this table — it computes the mask
-/// via `u64::MAX >> (64 - n)` to save a load. The table is still used
-/// by the BMI2 PEXT triple-extract path on x86-64 (where the mask is
-/// constructed once per call and then fed to `_pext_u64`), and by the
-/// tests that verify mask values directly.
-#[cfg(any(test, all(target_arch = "x86_64", feature = "kernel-bmi2")))]
+/// i.e. `(1u64 << n) - 1` for `n` in `0..=64`. Kept for the tests that verify
+/// mask values directly; `mask_lower_bits` computes the mask instead of
+/// loading it.
+#[cfg(test)]
 const BIT_MASK: [u64; 65] = {
     let mut table = [0u64; 65];
     let mut i: u32 = 1;
@@ -24,72 +18,6 @@ const BIT_MASK: [u64; 65] = {
     table
 };
 
-#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))]
-#[derive(Copy, Clone)]
-struct TripleExtractDispatch {
-    use_pext: bool,
-}
-
-#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))]
-static TRIPLE_EXTRACT_DISPATCH: OnceLock<TripleExtractDispatch> = OnceLock::new();
-
-#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))]
-#[inline(always)]
-fn should_use_pext(vendor: [u8; 12], family: u32) -> bool {
-    vendor != *b"AuthenticAMD" || family != 0x17
-}
-
-#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))]
-#[inline(always)]
-fn triple_extract_dispatch() -> &'static TripleExtractDispatch {
-    TRIPLE_EXTRACT_DISPATCH.get_or_init(detect_triple_extract_dispatch)
-}
-
-#[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))]
-fn detect_triple_extract_dispatch() -> TripleExtractDispatch {
-    use core::arch::x86_64::__cpuid;
-    use std::arch::is_x86_feature_detected;
-
-    if !is_x86_feature_detected!("bmi2") {
-        return TripleExtractDispatch { use_pext: false };
-    }
-
-    // AMD Zen1/Zen2 execute PEXT/PDEP through a slow microcode path.
-    // Keep scalar extraction there and enable PEXT on Intel and newer AMD.
-    let leaf0 = __cpuid(0);
-    let mut vendor = [0u8; 12];
-    vendor[0..4].copy_from_slice(&leaf0.ebx.to_le_bytes());
-    vendor[4..8].copy_from_slice(&leaf0.edx.to_le_bytes());
-    vendor[8..12].copy_from_slice(&leaf0.ecx.to_le_bytes());
-    let eax = __cpuid(1).eax;
-    let base_family = (eax >> 8) & 0xF;
-    let ext_family = (eax >> 20) & 0xFF;
-    let family = if base_family == 0xF {
-        base_family + ext_family
-    } else {
-        base_family
-    };
-
-    TripleExtractDispatch {
-        use_pext: should_use_pext(vendor, family),
-    }
-}
-
-#[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))]
-#[target_feature(enable = "bmi2")]
-unsafe fn extract_triple_pext(all_three: u64, n1: u8, n2: u8, n3: u8) -> (u64, u64, u64) {
-    use core::arch::x86_64::_pext_u64;
-
-    let mask3 = BIT_MASK[n3 as usize];
-    let mask2 = BIT_MASK[n2 as usize].wrapping_shl(u32::from(n3));
-    let mask1 = BIT_MASK[n1 as usize].wrapping_shl(u32::from(n2) + u32::from(n3));
-
-    let val1 = _pext_u64(all_three, mask1);
-    let val2 = _pext_u64(all_three, mask2);
-    let val3 = _pext_u64(all_three, mask3);
-    (val1, val2, val3)
-}
-
 /// Zstandard encodes some types of data in a way that the data must be read
 /// back to front to decode it properly. `BitReaderReversed` provides a
 /// convenient interface to do that.
@@ -144,16 +72,6 @@ pub struct BitReaderReversed<'s, K: CpuKernel = ScalarKernel> {
     /// drives monomorphisation of methods that route through `K::mask_lower_bits`
     /// without forcing the struct itself to carry runtime kernel state.
     _kernel: PhantomData<K>,
-
-    /// Cached `triple_extract_dispatch().use_pext` snapshot, populated
-    /// once in `new()`. `peek_bits_triple` reads this field instead of
-    /// re-checking the global `OnceLock` on every sequence — the
-    /// per-call atomic load + dispatch-branch was paying ~3 cycles on
-    /// every sequence decode (thousands per block × many blocks per
-    /// frame). One bool per `BitReaderReversed` lifetime, amortised
-    /// across every `peek_bits_triple` in the same decode pass.
-    #[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))]
-    pub(crate) use_pext_triple: bool,
 }
 
 impl<'s, K: CpuKernel> BitReaderReversed<'s, K> {
@@ -162,32 +80,6 @@ impl<'s, K: CpuKernel> BitReaderReversed<'s, K> {
         self.index as isize * 8 + (64 - self.bits_consumed as isize) - self.extra_bits as isize
     }
 
-    /// Returns `true` when the cached vendor policy says PEXT is fast
-    /// on the running CPU (Intel + AMD Zen3+) and the bmi2-direct
-    /// triple-extract path should be used. AMD Zen1/Zen2 microcode
-    /// PEXT is slower than the scalar 3× shift+mask path, so
-    /// [`should_use_pext`] caches `false` for those vendors.
-    ///
-    /// `no_std` x86_64 builds lack the runtime detection (`use_pext_triple`
-    /// is std-gated), so this falls back to `true`: callers on
-    /// `no_std` rely on compile-time `target_feature = "bmi2"` and
-    /// implicitly trust that the chosen target CPU advertises fast
-    /// PEXT. Vendor-specific microcode regression remains a
-    /// build-time concern there — pin a known-good target with
-    /// `RUSTFLAGS="-C target-cpu=..."`.
-    #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))]
-    #[inline(always)]
-    pub(crate) fn use_pext_triple_fast(&self) -> bool {
-        #[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))]
-        {
-            self.use_pext_triple
-        }
-        #[cfg(not(all(feature = "std", target_arch = "x86_64")))]
-        {
-            true
-        }
-    }
-
     pub fn new(source: &'s [u8]) -> BitReaderReversed<'s, K> {
         BitReaderReversed {
             index: source.len(),
@@ -196,8 +88,6 @@ impl<'s, K: CpuKernel> BitReaderReversed<'s, K> {
             bit_container: 0,
             extra_bits: 0,
             _kernel: PhantomData,
-            #[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))]
-            use_pext_triple: triple_extract_dispatch().use_pext,
         }
     }
 
@@ -366,21 +256,10 @@ impl<'s, K: CpuKernel> BitReaderReversed<'s, K> {
         let shift_by = (64u8 - self.bits_consumed).wrapping_sub(sum);
         let all_three = self.bit_container.wrapping_shr(shift_by as u32);
 
-        #[cfg(all(feature = "std", target_arch = "x86_64", feature = "kernel-bmi2"))]
-        if self.use_pext_triple {
-            // SAFETY: `use_pext_triple` was set in `new()` from
-            // `triple_extract_dispatch().use_pext`, which only returns
-            // `true` when BMI2 is runtime-detected; the unsafe call is
-            // gated on the same runtime check that the inline-form
-            // `try_extract_triple_with_pext` used to perform per-call.
-            return unsafe { extract_triple_pext(all_three, n1, n2, n3) };
-        }
-
-        let val1 = K::mask_lower_bits(all_three.wrapping_shr(u32::from(n3) + u32::from(n2)), n1);
-        let val2 = K::mask_lower_bits(all_three.wrapping_shr(u32::from(n3)), n2);
-        let val3 = K::mask_lower_bits(all_three, n3);
-
-        (val1, val2, val3)
+        // The kernel was chosen where this decode was dispatched, so the split
+        // is the monomorph's own instruction sequence. The reader used to carry
+        // the choice as a flag and branch on it here, once per sequence.
+        K::extract_triple(all_three, n1, n2, n3)
     }
 
     /// BMI2-scoped variant of [`peek_bits`]. The whole body executes
@@ -411,52 +290,6 @@ impl<'s, K: CpuKernel> BitReaderReversed<'s, K> {
         core::arch::x86_64::_bzhi_u64(self.bit_container.wrapping_shr(shift_by as u32), n as u32)
     }
 
-    /// BMI2-scoped variant of [`peek_bits_triple`]. Mirrors the
-    /// scalar/K-trait variant but inlines `_pext_u64` directly instead
-    /// of crossing the `extract_triple_pext` CALL boundary.
-    ///
-    /// On AMD Zen1/Zen2 (vendor=AuthenticAMD family=0x17) `_pext_u64`
-    /// goes through slow microcode; callers should still consult
-    /// `self.use_pext_triple` (populated at construction from the
-    /// global dispatch cache) and route to the scalar variant on
-    /// those CPUs. This method assumes the caller already gated on
-    /// `use_pext_triple == true`.
-    ///
-    /// # Safety
-    /// Caller MUST ensure BMI2 is available AND the running CPU
-    /// benefits from `_pext_u64` (i.e. not Zen1/Zen2).
-    #[cfg(all(target_arch = "x86_64", feature = "kernel-bmi2"))]
-    #[target_feature(enable = "bmi2")]
-    #[inline]
-    pub(crate) unsafe fn peek_bits_triple_bmi2(
-        &mut self,
-        sum: u8,
-        n1: u8,
-        n2: u8,
-        n3: u8,
-    ) -> (u64, u64, u64) {
-        debug_assert_eq!(
-            u16::from(sum),
-            u16::from(n1) + u16::from(n2) + u16::from(n3),
-            "peek_bits_triple_bmi2: sum ({}) must equal n1+n2+n3 ({}+{}+{})",
-            sum,
-            n1,
-            n2,
-            n3
-        );
-        debug_assert!(
-            sum == 0 || self.bits_consumed + sum <= 64,
-            "peek_bits_triple_bmi2: not enough bits (consumed={}, requested={})",
-            self.bits_consumed,
-            sum
-        );
-        let shift_by = (64u8 - self.bits_consumed).wrapping_sub(sum);
-        let all_three = self.bit_container.wrapping_shr(shift_by as u32);
-        // SAFETY: caller's target_feature includes BMI2 per `# Safety`
-        // contract; same scope as the enclosing fn.
-        unsafe { extract_triple_pext(all_three, n1, n2, n3) }
-    }
-
     /// Consume `n` bits from the source.
     #[inline(always)]
     pub fn consume(&mut self, n: u8) {