Skip to main content

keyhog_scanner/gpu/
policy.rs

1//! GPU runtime policy + require-GPU preflight policy.
2//!
3//! Split out of `gpu.rs` (Law 5 / 500-LOC modularity cap): these are the
4//! explicit runtime-policy readers plus the `require-GPU` preflight that fails
5//! closed when a GPU is demanded but absent. Re-exported from `gpu` via
6//! `pub use policy::*`.
7
8use std::{
9    fmt,
10    sync::atomic::{AtomicU8, Ordering},
11};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14#[repr(u8)]
15pub enum GpuRuntimePolicy {
16    Auto = 0,
17    Disabled = 1,
18    Required = 2,
19}
20
21impl GpuRuntimePolicy {
22    #[must_use]
23    pub(crate) fn label(self) -> &'static str {
24        match self {
25            Self::Auto => "auto",
26            Self::Disabled => "off",
27            Self::Required => "required",
28        }
29    }
30
31    fn from_u8(value: u8) -> Self {
32        match value {
33            1 => Self::Disabled,
34            2 => Self::Required,
35            _ => Self::Auto,
36        }
37    }
38}
39
40impl fmt::Display for GpuRuntimePolicy {
41    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42        f.write_str((*self).label())
43    }
44}
45
46static GPU_RUNTIME_POLICY: AtomicU8 = AtomicU8::new(GpuRuntimePolicy::Auto as u8);
47
48#[derive(Debug, Clone, Default)]
49pub(crate) struct GpuRuntimeProbe {
50    pub(crate) available: bool,
51    pub(crate) name: Option<String>,
52    pub(crate) buffer_limit_mb: Option<u64>,
53    pub(crate) runtime_identity: Option<String>,
54    pub(crate) is_software: bool,
55}
56
57/// Set the process-wide GPU runtime policy (`Auto`/`On`/`Off`) consulted by
58/// backend routing and GPU init.
59pub fn set_gpu_runtime_policy(policy: GpuRuntimePolicy) {
60    GPU_RUNTIME_POLICY.store(policy as u8, Ordering::SeqCst);
61}
62
63/// The current process-wide GPU runtime policy.
64#[must_use]
65pub fn gpu_runtime_policy() -> GpuRuntimePolicy {
66    GpuRuntimePolicy::from_u8(GPU_RUNTIME_POLICY.load(Ordering::SeqCst))
67}
68
69/// Probe GPU availability and adapter metadata without panicking.
70///
71/// Honours the explicit disabled GPU policy by reporting "no GPU available"
72/// without ever calling `backend::get_gpu()`. The MoE compute-shader init
73/// happens lazily inside `get_gpu()`, so this short-circuit is the difference
74/// between "adapter request blocks for minutes on broken driver stacks" and
75/// "scanner starts like every other CPU-only tool".
76#[must_use]
77pub(crate) fn gpu_probe() -> GpuRuntimeProbe {
78    if gpu_disabled_by_policy() {
79        return GpuRuntimeProbe::default();
80    }
81    #[cfg(feature = "gpu")]
82    if let Some(gpu) = super::gpu_adapter_probe() {
83        return GpuRuntimeProbe {
84            available: !gpu.is_software,
85            name: Some(gpu.name.clone()),
86            buffer_limit_mb: Some(gpu.buffer_limit_mb),
87            runtime_identity: Some(gpu.runtime_identity.clone()),
88            is_software: gpu.is_software,
89        };
90    }
91    GpuRuntimeProbe::default()
92}
93
94/// True when the resolved runtime policy demands a usable GPU and a silent CPU
95/// fallback is forbidden.
96#[must_use]
97pub fn gpu_required_by_policy() -> bool {
98    gpu_runtime_policy() == GpuRuntimePolicy::Required
99}
100
101/// Require-GPU preflight, independent of backend routing.
102///
103/// When the policy is not [`GpuRuntimePolicy::Required`] this is a no-op and
104/// returns `Ok(())`. When it is required, the contract is to refuse to run when
105/// no usable GPU adapter is detected. This check fires on the no-GPU path the
106/// flag exists for; it does not depend on `select_backend` having chosen GPU
107/// first.
108///
109/// Returns `Err(diagnostic)` when no acquired CUDA or WGPU peer passes the
110/// production region-presence parity self-test. The caller maps that to the
111/// documented exit code 12. Returning an `Err` here - rather than calling
112/// `std::process::exit` from the library - keeps embedders alive (finding
113/// M12).
114pub fn require_gpu_preflight() -> Result<(), String> {
115    if !gpu_required_by_policy() {
116        return Ok(());
117    }
118
119    if let Err(reason) = super::gpu_region_presence_self_test() {
120        return Err(format!(
121            "--require-gpu requested but no complete production GPU peer set passed region-presence parity ({reason}); \
122             refusing to run on CPU. Fix the GPU stack or run without \
123             --require-gpu."
124        ));
125    }
126
127    Ok(())
128}
129
130pub(crate) fn gpu_disabled_by_policy() -> bool {
131    matches!(gpu_runtime_policy(), GpuRuntimePolicy::Disabled)
132}