ashwa/
lib.rs

1//! Hardware accelerated routines for single substring search
2//!
3//! ## Example
4//!
5//! ```
6//! use ashwa::search_one;
7//!
8//! let haystack = b"The quick brown fox jumps over the lazy dog";
9//! assert_eq!(search_one(haystack, b'f'), Some(0x10));
10//! assert_eq!(search_one(haystack, b'z'), Some(0x25));
11//! assert_eq!(search_one(haystack, b'!'), None);
12//! ```
13
14#![cfg_attr(not(test), no_std)]
15#![allow(unsafe_op_in_unsafe_fn)]
16#![deny(
17    missing_docs,
18    trivial_casts,
19    trivial_numeric_casts,
20    unused_extern_crates,
21    unused_import_braces,
22    unused_results
23)]
24
25#[cfg(not(any(target_pointer_width = "64", target_pointer_width = "32")))]
26compile_error!("ashwa is only supported on 64 and 32 bit targets");
27
28#[cfg(all(target_arch = "x86_64"))]
29use core::{arch::x86_64, sync::atomic};
30
31mod one;
32
33pub use one::search_one;
34
35#[repr(u8)]
36#[cfg(target_arch = "x86_64")]
37pub(crate) enum ISA {
38    NONE,
39    SWAR,
40    SSE2,
41    SSSE3,
42    SSE4_2,
43    AVX2,
44    AVX512BW,
45}
46
47#[cfg(target_arch = "x86_64")]
48impl From<u8> for ISA {
49    #[inline(always)]
50    fn from(value: u8) -> Self {
51        match value {
52            0 => ISA::NONE,
53            1 => ISA::SWAR,
54            2 => ISA::SSE2,
55            3 => ISA::SSSE3,
56            4 => ISA::SSE4_2,
57            5 => ISA::AVX2,
58            6 => ISA::AVX512BW,
59            _ => unreachable!("invalid ISA {}", value),
60        }
61    }
62}
63
64/// Best available ISA on the target microarchitecture
65#[cfg(target_arch = "x86_64")]
66static CPU_FEATURE: atomic::AtomicU8 = atomic::AtomicU8::new(ISA::NONE as u8);
67
68#[inline(always)]
69#[allow(unreachable_code)]
70#[cfg(target_arch = "x86_64")]
71pub(crate) fn get_cpu_feature() -> ISA {
72    #[cfg(forced_swar_backend)]
73    return ISA::SWAR;
74
75    #[cfg(target_feature = "avx512bw")]
76    return ISA::AVX512BW;
77
78    #[cfg(target_feature = "avx2")]
79    return ISA::AVX2;
80
81    #[cfg(target_feature = "sse4.2")]
82    return ISA::SSE4_2;
83
84    #[cfg(target_feature = "ssse3")]
85    return ISA::SSSE3;
86
87    #[cfg(target_feature = "sse2")]
88    return ISA::SSE2;
89
90    let feature = CPU_FEATURE.load(atomic::Ordering::Relaxed);
91    if feature != 0 {
92        return feature.into();
93    }
94
95    let detected = unsafe { detect_features_x86_64() };
96    CPU_FEATURE.store(detected as u8, atomic::Ordering::Relaxed);
97
98    detected.into()
99}
100
101#[cold]
102#[inline(never)]
103#[cfg(target_arch = "x86_64")]
104unsafe fn detect_features_x86_64() -> ISA {
105    let cpuid1 = x86_64::__cpuid(1);
106
107    let has_sse2 = (cpuid1.edx & (1 << 0x1A)) != 0;
108    let has_ssse3 = (cpuid1.ecx & (1 << 9)) != 0;
109    let has_sse4_2 = (cpuid1.ecx & (1 << 0x14)) != 0;
110
111    let osxsave = (cpuid1.ecx & (1 << 0x1B)) != 0;
112    if osxsave {
113        let xcr0 = x86_64::_xgetbv(0);
114        let xmm_ymm_enabled = (xcr0 & 0b110) == 0b110;
115
116        if xmm_ymm_enabled {
117            let cpuid7 = x86_64::__cpuid_count(7, 0);
118
119            #[cfg(target_feature = "avx512bw")]
120            {
121                let avx512f_bw = (1 << 16) | (1 << 30);
122                let vbmi_vbmi2 = (1 << 1) | (1 << 6);
123                if (cpuid7.ebx & avx512f_bw) == avx512f_bw
124                    && (cpuid7.ecx & vbmi_vbmi2) == vbmi_vbmi2
125                {
126                    return ISA::AVX512BW;
127                }
128            }
129
130            if (cpuid7.ebx & (1 << 5)) != 0 {
131                return ISA::AVX2;
132            }
133        }
134    }
135
136    if has_sse4_2 {
137        return ISA::SSE4_2;
138    }
139
140    if has_ssse3 {
141        return ISA::SSSE3;
142    }
143
144    if has_sse2 {
145        return ISA::SSE2;
146    }
147
148    ISA::SWAR
149}
ashwa/lib.rs

ashwa/
lib.rs
