Skip to main content

keyhog_scanner/hw_probe/
select.rs

1//! Workload-aware backend routing. File and batch entry points delegate to one
2//! private workload selector so explicit test overrides, GPU suppression, GPU
3//! thresholds, and CPU-tier fallback cannot drift.
4
5use super::tier::{
6    classify_gpu_tier, gpu_min_bytes_for_tier, gpu_pattern_breakeven_for_tier, gpu_routing_profile,
7    gpu_solo_bytes_for_tier,
8};
9use super::{HardwareCaps, ScanBackend};
10
11thread_local! {
12    pub(crate) static TEST_BACKEND_OVERRIDE: std::cell::RefCell<Option<Option<ScanBackend>>> = const { std::cell::RefCell::new(None) };
13}
14
15#[cfg(test)]
16pub(crate) fn set_test_backend_override(val: Option<ScanBackend>) {
17    TEST_BACKEND_OVERRIDE.with(|cell| {
18        *cell.borrow_mut() = Some(val);
19    });
20}
21
22#[cfg(test)]
23pub(crate) fn clear_test_backend_override() {
24    TEST_BACKEND_OVERRIDE.with(|cell| {
25        *cell.borrow_mut() = None;
26    });
27}
28
29/// The CPU-only backend tier for this hardware: `SimdCpu` only when the
30/// Hyperscan/Vectorscan prefilter is compiled in and live, otherwise the
31/// pure-scalar `CpuFallback`. CPU ISA flags are reported for operator
32/// visibility, but they do not by themselves prove the `simd-regex` backend
33/// exists. This is the SINGLE source of truth for the "no GPU in play"
34/// decision - every router that needs a non-GPU backend (`select_backend`,
35/// `select_backend_for_batch`, and the CLI's measured autoroute default)
36/// routes through here so the four-way ladder can never drift between sites.
37#[must_use]
38pub(crate) fn cpu_tier_backend(caps: &HardwareCaps) -> ScanBackend {
39    if caps.hyperscan_available {
40        ScanBackend::SimdCpu
41    } else {
42        ScanBackend::CpuFallback
43    }
44}
45
46#[derive(Debug, Clone, Copy)]
47struct BackendWorkload {
48    bytes: u64,
49    pattern_count: usize,
50    large_chunk_bytes: Option<u64>,
51}
52
53impl BackendWorkload {
54    fn file(bytes: u64, pattern_count: usize) -> Self {
55        Self {
56            bytes,
57            pattern_count,
58            large_chunk_bytes: None,
59        }
60    }
61
62    #[cfg(test)]
63    fn batch(bytes: u64, pattern_count: usize, large_chunk_bytes: u64) -> Self {
64        Self {
65            bytes,
66            pattern_count,
67            large_chunk_bytes: Some(large_chunk_bytes),
68        }
69    }
70
71    fn gpu_dominates_dispatch_cost(self) -> bool {
72        match self.large_chunk_bytes {
73            None => true,
74            Some(large_chunk_bytes) => {
75                large_chunk_bytes > 0 && large_chunk_bytes.saturating_mul(2) >= self.bytes
76            }
77        }
78    }
79}
80
81#[derive(Debug, Clone, Copy, PartialEq, Eq)]
82pub enum BackendRoutingReason {
83    TestOverride,
84    GpuDisabledByPolicy,
85    GpuProbeMiss,
86    GpuSoftwareRenderer,
87    GpuBatchNotDominant,
88    GpuThresholdNotMet,
89    GpuSelected,
90}
91
92impl BackendRoutingReason {
93    /// Stable machine-readable label for this routing reason (telemetry/logs).
94    #[must_use]
95    pub fn label(self) -> &'static str {
96        match self {
97            Self::TestOverride => "test_override",
98            Self::GpuDisabledByPolicy => "gpu_disabled_by_policy",
99            Self::GpuProbeMiss => "gpu_probe_miss",
100            Self::GpuSoftwareRenderer => "gpu_software_renderer",
101            Self::GpuBatchNotDominant => "gpu_batch_not_dominant",
102            Self::GpuThresholdNotMet => "gpu_threshold_not_met",
103            Self::GpuSelected => "gpu_selected",
104        }
105    }
106}
107
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
109pub struct BackendRoutingVerdict {
110    pub backend: ScanBackend,
111    pub reason: BackendRoutingReason,
112    pub workload_bytes: u64,
113    pub pattern_count: usize,
114    pub large_chunk_bytes: Option<u64>,
115    pub gpu_available: bool,
116    pub gpu_is_software: bool,
117    pub gpu_tier: &'static str,
118    pub gpu_min_bytes: u64,
119    pub gpu_solo_bytes: u64,
120    pub gpu_pattern_breakeven: usize,
121}
122
123impl BackendRoutingVerdict {
124    fn new(
125        caps: &HardwareCaps,
126        workload: BackendWorkload,
127        backend: ScanBackend,
128        reason: BackendRoutingReason,
129    ) -> Self {
130        let profile = gpu_routing_profile(caps.gpu_name.as_deref());
131        Self {
132            backend,
133            reason,
134            workload_bytes: workload.bytes,
135            pattern_count: workload.pattern_count,
136            large_chunk_bytes: workload.large_chunk_bytes,
137            gpu_available: caps.gpu_available,
138            gpu_is_software: caps.gpu_is_software,
139            gpu_tier: profile.tier,
140            gpu_min_bytes: profile.min_bytes,
141            gpu_solo_bytes: profile.solo_bytes,
142            gpu_pattern_breakeven: profile.pattern_breakeven,
143        }
144    }
145
146    /// Human-readable one-line explanation of this backend-routing verdict.
147    #[must_use]
148    pub fn reason_detail(self) -> String {
149        match self.reason {
150            BackendRoutingReason::TestOverride => "forced by test override".to_string(),
151            BackendRoutingReason::GpuDisabledByPolicy => {
152                "GPU disabled by resolved runtime policy".to_string()
153            }
154            BackendRoutingReason::GpuProbeMiss => {
155                "no usable GPU adapter reported by hardware probe".to_string()
156            }
157            BackendRoutingReason::GpuSoftwareRenderer => {
158                "GPU adapter is a software renderer and is slower than CPU/SIMD".to_string()
159            }
160            BackendRoutingReason::GpuBatchNotDominant => {
161                let Some(large) = self.large_chunk_bytes else {
162                    return format!(
163                        "large-chunk byte share is unavailable for workload bytes ({})",
164                        self.workload_bytes
165                    );
166                };
167                format!(
168                    "large-chunk bytes ({large}) do not dominate workload bytes ({})",
169                    self.workload_bytes
170                )
171            }
172            BackendRoutingReason::GpuThresholdNotMet => format!(
173                "GPU thresholds not met for tier {}: bytes={} min={} solo={} patterns={} pattern_floor={}",
174                self.gpu_tier,
175                self.workload_bytes,
176                self.gpu_min_bytes,
177                self.gpu_solo_bytes,
178                self.pattern_count,
179                self.gpu_pattern_breakeven
180            ),
181            BackendRoutingReason::GpuSelected => format!(
182                "GPU thresholds met for tier {}: bytes={} min={} solo={} patterns={} pattern_floor={}",
183                self.gpu_tier,
184                self.workload_bytes,
185                self.gpu_min_bytes,
186                self.gpu_solo_bytes,
187                self.pattern_count,
188                self.gpu_pattern_breakeven
189            ),
190        }
191    }
192}
193
194fn select_backend_for_workload(
195    caps: &HardwareCaps,
196    workload: BackendWorkload,
197) -> BackendRoutingVerdict {
198    if let Some(forced) = test_backend_override() {
199        return BackendRoutingVerdict::new(
200            caps,
201            workload,
202            forced,
203            BackendRoutingReason::TestOverride,
204        );
205    }
206
207    // Skip GPU consideration when the resolved scanner runtime policy disables
208    // GPU init, so the routing decision matches what the GPU init paths will
209    // actually do.
210    let cpu_backend = cpu_tier_backend(caps);
211    if crate::gpu::gpu_disabled_by_policy() {
212        return BackendRoutingVerdict::new(
213            caps,
214            workload,
215            cpu_backend,
216            BackendRoutingReason::GpuDisabledByPolicy,
217        );
218    }
219
220    if !caps.gpu_available {
221        return BackendRoutingVerdict::new(
222            caps,
223            workload,
224            cpu_backend,
225            BackendRoutingReason::GpuProbeMiss,
226        );
227    }
228
229    if caps.gpu_is_software {
230        return BackendRoutingVerdict::new(
231            caps,
232            workload,
233            cpu_backend,
234            BackendRoutingReason::GpuSoftwareRenderer,
235        );
236    }
237
238    if !workload.gpu_dominates_dispatch_cost() {
239        return BackendRoutingVerdict::new(
240            caps,
241            workload,
242            cpu_backend,
243            BackendRoutingReason::GpuBatchNotDominant,
244        );
245    }
246
247    if gpu_could_engage(caps, workload.bytes, workload.pattern_count) {
248        return BackendRoutingVerdict::new(
249            caps,
250            workload,
251            ScanBackend::GpuWgpu,
252            BackendRoutingReason::GpuSelected,
253        );
254    }
255
256    BackendRoutingVerdict::new(
257        caps,
258        workload,
259        cpu_backend,
260        BackendRoutingReason::GpuThresholdNotMet,
261    )
262}
263
264/// Auto-route a scan to the best backend for this hardware + workload.
265///
266/// Routing rules (highest-priority match wins):
267///
268/// 0. **Test override** - scanner tests may force a backend through the
269///    race-free testing facade. Shipped CLI scans pass explicit `--backend`
270///    choices directly to `scan_with_backend` instead of mutating process env.
271/// 1. **GPU** - discrete non-software adapter is present AND the workload is
272///    large enough to amortize device-dispatch overhead AND we have either
273///    enough patterns to benefit from massively-parallel literal matching, OR
274///    a single very large file at or above the tier solo cap where one device
275///    dispatch can beat saturating one CPU core with Hyperscan.
276/// 2. **SimdCpu** - Hyperscan/Vectorscan is compiled in and live. This is the
277///    default high-throughput path for most deployments.
278/// 3. **CpuFallback** - pure scalar AC + regex. Works everywhere.
279///
280/// The crossover thresholds were tuned against the standard corpus (Django +
281/// kubernetes/kubernetes + linux/linux). See [`super::thresholds`].
282#[must_use]
283pub fn select_backend(
284    caps: &HardwareCaps,
285    workload_bytes: u64,
286    pattern_count: usize,
287) -> ScanBackend {
288    select_backend_verdict(caps, workload_bytes, pattern_count).backend
289}
290
291/// Backend-routing verdict for a single-file workload of `workload_bytes`
292/// scanned across `pattern_count` patterns (the chosen backend plus its reason).
293#[must_use]
294pub fn select_backend_verdict(
295    caps: &HardwareCaps,
296    workload_bytes: u64,
297    pattern_count: usize,
298) -> BackendRoutingVerdict {
299    select_backend_for_workload(caps, BackendWorkload::file(workload_bytes, pattern_count))
300}
301
302/// Batch-aware backend routing (a pure, hardware-only library router).
303///
304/// NOTE on the live CLI path: the shipped scan dispatcher does NOT call this;
305/// it uses the measured, parity-checked `MeasuredBackendRouter`
306/// (`crates/cli/src/orchestrator/dispatch/backend.rs`), which benchmarks the
307/// candidate backends on a real sample and gates the GPU behind explicit
308/// `--autoroute-gpu` calibration eligibility (GPU region presence is slower
309/// than SIMD on keyhog's workload through the measured range). This function is the deterministic,
310/// side-effect-free dominance heuristic used by the `keyhog backend` report and
311/// by callers that want a backend decision without running the scanner, it
312/// shares [`cpu_tier_backend`] and [`gpu_could_engage`] with the live router so
313/// the CPU-tier verdict never diverges.
314///
315/// Identical to [`select_backend`] for the CPU tiers, but adds a structural
316/// guard before the GPU branch: `large_chunk_bytes`
317/// is the number of bytes in the batch that live in *large* chunks - chunks at
318/// or above the tier's `gpu_min_bytes` floor (the per-file size below which a
319/// chunk can never carry its share of the device-dispatch cost).
320///
321/// `select_backend` decides on `workload_bytes` alone - the coalesced batch
322/// total. That conflates two workloads the GPU treats very differently:
323///
324///   * a batch *dominated* by genuinely large files (e.g. minified bundles,
325///     data blobs, generated headers) - the GPU's massively-parallel literal/
326///     AC kernel scans those contiguous regions far faster than one Hyperscan
327///     core, amortizing the fixed per-batch device-dispatch + PCIe-copy +
328///     readback + host-side match-attribution cost; and
329///   * a *swarm* of tiny files whose sizes merely SUM past the GPU floor
330///     (the Linux kernel: 94k files, 1.5 GiB, but only 55 files >= 2 MiB and a
331///     single 22 MiB max - the tiny files coalesce into 256 MiB batches). Here
332///     the GPU re-scans every byte, surfaces a literal hit for every detector-
333///     prefix occurrence across the whole buffer, then hands the CPU the SAME
334///     per-chunk phase-2 confirmation it would have run anyway - plus the
335///     coalesce/copy/readback the SIMD path never pays. Measured on the kernel
336///     this routes ~2.1x SLOWER (204 s vs 96 s) at ~3x peak RSS (4.1 vs 2.3
337///     GiB), and the unbounded device wait can stall the whole scan when the
338///     driver drops a completion.
339///
340/// A largest-chunk guard is not enough: the kernel's 55 large files are
341/// sprinkled through the walk, so nearly every 4096-file batch catches one and
342/// would still route to GPU. The robust signal is DOMINANCE - GPU engages only
343/// when large-chunk bytes are at least half the batch, so a tiny-file swarm
344/// never qualifies no matter how the large files cluster, while a batch that is
345/// mostly big-file data still gets the device. An explicit CLI backend override
346/// still wins (forced/diagnostic GPU path unchanged), and benchmarks should pin
347/// `--backend simd`, so this only changes the *default* routing for many-small-
348/// file trees - the common real-world scan.
349#[must_use]
350#[cfg(test)]
351pub(crate) fn select_backend_for_batch(
352    caps: &HardwareCaps,
353    workload_bytes: u64,
354    pattern_count: usize,
355    large_chunk_bytes: u64,
356) -> ScanBackend {
357    select_backend_for_batch_verdict(caps, workload_bytes, pattern_count, large_chunk_bytes).backend
358}
359
360#[must_use]
361#[cfg(test)]
362pub(crate) fn select_backend_for_batch_verdict(
363    caps: &HardwareCaps,
364    workload_bytes: u64,
365    pattern_count: usize,
366    large_chunk_bytes: u64,
367) -> BackendRoutingVerdict {
368    select_backend_for_workload(
369        caps,
370        BackendWorkload::batch(workload_bytes, pattern_count, large_chunk_bytes),
371    )
372}
373
374/// Cheap, side-effect-free pre-check: could a scan of `workload_bytes` over
375/// `pattern_count` patterns ever route to a GPU backend on this
376/// hardware? This is exactly the GPU branch condition inside
377/// [`select_backend`], factored out so cold-path callers can gate the
378/// expensive wgpu/CUDA device acquisition (the ~250 ms adapter-enumeration
379/// cold-start in `compiled_scanner::compile`) on whether the workload can clear the
380/// tier's GPU floor at all.
381///
382/// On a many-tiny-file corpus the per-batch byte total never reaches the
383/// high-tier measured-safe floor (see [`super::thresholds`]), so this returns
384/// `false` and the caller can skip paying for a device no chunk will ever touch.
385/// It does **not** consult explicit backend overrides or `--no-gpu`;
386/// callers that need an override should pass it through their own resolved
387/// config before falling back to this hardware-only predicate.
388#[must_use]
389pub fn gpu_could_engage(caps: &HardwareCaps, workload_bytes: u64, pattern_count: usize) -> bool {
390    if !caps.gpu_available || caps.gpu_is_software {
391        return false;
392    }
393    let tier = classify_gpu_tier(caps.gpu_name.as_deref());
394    let solo = gpu_solo_bytes_for_tier(tier);
395    let min = gpu_min_bytes_for_tier(tier);
396    let pattern_floor = gpu_pattern_breakeven_for_tier(tier);
397    workload_bytes >= solo || (workload_bytes >= min && pattern_count >= pattern_floor)
398}
399
400/// Test-only forced backend override.
401#[cfg(test)]
402pub(crate) fn forced_backend_override_for_test() -> Option<ScanBackend> {
403    test_backend_override()
404}
405
406pub(super) fn test_backend_override() -> Option<ScanBackend> {
407    // The thread-local holds `Option<Option<ScanBackend>>` (outer = "is an
408    // override set", inner = "the forced backend, or None for forced-CPU);
409    // collapsing the two layers is exactly `Option::flatten`.
410    TEST_BACKEND_OVERRIDE.with(|cell| *cell.borrow()).flatten()
411}
412
413/// Operator-facing `--backend` values accepted by the CLI.
414///
415/// Keep this list at the parser owner so Clap validation, error messages, docs
416/// gates, and `parse_backend_str` cannot drift into rejecting canonical labels
417/// before routing sees them.
418pub const BACKEND_OVERRIDE_VALUES: [&str; 9] = [
419    "auto",
420    "gpu-cuda",
421    "gpu-cuda-region-presence",
422    "gpu-wgpu",
423    "gpu-wgpu-region-presence",
424    "simd",
425    "simd-regex",
426    "cpu",
427    "cpu-fallback",
428];
429
430/// Pure backend string → [`ScanBackend`] mapping, with no env or
431/// thread-local override read. Tests that only verify the string→backend
432/// mapping MUST call this directly rather than mutating global process state.
433/// Keeping the mapping pure removes parallel-test hazards. The CLI/config
434/// boundary admits only [`BACKEND_OVERRIDE_VALUES`]; this parser additionally
435/// reads the stable descriptive labels stored in autoroute evidence.
436pub fn parse_backend_str(raw: &str) -> Option<ScanBackend> {
437    match raw.trim().to_ascii_lowercase().as_str() {
438        "gpu-cuda" | "gpu-cuda-region-presence" => Some(ScanBackend::GpuCuda),
439        "gpu-wgpu" | "gpu-wgpu-region-presence" => Some(ScanBackend::GpuWgpu),
440        "simd" | "simd-regex" => Some(ScanBackend::SimdCpu),
441        "cpu" | "cpu-fallback" => Some(ScanBackend::CpuFallback),
442        _ => None,
443    }
444}