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

keyhog_scanner/gpu/
policy.rs
