simoom_sys/
bits.rs

1//! Codec kernels with CPU features picked at run time (owned by the simoom-compress work).
2//!
3//! [`bmi2`] says whether the CPU has BMI1, BMI2 and LZCNT (x86-64); the loops in
4//! [`crate::codec`] then run a copy compiled with them (`shlx`/`shrx`/`tzcnt`: no shift-count
5//! moves), as libzstd's `DYNAMIC_BMI2` builds do. [`set_scalar_only`] forces the base-target copy
6//! on a thread, so tests can compare the two.
7
8use std::cell::Cell;
9use std::sync::atomic::{AtomicU8, Ordering};
10
11/// 0: not checked yet, 1: absent, 2: present.
12static BMI2: AtomicU8 = AtomicU8::new(0);
13
14thread_local! {
15    static SCALAR_ONLY: Cell<bool> = const { Cell::new(false) };
16}
17
18/// The CPU has BMI1, BMI2 and LZCNT (checked once), and [`set_scalar_only`] is off.
19#[inline]
20pub fn bmi2() -> bool {
21    if SCALAR_ONLY.with(Cell::get) {
22        return false;
23    }
24    match BMI2.load(Ordering::Relaxed) {
25        2 => true,
26        1 => false,
27        _ => {
28            let yes = detect();
29            BMI2.store(if yes { 2 } else { 1 }, Ordering::Relaxed);
30            yes
31        }
32    }
33}
34
35#[cfg(target_arch = "x86_64")]
36fn detect() -> bool {
37    std::arch::is_x86_feature_detected!("bmi1") && std::arch::is_x86_feature_detected!("bmi2") && std::arch::is_x86_feature_detected!("lzcnt")
38}
39
40#[cfg(not(target_arch = "x86_64"))]
41fn detect() -> bool {
42    false
43}
44
45/// Makes the codec loops on this thread use the base-target code from now on (`true`) or pick by
46/// CPU again (`false`): for tests and fuzzing that check both paths give the same results.
47pub fn set_scalar_only(on: bool) {
48    SCALAR_ONLY.with(|c| c.set(on));
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54
55    #[test]
56    fn scalar_switch() {
57        set_scalar_only(true);
58        assert!(!bmi2());
59        set_scalar_only(false);
60    }
61}
simoom_sys/bits.rs

simoom_sys/
bits.rs
