lgz/src/bitsplit/simd.rs
//! **The bit-level candidate filter, per CPU.** One step tests 512 bit
//! offsets: 8 lanes of 64, each lane's word shifted against the next word
//! (`s(k) = cur >> k | nxt << (64 − k)`), and the first 13 bits of a
//! non-final dynamic block checked at every offset at once:
//!
//! ```text
//! mask = !s0 & !s1 & s2 & !(s4 & s5 & s6 & s7) & !(s9 & s10 & s11 & s12)
//! BFINAL=0 BTYPE=2 HLIT ≤ 29 HDIST ≤ 29
//! ```
//!
//! Versions: AVX-512 (8 lanes in one register), AVX2 (2 × 4 lanes), NEON
//! (4 × 2 lanes) and the portable 64-bit-integer one, which is the reference.
//! The version is chosen once, at first use (`LGZ_SIMD=portable|avx2|avx512|neon`
//! forces one, for tests and benches).
use std::sync::OnceLock;
/// Words per step: 8 lanes plus the next word for the shifts.
pub const WORDS: usize = 9;
/// Lanes (64-bit masks) per step.
pub const LANES: usize = 8;
/// One version of the filter.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Impl {
Portable,
Avx2,
Avx512,
Neon,
}
impl Impl {
pub const ALL: [Impl; 4] = [Impl::Portable, Impl::Avx2, Impl::Avx512, Impl::Neon];
/// Whether this CPU runs it.
pub fn available(self) -> bool {
match self {
Impl::Portable => true,
#[cfg(target_arch = "x86_64")]
Impl::Avx2 => std::arch::is_x86_feature_detected!("avx2"),
#[cfg(target_arch = "x86_64")]
Impl::Avx512 => {
std::arch::is_x86_feature_detected!("avx512f")
&& std::arch::is_x86_feature_detected!("avx512bw")
}
#[cfg(target_arch = "aarch64")]
Impl::Neon => std::arch::is_aarch64_feature_detected!("neon"),
#[allow(unreachable_patterns)]
_ => false,
}
}
}
/// The version in use: `LGZ_SIMD` if set (and available), else the best one.
pub fn chosen() -> Impl {
static C: OnceLock<Impl> = OnceLock::new();
*C.get_or_init(|| {
let forced = std::env::var("LGZ_SIMD")
.ok()
.and_then(|v| match v.as_str() {
"portable" => Some(Impl::Portable),
"avx2" => Some(Impl::Avx2),
"avx512" => Some(Impl::Avx512),
"neon" => Some(Impl::Neon),
_ => None,
});
match forced {
Some(i) if i.available() => i,
_ => [Impl::Avx512, Impl::Avx2, Impl::Neon]
.into_iter()
.find(|i| i.available())
.unwrap_or(Impl::Portable),
}
})
}
/// The 8 lane masks of one step, by the given version (which must be
/// available; an unavailable one falls back to portable).
#[inline]
pub fn masks_with(i: Impl, w: &[u64; WORDS]) -> [u64; LANES] {
match i {
#[cfg(target_arch = "x86_64")]
Impl::Avx512 if i.available() => unsafe { masks_avx512(w) },
#[cfg(target_arch = "x86_64")]
Impl::Avx2 if i.available() => unsafe { masks_avx2(w) },
#[cfg(target_arch = "aarch64")]
Impl::Neon if i.available() => unsafe { masks_neon(w) },
_ => masks_portable(w),
}
}
#[inline(always)]
fn lane(cur: u64, nxt: u64) -> u64 {
let s = |k: u32| (cur >> k) | (nxt << (64 - k));
let hlit_bad = s(4) & s(5) & s(6) & s(7);
let hdist_bad = s(9) & s(10) & s(11) & s(12);
!cur & !s(1) & s(2) & !hlit_bad & !hdist_bad
}
/// The reference.
pub fn masks_portable(w: &[u64; WORDS]) -> [u64; LANES] {
let mut m = [0u64; LANES];
for j in 0..LANES {
m[j] = lane(w[j], w[j + 1]);
}
m
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn masks_avx2(w: &[u64; WORDS]) -> [u64; LANES] {
use std::arch::x86_64::*;
let mut m = [0u64; LANES];
for h in 0..2 {
// SAFETY: w has 9 words; lanes 4h..4h+4 and 4h+1..4h+5 are in bounds
let cur = unsafe { _mm256_loadu_si256(w.as_ptr().add(4 * h) as *const __m256i) };
let nxt = unsafe { _mm256_loadu_si256(w.as_ptr().add(4 * h + 1) as *const __m256i) };
macro_rules! s {
($k:literal) => {
_mm256_or_si256(
_mm256_srli_epi64::<$k>(cur),
_mm256_slli_epi64::<{ 64 - $k }>(nxt),
)
};
}
let s1 = s!(1);
let s2 = s!(2);
let hlit_bad = _mm256_and_si256(
_mm256_and_si256(s!(4), s!(5)),
_mm256_and_si256(s!(6), s!(7)),
);
let hdist_bad = _mm256_and_si256(
_mm256_and_si256(s!(9), s!(10)),
_mm256_and_si256(s!(11), s!(12)),
);
// andnot(a, b) = !a & b
let x = _mm256_andnot_si256(cur, _mm256_andnot_si256(s1, s2));
let x = _mm256_andnot_si256(hlit_bad, x);
let x = _mm256_andnot_si256(hdist_bad, x);
// SAFETY: m has 8 lanes; 4h..4h+4 is in bounds
unsafe { _mm256_storeu_si256(m.as_mut_ptr().add(4 * h) as *mut __m256i, x) };
}
m
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f,avx512bw")]
unsafe fn masks_avx512(w: &[u64; WORDS]) -> [u64; LANES] {
use std::arch::x86_64::*;
let mut m = [0u64; LANES];
// SAFETY: w has 9 words: 0..8 and 1..9 are in bounds
let cur = unsafe { _mm512_loadu_si512(w.as_ptr() as *const __m512i) };
let nxt = unsafe { _mm512_loadu_si512(w.as_ptr().add(1) as *const __m512i) };
macro_rules! s {
($k:literal) => {
_mm512_or_si512(
_mm512_srli_epi64::<$k>(cur),
_mm512_slli_epi64::<{ 64 - $k }>(nxt),
)
};
}
let s1 = s!(1);
let s2 = s!(2);
let hlit_bad = _mm512_and_si512(
_mm512_and_si512(s!(4), s!(5)),
_mm512_and_si512(s!(6), s!(7)),
);
let hdist_bad = _mm512_and_si512(
_mm512_and_si512(s!(9), s!(10)),
_mm512_and_si512(s!(11), s!(12)),
);
let x = _mm512_andnot_si512(cur, _mm512_andnot_si512(s1, s2));
let x = _mm512_andnot_si512(hlit_bad, x);
let x = _mm512_andnot_si512(hdist_bad, x);
// SAFETY: m has 8 lanes
unsafe { _mm512_storeu_si512(m.as_mut_ptr() as *mut __m512i, x) };
m
}
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
unsafe fn masks_neon(w: &[u64; WORDS]) -> [u64; LANES] {
use std::arch::aarch64::*;
let mut m = [0u64; LANES];
for h in 0..4 {
// SAFETY: w has 9 words; 2h..2h+2 and 2h+1..2h+3 are in bounds
let cur = unsafe { vld1q_u64(w.as_ptr().add(2 * h)) };
let nxt = unsafe { vld1q_u64(w.as_ptr().add(2 * h + 1)) };
macro_rules! s {
($k:literal) => {
vorrq_u64(vshrq_n_u64::<$k>(cur), vshlq_n_u64::<{ 64 - $k }>(nxt))
};
}
let hlit_bad = vandq_u64(vandq_u64(s!(4), s!(5)), vandq_u64(s!(6), s!(7)));
let hdist_bad = vandq_u64(vandq_u64(s!(9), s!(10)), vandq_u64(s!(11), s!(12)));
// vbicq(a, b) = a & !b
let x = vbicq_u64(vbicq_u64(s!(2), s!(1)), cur);
let x = vbicq_u64(vbicq_u64(x, hlit_bad), hdist_bad);
// SAFETY: m has 8 lanes; 2h..2h+2 is in bounds
unsafe { vst1q_u64(m.as_mut_ptr().add(2 * h), x) };
}
m
}