keyhog_scanner/
gpu_input_budget.rs

1//! VRAM-adaptive GPU batch-input sizing.
2//!
3//! This module owns the live GPU region-presence byte-budget selector used for
4//! routing and cache-key stability.
5
6use std::sync::atomic::{AtomicUsize, Ordering};
7
8// ---------------------------------------------------------------------------
9// VRAM sizing table: ONE owner for every threshold and byte budget.
10//
11// The adaptive `gpu_batch_input_limit_for_vram_mb` match arms below are the SOLE
12// readers of these; nothing is a bare magic number inline. The `_UNKNOWN` floor
13// (128 MiB) doubles as the lower clamp bound and the `_HIGH` budget (1 GiB) as
14// the upper clamp bound for any Tier-A override, so the operator can never drive
15// the buffer outside the range the table itself honors.
16// ---------------------------------------------------------------------------
17
18/// `>= 24 GiB` VRAM (RTX 4090 / 5090, A100 / H100) -> 1 GiB input.
19pub(crate) const VRAM_MB_TIER_HIGH: u64 = 24 * 1024;
20pub(crate) const GPU_BATCH_INPUT_LIMIT_HIGH: usize = 1024 * 1024 * 1024;
21/// `12 - 23 GiB` VRAM (RTX 3090, RTX 4080, M-Max) -> 512 MiB input.
22pub(crate) const VRAM_MB_TIER_MID: u64 = 12 * 1024;
23pub(crate) const GPU_BATCH_INPUT_LIMIT_MID: usize = 512 * 1024 * 1024;
24/// `8 - 11 GiB` VRAM (RTX 3080, RTX 4070, M-Pro) -> 256 MiB input.
25pub(crate) const VRAM_MB_TIER_LOW: u64 = 8 * 1024;
26pub(crate) const GPU_BATCH_INPUT_LIMIT_LOW: usize = 256 * 1024 * 1024;
27
28/// Conservative floor for hosts with low or unknown VRAM. Unknown must not
29/// inherit the 8-11 GiB tier: absence of adapter memory evidence is the same
30/// safety class as low-memory/iGPU/software adapters. Also the lower clamp bound
31/// for a Tier-A override (see [`set_gpu_batch_input_limit`]).
32pub(crate) const GPU_BATCH_INPUT_LIMIT_UNKNOWN: usize = 128 * 1024 * 1024;
33
34/// Process-wide GPU batch-input override in bytes. `0` = unset (use the
35/// VRAM-adaptive table). Set ONCE at scan startup, before the first
36/// [`gpu_batch_input_limit`] call caches the value, from resolved config (Tier-A:
37/// compiled default -> `.keyhog.toml` -> `--gpu-batch-input-limit`). Mirrors the
38/// `REGEX_DFA_LIMIT_OVERRIDE` process-global pattern so the routing/cache-key
39/// path needs no per-call plumbing.
40static GPU_BATCH_INPUT_LIMIT_OVERRIDE: AtomicUsize = AtomicUsize::new(0);
41
42/// The `[floor, cap]` the resolved GPU batch input limit is clamped into: the
43/// 128 MiB unknown-host floor and the 1 GiB pre-compile-time ceiling that bound
44/// the VRAM table. A Tier-A override is clamped into this range so no config/CLI
45/// value can request a buffer the sizing contract forbids.
46#[must_use]
47pub fn gpu_batch_input_limit_bounds() -> (usize, usize) {
48    (GPU_BATCH_INPUT_LIMIT_UNKNOWN, GPU_BATCH_INPUT_LIMIT_HIGH)
49}
50
51/// Override the GPU batch input limit for this process. Call before scanning.
52/// `0` resets to the VRAM-adaptive default; any other value is clamped into
53/// [`gpu_batch_input_limit_bounds`] at read time. Tier-A config knob
54/// (compiled default -> TOML -> CLI), the sizing analogue of
55/// [`crate::types::set_regex_dfa_limit`].
56pub fn set_gpu_batch_input_limit(bytes: usize) {
57    GPU_BATCH_INPUT_LIMIT_OVERRIDE.store(bytes, Ordering::Relaxed);
58}
59
60/// Clamp a raw Tier-A override into [`gpu_batch_input_limit_bounds`]. Pure, testable
61/// without the process-global, so the clamp contract is proven deterministically.
62pub(crate) fn clamp_gpu_batch_input_limit(bytes: usize) -> usize {
63    let (floor, cap) = gpu_batch_input_limit_bounds();
64    bytes.clamp(floor, cap)
65}
66
67/// Resolve the Tier-A override into an effective byte budget, or `None` when
68/// unset (`0`). Split out so the cached entry point stays thin. Reads the
69/// process-global; the clamp itself lives in [`clamp_gpu_batch_input_limit`].
70pub(crate) fn gpu_batch_input_limit_override() -> Option<usize> {
71    match GPU_BATCH_INPUT_LIMIT_OVERRIDE.load(Ordering::Relaxed) {
72        0 => None,
73        n => Some(clamp_gpu_batch_input_limit(n)),
74    }
75}
76
77/// VRAM-adaptive GPU batch-input limit. Bigger buffers mean fewer
78/// device dispatches per multi-TB scan; each kernel launch is a fixed
79/// ~50-300 µs cost regardless of payload, so doubling the input
80/// halves dispatch overhead. Capped by host VRAM (input + transition
81/// tables + match output must fit) and by a 1 GiB upper bound so the
82/// pre-compile time stays bounded.
83///
84/// | VRAM detected     | Input length | Adapter examples                 |
85/// |-------------------|--------------|----------------------------------|
86/// | >= 24 GiB         | 1 GiB        | RTX 4090 / 5090, A100 / H100     |
87/// | 12 - 23 GiB       | 512 MiB      | RTX 3090, RTX 4080, M-Max        |
88/// | 8 - 11 GiB        | 256 MiB      | RTX 3080, RTX 4070, M-Pro        |
89/// |  < 8 GiB / Unknown| 128 MiB      | iGPU, software, no-GPU CI runner |
90///
91/// Cached on first call; the result is stable for the process
92/// lifetime so routing and cache identities stay consistent across
93/// every batch.
94pub fn gpu_batch_input_limit() -> usize {
95    // Read the explicit Tier-A value on every call. Only the hardware-derived
96    // default is cached, so setting or clearing the override can never be
97    // silently ignored merely because another caller resolved the default first.
98    if let Some(len) = gpu_batch_input_limit_override() {
99        tracing::debug!(
100            target: "keyhog::routing",
101            gpu_batch_input_limit = len,
102            "GPU batch input limit set from Tier-A override"
103        );
104        return len;
105    }
106    use std::sync::OnceLock;
107    static CACHED: OnceLock<usize> = OnceLock::new();
108    *CACHED.get_or_init(|| {
109        let caps = crate::hw_probe::probe_hardware();
110        let len = gpu_batch_input_limit_for_vram_mb(caps.gpu_vram_mb);
111        tracing::debug!(
112            target: "keyhog::routing",
113            gpu_vram_mb = ?caps.gpu_vram_mb,
114            gpu_batch_input_limit = len,
115            "GPU batch input limit sized for VRAM"
116        );
117        len
118    })
119}
120
121pub(crate) fn gpu_batch_input_limit_for_vram_mb(gpu_vram_mb: Option<u64>) -> usize {
122    match gpu_vram_mb {
123        Some(mb) if mb >= VRAM_MB_TIER_HIGH => GPU_BATCH_INPUT_LIMIT_HIGH,
124        Some(mb) if mb >= VRAM_MB_TIER_MID => GPU_BATCH_INPUT_LIMIT_MID,
125        Some(mb) if mb >= VRAM_MB_TIER_LOW => GPU_BATCH_INPUT_LIMIT_LOW,
126        Some(_) | None => GPU_BATCH_INPUT_LIMIT_UNKNOWN,
127    }
128}
keyhog_scanner/gpu_input_budget.rs

keyhog_scanner/
gpu_input_budget.rs
