diff --git a/Cargo.toml b/Cargo.toml index 044ca20d..9a1bca5d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,7 @@ bytes = "1.10" cfg-if = "1.0" faststr = { version = "0.2", features = ["serde"] } itoa = "1.0" +multiversion = "0.8.0" ref-cast = "1.0" serde = { version = "1.0", features = ["rc", "derive"] } simdutf8 = "0.1" diff --git a/sonic-number/src/arch/mod.rs b/sonic-number/src/arch/mod.rs index ae8568ed..3e9dc4d0 100644 --- a/sonic-number/src/arch/mod.rs +++ b/sonic-number/src/arch/mod.rs @@ -1,5 +1,5 @@ cfg_if::cfg_if! { - if #[cfg(all(target_arch = "x86_64", target_feature = "avx2", target_feature = "sse2"))] { + if #[cfg(all(target_arch = "x86_64", target_feature = "sse2"))] { mod x86_64; pub use x86_64::*; } else if #[cfg(all(target_feature="neon", target_arch="aarch64"))] { diff --git a/src/util/arch/mod.rs b/src/util/arch/mod.rs index 84808b61..57a5c342 100644 --- a/src/util/arch/mod.rs +++ b/src/util/arch/mod.rs @@ -2,7 +2,12 @@ cfg_if::cfg_if! { if #[cfg(all(target_arch = "x86_64", target_feature = "pclmulqdq", target_feature = "avx2", target_feature = "sse2"))] { mod x86_64; pub use x86_64::*; - } else if #[cfg(all(target_feature="neon", target_arch="aarch64"))] { + } else if #[cfg(target_arch = "x86_64")] { + mod x86_64; + mod fallback; + mod runtime; + pub use runtime::*; + } else if #[cfg(all(target_feature = "neon", target_arch = "aarch64"))] { mod aarch64; pub use aarch64::*; } else { @@ -22,4 +27,26 @@ mod test { let expected_bits = 0b1111111111111111111111111111111111111111111111111111111111110000; assert_eq!(non_space_bits, expected_bits, "bits is {non_space_bits:b}"); } + + #[test] + fn test_prefix_xor() { + // prefix_xor computes a running XOR: each bit in the output is the XOR of itself + // and all lower bits in the input. For input 0b0100, the result is 0b1100...0 (all + // bits from position 2 upward are flipped). + let result = unsafe { prefix_xor(0b0100) }; + assert_eq!(result, 0xFFFF_FFFF_FFFF_FFFC); + + // Two set bits cancel each other in the prefix XOR — the region between them is 1, + // outside is 0. + let result = unsafe { prefix_xor(0b1010) }; + assert_eq!(result, 0x0000_0000_0000_0006); + + // Single bit at position 0. + let result = unsafe { prefix_xor(1) }; + assert_eq!(result, u64::MAX); + + // Zero input → zero output. + let result = unsafe { prefix_xor(0) }; + assert_eq!(result, 0); + } } diff --git a/src/util/arch/runtime.rs b/src/util/arch/runtime.rs new file mode 100644 index 00000000..5195cc19 --- /dev/null +++ b/src/util/arch/runtime.rs @@ -0,0 +1,43 @@ +//! Runtime CPU feature detection for x86_64. +//! +//! This module is only compiled when targeting x86_64 without compile-time AVX2/PCLMULQDQ +//! features. It uses the `multiversion` crate to detect CPU features at runtime and dispatch +//! to either the optimized x86_64 implementations (AVX2/PCLMULQDQ) or the scalar fallback. +//! +//! When the features *are* known at compile time (e.g. via `-C target-cpu=native`), this +//! module is not compiled at all — the optimized implementations are used directly with +//! zero overhead. See the `cfg_if` dispatch in `mod.rs`. + +use multiversion::multiversion; +use multiversion::target::match_target; + +/// Detect non-whitespace bytes in a 64-byte block and return a bitmask. +/// +/// At runtime, dispatches to: +/// - AVX2 shuffle-based implementation if the CPU supports AVX2 +/// - Scalar byte-by-byte fallback otherwise +/// +/// The `multiversion` macro compiles this function once per listed target. The `match_target!` +/// macro resolves at compile time *within each clone*, selecting the appropriate implementation. +/// After the first call the selected function pointer is cached in a static atomic — subsequent +/// calls are a single atomic load + indirect call (~1-2 ns overhead). +#[multiversion(targets("x86_64+avx2", "x86_64+sse2"))] +pub unsafe fn get_nonspace_bits(data: &[u8; 64]) -> u64 { + match_target! { + "x86_64+avx2" => super::x86_64::get_nonspace_bits(data), + _ => super::fallback::get_nonspace_bits(data), + } +} + +/// Compute prefix XOR of a 64-bit bitmask. +/// +/// At runtime, dispatches to: +/// - PCLMULQDQ carryless-multiply implementation if the CPU supports it +/// - Scalar shift-cascade fallback otherwise +#[multiversion(targets("x86_64+pclmulqdq", "x86_64+sse2"))] +pub unsafe fn prefix_xor(bitmask: u64) -> u64 { + match_target! { + "x86_64+pclmulqdq" => super::x86_64::prefix_xor(bitmask), + _ => super::fallback::prefix_xor(bitmask), + } +} diff --git a/src/util/arch/x86_64.rs b/src/util/arch/x86_64.rs index 7bcda4b5..d50a3e4e 100644 --- a/src/util/arch/x86_64.rs +++ b/src/util/arch/x86_64.rs @@ -1,5 +1,9 @@ use std::arch::x86_64::*; +#[cfg_attr( + not(target_feature = "pclmulqdq"), + target_feature(enable = "pclmulqdq") +)] #[inline(always)] pub unsafe fn prefix_xor(bitmask: u64) -> u64 { unsafe { @@ -9,6 +13,7 @@ pub unsafe fn prefix_xor(bitmask: u64) -> u64 { } } +#[cfg_attr(not(target_feature = "avx2"), target_feature(enable = "avx2"))] #[inline(always)] pub unsafe fn get_nonspace_bits(data: &[u8; 64]) -> u64 { unsafe {