Skip to main content

keyhog_scanner/hw_probe/
mod.rs

1//! Hardware capability probing with once-cached results.
2//!
3//! Detects CPU features (AVX-512, AVX2, NEON), GPU compute (wgpu/Vulkan),
4//! Hyperscan availability, io_uring support, memory, and core counts.
5//! All detection is done once at startup and cached for the process
6//! lifetime.
7//!
8//! Split into focused submodules by hardware-probe responsibility:
9//!
10//!   * `thresholds` - GPU routing crossover constants consumed through
11//!     the public tier lookup functions.
12//!   * [`tier`] - GPU adapter classification + tier threshold profiles.
13//!   * [`select`] - [`select_backend`] routing logic + env-override
14//!     parsing.
15//!   * [`banner`] - `startup_banner` formatter for the CLI header.
16//!   * [`platform`] - per-OS detection of physical cores, memory,
17//!     and io_uring availability.
18
19use std::sync::OnceLock;
20
21mod banner;
22pub(crate) mod platform;
23pub(crate) mod select;
24mod tier;
25
26pub(crate) mod thresholds;
27
28pub use banner::startup_banner;
29pub use select::{
30    gpu_could_engage, parse_backend_str, select_backend, select_backend_verdict,
31    BackendRoutingReason, BackendRoutingVerdict, BACKEND_OVERRIDE_VALUES,
32};
33pub use tier::{gpu_routing_profile, gpu_routing_profiles, GpuRoutingProfile};
34
35/// Scan execution backend selected for a given workload.
36#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
37#[non_exhaustive]
38pub enum ScanBackend {
39    /// GPU region-presence phase 1 through VYRE's CUDA driver.
40    GpuCuda,
41    /// GPU region-presence phase 1 through VYRE's native Metal driver.
42    GpuMetal,
43    /// GPU region-presence phase 1 through VYRE's WGPU driver.
44    GpuWgpu,
45    /// Hyperscan NFA multi-pattern matching + SIMD prefilter.
46    /// This is the primary high-throughput path on all platforms.
47    SimdCpu,
48    /// Pure CPU: Aho-Corasick literals plus Rust regex extraction. No
49    /// Hyperscan or GPU execution.
50    CpuFallback,
51}
52
53impl ScanBackend {
54    /// Stable label for logs and CLI startup banner.
55    #[must_use]
56    pub fn label(self) -> &'static str {
57        match self {
58            Self::GpuCuda => "gpu-cuda-region-presence",
59            Self::GpuMetal => "gpu-metal-region-presence",
60            Self::GpuWgpu => "gpu-wgpu-region-presence",
61            Self::SimdCpu => "simd-regex",
62            Self::CpuFallback => "cpu-fallback",
63        }
64    }
65
66    /// Whether this route executes on a physical GPU.
67    #[must_use]
68    pub const fn is_gpu(self) -> bool {
69        matches!(self, Self::GpuCuda | Self::GpuMetal | Self::GpuWgpu)
70    }
71}
72
73/// True when this build compiled a scan backend beyond the always-present scalar
74/// [`ScanBackend::CpuFallback`]. Hyperscan (`simd`) and/or the GPU stack (`gpu`).
75///
76/// When this is `false` there is no routing choice: every scan can only run
77/// `CpuFallback`, so autoroute calibration is neither needed nor possible and the
78/// caller resolves the lone backend directly instead of failing closed. This is a
79/// COMPILE-time fact, owned here in the scanner where the `simd`/`gpu` feature gates
80/// actually live: a consumer's own feature flags can diverge (e.g. the CLI's
81/// `ci-lean` enables `keyhog-scanner/simd` without the CLI's own `simd` feature), so
82/// consumers MUST ask the scanner rather than checking their own `cfg!`.
83#[must_use]
84pub const fn multiple_backends_compiled() -> bool {
85    simd_backend_compiled() || gpu_backend_compiled()
86}
87
88/// True when this scanner crate was compiled with the Hyperscan/SIMD backend.
89/// Consumers must query this owner instead of their own Cargo feature namespace:
90/// workspace feature unification can enable a dependency backend without
91/// enabling a same-named feature on the consuming crate.
92#[must_use]
93pub const fn simd_backend_compiled() -> bool {
94    cfg!(feature = "simd")
95}
96
97/// Runtime identity of the dynamically linked Hyperscan/Vectorscan library.
98///
99/// Autoroute persistence must bind SIMD measurements to the library that
100/// actually executed them, not only to the fact that SIMD support compiled.
101#[must_use]
102pub fn hyperscan_runtime_identity() -> Option<String> {
103    #[cfg(feature = "simd")]
104    {
105        Some(hyperscan::version().to_string())
106    }
107    #[cfg(not(feature = "simd"))]
108    {
109        None
110    }
111}
112
113/// True when this scanner crate was compiled with the GPU backend stack.
114/// Autoroute host identity and persisted build evidence use this dependency-owned
115/// fact so a GPU-selected calibration can never be stored without GPU identity.
116#[must_use]
117pub const fn gpu_backend_compiled() -> bool {
118    cfg!(feature = "gpu")
119}
120
121/// Single owner of the SIMD-tier label precedence chain.
122///
123/// The label reported by the startup banner, `keyhog backend`, `keyhog doctor`,
124/// and the backend store must always agree, so the `"AVX-512" > "AVX2" > "NEON"
125/// > "scalar"` precedence lives in exactly ONE place. Callers pass the three
126/// probed CPU-feature booleans (typically `caps.has_avx512`, `caps.has_avx2`,
127/// `caps.has_neon`) and receive the highest-priority label that is available.
128///
129/// Precedence is strict and independent of the lower bits: if `has_avx512` is
130/// true the result is `"AVX-512"` regardless of the other two, and so on down
131/// to `"scalar"` when none are present.
132#[must_use]
133pub const fn simd_label(has_avx512: bool, has_avx2: bool, has_neon: bool) -> &'static str {
134    if has_avx512 {
135        "AVX-512"
136    } else if has_avx2 {
137        "AVX2"
138    } else if has_neon {
139        "NEON"
140    } else {
141        "scalar"
142    }
143}
144
145/// Hardware capabilities detected at startup.
146#[derive(Debug, Clone)]
147pub struct HardwareCaps {
148    pub physical_cores: usize,
149    pub logical_cores: usize,
150    pub has_avx2: bool,
151    pub has_avx512: bool,
152    pub has_neon: bool,
153    pub gpu_available: bool,
154    pub gpu_name: Option<String>,
155    /// WGPU's portable `max_buffer_size`, expressed in MiB and capped by
156    /// KeyHog. The historical field name does not mean physical VRAM; WGPU can
157    /// expose a virtual buffer limit larger than installed device memory.
158    pub gpu_vram_mb: Option<u64>,
159    pub gpu_runtime_identity: Option<String>,
160    /// True when the GPU is a software renderer (llvmpipe/lavapipe) - always slower than CPU.
161    pub gpu_is_software: bool,
162    pub total_memory_mb: Option<u64>,
163    pub io_uring_available: bool,
164    /// True when the `simd` feature is compiled in AND Hyperscan initialized.
165    pub hyperscan_available: bool,
166    /// Runtime identity of the Hyperscan/Vectorscan library that produced these
167    /// capabilities. `None` when `hyperscan_available` is false or the identity
168    /// has not been probed yet.
169    pub hyperscan_runtime_identity: Option<String>,
170}
171
172static HW_PROBE: OnceLock<HardwareCaps> = OnceLock::new();
173
174/// Probe hardware once and cache the result.
175pub fn probe_hardware() -> &'static HardwareCaps {
176    HW_PROBE.get_or_init(|| {
177        let logical_cores = std::thread::available_parallelism()
178            .map(|n| n.get())
179            .unwrap_or(1); // LAW10: host/OS hardware probe parse failure => None/conservative default; perf-only, recall-irrelevant
180        let physical_cores = platform::physical_core_count().unwrap_or(logical_cores); // LAW10: host/OS hardware probe parse failure => None/conservative default; perf-only, recall-irrelevant
181
182        #[cfg(target_arch = "x86_64")]
183        let (has_avx2, has_avx512, has_neon) = (
184            std::arch::is_x86_feature_detected!("avx2"),
185            std::arch::is_x86_feature_detected!("avx512f"),
186            false,
187        );
188        #[cfg(target_arch = "aarch64")]
189        let (has_avx2, has_avx512, has_neon) = (false, false, true);
190        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
191        let (has_avx2, has_avx512, has_neon) = (false, false, false);
192
193        let gpu_probe = crate::gpu::gpu_probe();
194        let gpu_available = gpu_probe.available;
195        let gpu_name = gpu_probe.name;
196        let gpu_vram_mb = gpu_probe.buffer_limit_mb;
197        let gpu_runtime_identity = gpu_probe.runtime_identity;
198        let gpu_is_software = gpu_probe.is_software;
199        if gpu_is_software {
200            tracing::warn!(
201                gpu = ?gpu_name,
202                "Software GPU detected: GPU scanning disabled (slower than CPU)"
203            );
204        }
205
206        let hyperscan_available = cfg!(feature = "simd");
207        let hyperscan_runtime_identity = hyperscan_available
208            .then(hyperscan_runtime_identity)
209            .flatten();
210        let total_memory_mb = platform::detect_total_memory_mb();
211        let io_uring_available = platform::detect_io_uring();
212
213        let caps = HardwareCaps {
214            physical_cores,
215            logical_cores,
216            has_avx2,
217            has_avx512,
218            has_neon,
219            gpu_available,
220            gpu_name: gpu_name.clone(),
221            gpu_vram_mb,
222            gpu_runtime_identity,
223            gpu_is_software,
224            total_memory_mb,
225            io_uring_available,
226            hyperscan_available,
227            hyperscan_runtime_identity,
228        };
229
230        tracing::info!(
231            physical_cores,
232            logical_cores,
233            gpu_available,
234            gpu_name = ?gpu_name,
235            has_avx512 = caps.has_avx512,
236            has_avx2 = caps.has_avx2,
237            has_neon = caps.has_neon,
238            hyperscan = hyperscan_available,
239            io_uring = io_uring_available,
240            "hardware probe complete"
241        );
242
243        caps
244    })
245}
246
247#[cfg(test)]
248#[doc(hidden)]
249pub mod testing {
250    pub use super::{
251        gpu_could_engage, parse_backend_str, probe_hardware, select_backend,
252        select_backend_verdict, startup_banner, BackendRoutingReason, BackendRoutingVerdict,
253        HardwareCaps, ScanBackend,
254    };
255
256    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
257    pub enum GpuTier {
258        High,
259        Mid,
260        Low,
261    }
262
263    fn from_inner(tier: super::tier::GpuTier) -> GpuTier {
264        match tier {
265            super::tier::GpuTier::High => GpuTier::High,
266            super::tier::GpuTier::Mid => GpuTier::Mid,
267            super::tier::GpuTier::Low => GpuTier::Low,
268        }
269    }
270
271    fn to_inner(tier: GpuTier) -> super::tier::GpuTier {
272        match tier {
273            GpuTier::High => super::tier::GpuTier::High,
274            GpuTier::Mid => super::tier::GpuTier::Mid,
275            GpuTier::Low => super::tier::GpuTier::Low,
276        }
277    }
278
279    /// Select the CPU scan backend (SIMD/scalar tier) for the given hardware
280    /// capabilities, ignoring any GPU. Delegates to [`super::select`].
281    pub fn cpu_tier_backend(caps: &HardwareCaps) -> ScanBackend {
282        super::select::cpu_tier_backend(caps)
283    }
284
285    /// Classify a GPU adapter name into a routing [`GpuTier`] (High/Mid/Low).
286    /// `None` (no adapter) classifies to the lowest tier.
287    pub fn classify_gpu_tier(adapter_name: Option<&str>) -> GpuTier {
288        from_inner(super::tier::classify_gpu_tier(adapter_name))
289    }
290
291    /// Minimum workload size (bytes) at which the GPU backend is allowed to
292    /// engage at all for this tier (below it, CPU always wins).
293    pub fn gpu_min_bytes_for_tier(tier: GpuTier) -> u64 {
294        super::tier::gpu_min_bytes_for_tier(to_inner(tier))
295    }
296
297    /// Minimum workload size (bytes) at which the GPU runs *solo* (no CPU
298    /// co-scan) for this tier.
299    pub fn gpu_solo_bytes_for_tier(tier: GpuTier) -> u64 {
300        super::tier::gpu_solo_bytes_for_tier(to_inner(tier))
301    }
302
303    /// Pattern-count break-even above which GPU scanning beats CPU at this tier.
304    pub fn gpu_pattern_breakeven_for_tier(tier: GpuTier) -> usize {
305        super::tier::gpu_pattern_breakeven_for_tier(to_inner(tier))
306    }
307
308    /// Choose the scan backend for a batch from hardware caps, total workload
309    /// size, pattern count, and the largest single-chunk size.
310    pub fn select_backend_for_batch(
311        caps: &HardwareCaps,
312        workload_bytes: u64,
313        pattern_count: usize,
314        large_chunk_bytes: u64,
315    ) -> ScanBackend {
316        super::select::select_backend_for_batch(
317            caps,
318            workload_bytes,
319            pattern_count,
320            large_chunk_bytes,
321        )
322    }
323
324    /// Like [`select_backend_for_batch`] but returns the full
325    /// [`BackendRoutingVerdict`] (the chosen backend plus the inputs and reason
326    /// behind the decision) for diagnostics/telemetry.
327    pub fn select_backend_for_batch_verdict(
328        caps: &HardwareCaps,
329        workload_bytes: u64,
330        pattern_count: usize,
331        large_chunk_bytes: u64,
332    ) -> BackendRoutingVerdict {
333        super::select::select_backend_for_batch_verdict(
334            caps,
335            workload_bytes,
336            pattern_count,
337            large_chunk_bytes,
338        )
339    }
340
341    /// Test-only forced backend override (from the `KEYHOG_*` routing env), or
342    /// `None` when routing is not overridden.
343    pub fn forced_backend_override_for_test() -> Option<ScanBackend> {
344        super::select::forced_backend_override_for_test()
345    }
346
347    /// Parse the physical CPU core count from `/proc/cpuinfo` contents (Linux).
348    #[cfg(target_os = "linux")]
349    pub fn linux_physical_cores_from_cpuinfo(content: &str) -> Option<usize> {
350        super::platform::linux_physical_cores_from_cpuinfo(content)
351    }
352
353    /// Parse total system memory (MiB) from `/proc/meminfo` contents (Linux).
354    #[cfg(target_os = "linux")]
355    pub fn linux_total_memory_mb_from_meminfo(content: &str) -> Option<u64> {
356        super::platform::linux_total_memory_mb_from_meminfo(content)
357    }
358}