keyhog_scanner/gpu/backend/
acquisition.rs

1//! Lazy CUDA, Metal, and WGPU acquisition behind the scanner GPU boundary.
2
3#[cfg(feature = "gpu")]
4use super::artifact::{load_moe_artifacts, MoeArtifacts};
5#[cfg(feature = "gpu")]
6use super::diagnostics::{on_gpu_init_failed, GpuBackendError, GpuInitError};
7use crate::hw_probe::ScanBackend;
8#[cfg(feature = "gpu")]
9use std::sync::LazyLock;
10use std::sync::{Arc, OnceLock};
11
12pub(crate) struct AcquiredGpuPeer {
13    pub(crate) backend: Arc<dyn vyre::VyreBackend>,
14    pub(crate) device_identity: Option<String>,
15    pub(crate) is_software: bool,
16    pub(crate) resident_timed_dispatch_supported: bool,
17}
18
19#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
20pub struct GpuBackendAvailability {
21    pub cuda: bool,
22    pub metal: bool,
23    pub wgpu: bool,
24}
25
26impl GpuBackendAvailability {
27    #[must_use]
28    pub const fn any(self) -> bool {
29        self.cuda || self.metal || self.wgpu
30    }
31}
32
33#[derive(Clone, Debug, Eq, PartialEq)]
34pub(crate) struct GpuBackendAcquisitionFailure {
35    pub backend: &'static str,
36    pub diagnostic: String,
37}
38
39pub(crate) struct GpuBackendPeers {
40    cuda: OnceLock<Result<AcquiredGpuPeer, String>>,
41    metal: OnceLock<Result<AcquiredGpuPeer, String>>,
42    wgpu: OnceLock<Result<AcquiredGpuPeer, String>>,
43    pub(crate) cuda_available: bool,
44    pub(crate) metal_available: bool,
45    pub(crate) wgpu_available: bool,
46    pub(crate) cuda_device_identity: Option<String>,
47    pub(crate) cuda_runtime_identity: Option<String>,
48    pub(crate) metal_device_identity: Option<String>,
49    pub(crate) metal_runtime_identity: Option<String>,
50    pub(crate) wgpu_device_identity: Option<String>,
51    pub(crate) wgpu_runtime_identity: Option<String>,
52    pub(crate) wgpu_is_software: bool,
53}
54
55impl Default for GpuBackendPeers {
56    fn default() -> Self {
57        Self {
58            cuda: OnceLock::new(),
59            metal: OnceLock::new(),
60            wgpu: OnceLock::new(),
61            cuda_available: false,
62            metal_available: false,
63            wgpu_available: false,
64            cuda_device_identity: None,
65            metal_device_identity: None,
66            metal_runtime_identity: None,
67            cuda_runtime_identity: None,
68            wgpu_device_identity: None,
69            wgpu_runtime_identity: None,
70            wgpu_is_software: false,
71        }
72    }
73}
74
75pub(super) fn lazy_acquire<T, E>(
76    available: bool,
77    slot: &OnceLock<Result<T, E>>,
78    acquire: impl FnOnce() -> Result<T, E>,
79) -> Option<&Result<T, E>> {
80    if !available {
81        return None;
82    }
83    Some(slot.get_or_init(acquire))
84}
85
86impl GpuBackendPeers {
87    pub(crate) fn get(&self, backend: ScanBackend) -> Option<&Arc<dyn vyre::VyreBackend>> {
88        let result = match backend {
89            ScanBackend::GpuCuda => {
90                lazy_acquire(self.cuda_available, &self.cuda, acquire_cuda_peer)
91            }
92            ScanBackend::GpuMetal => {
93                lazy_acquire(self.metal_available, &self.metal, acquire_metal_peer)
94            }
95            ScanBackend::GpuWgpu => {
96                lazy_acquire(self.wgpu_available, &self.wgpu, acquire_wgpu_peer)
97            }
98            _ => None,
99        }?;
100        match result {
101            Ok(peer) => Some(&peer.backend),
102            Err(error) => {
103                tracing::error!(
104                    target: "keyhog::routing",
105                    ?backend,
106                    diagnostic = %error,
107                    "selected GPU backend acquisition failed"
108                );
109                None
110            }
111        }
112    }
113
114    pub(crate) fn initialized(&self, backend: ScanBackend) -> Option<&AcquiredGpuPeer> {
115        let result = match backend {
116            ScanBackend::GpuCuda => self.cuda.get(),
117            ScanBackend::GpuMetal => self.metal.get(),
118            ScanBackend::GpuWgpu => self.wgpu.get(),
119            _ => None,
120        }?;
121        // LAW10: diagnostics retain the complete error in the slot; this status
122        // accessor distinguishes initialized success without consuming it.
123        match result {
124            Ok(peer) => Some(peer),
125            Err(_) => None, // LAW10: status projection only; initialization_error retains the typed diagnostic and execution logs it before refusing this backend.
126        }
127    }
128    pub(crate) fn resident_timed_dispatch_supported(&self, backend: ScanBackend) -> bool {
129        self.initialized(backend)
130            .is_some_and(|peer| peer.resident_timed_dispatch_supported)
131    }
132
133    pub(crate) fn initialization_error(&self, backend: ScanBackend) -> Option<&str> {
134        match backend {
135            ScanBackend::GpuCuda => self.cuda.get(),
136            ScanBackend::GpuMetal => self.metal.get(),
137            ScanBackend::GpuWgpu => self.wgpu.get(),
138            _ => None,
139        }
140        .and_then(|result| result.as_ref().err().map(String::as_str))
141    }
142
143    pub(crate) fn availability(&self) -> GpuBackendAvailability {
144        GpuBackendAvailability {
145            cuda: self.cuda_available,
146            metal: self.metal_available,
147            wgpu: self.wgpu_available,
148        }
149    }
150}
151
152#[cfg(all(feature = "gpu", target_os = "linux"))]
153pub(super) fn run_cuda_after_preflight<T>(
154    preflight: impl FnOnce() -> Result<(), String>,
155    acquire: impl FnOnce() -> Result<T, String>,
156    operation: &'static str,
157) -> Result<T, String> {
158    preflight()?;
159    std::panic::catch_unwind(std::panic::AssertUnwindSafe(acquire)).map_err(|panic| {
160        format!(
161            "CUDA {operation} panicked: {}. Fix: repair the CUDA driver/runtime{}",
162            crate::error::panic_payload_detail(panic),
163            if operation == "backend acquisition" {
164                " or select another calibrated backend"
165            } else {
166                " before enabling this backend"
167            }
168        )
169    })?
170}
171
172#[cfg(all(feature = "gpu", target_os = "linux"))]
173fn acquire_cuda_peer() -> Result<AcquiredGpuPeer, String> {
174    let backend = run_cuda_after_preflight(
175        ensure_cuda_driver_library_loadable,
176        || {
177            let cuda = vyre_driver_cuda::backend::CudaBackend::acquire()?;
178            let boxed: Box<dyn vyre::VyreBackend> =
179                Box::new(vyre_driver_cuda::CudaBackendRegistration::new(cuda));
180            Ok::<Arc<dyn vyre::VyreBackend>, String>(Arc::from(boxed))
181        },
182        "backend acquisition",
183    )?;
184    tracing::info!(target: "keyhog::routing", "selected CUDA peer backend acquired");
185    Ok(AcquiredGpuPeer {
186        backend,
187        device_identity: None,
188        is_software: false,
189        resident_timed_dispatch_supported: true,
190    })
191}
192
193#[cfg(not(all(feature = "gpu", target_os = "linux")))]
194fn acquire_cuda_peer() -> Result<AcquiredGpuPeer, String> {
195    Err("CUDA peer is not compiled for this platform".to_string())
196}
197
198#[cfg(all(feature = "gpu", target_os = "macos"))]
199fn acquire_metal_peer() -> Result<AcquiredGpuPeer, String> {
200    let backend = std::panic::catch_unwind(std::panic::AssertUnwindSafe(
201        vyre_driver_metal::acquire,
202    ))
203    .map_err(|panic| {
204        format!(
205            "Metal backend acquisition panicked: {}. Fix: repair Metal.framework or select another calibrated backend",
206            crate::error::panic_payload_detail(panic)
207        )
208    })?
209    .map_err(|error| error.to_string())?;
210    tracing::info!(target: "keyhog::routing", "selected native Metal peer backend acquired");
211    Ok(AcquiredGpuPeer {
212        backend: Arc::from(backend),
213        device_identity: Some("Apple Metal default device".to_string()),
214        is_software: false,
215        resident_timed_dispatch_supported: false,
216    })
217}
218
219#[cfg(not(all(feature = "gpu", target_os = "macos")))]
220fn acquire_metal_peer() -> Result<AcquiredGpuPeer, String> {
221    Err("native Metal peer is not compiled for this platform".to_string())
222}
223
224#[cfg(feature = "gpu")]
225fn wgpu_resident_timed_dispatch_supported(features: wgpu::Features) -> bool {
226    features
227        .contains(wgpu::Features::TIMESTAMP_QUERY | wgpu::Features::TIMESTAMP_QUERY_INSIDE_ENCODERS)
228}
229
230#[cfg(feature = "gpu")]
231fn acquire_wgpu_peer() -> Result<AcquiredGpuPeer, String> {
232    let backend = std::panic::catch_unwind(std::panic::AssertUnwindSafe(
233        vyre_driver_wgpu::WgpuBackend::shared,
234    ))
235    .map_err(|panic| {
236        format!(
237            "WGPU backend acquisition panicked: {}. Fix: repair the graphics driver/runtime or select another calibrated backend",
238            crate::error::panic_payload_detail(panic)
239        )
240    })?
241    .map_err(|error| error.to_string())?;
242    let info = backend.adapter_info();
243    let device_identity =
244        crate::gpu::gpu_adapter_device_identity(info, backend.device_limits().max_buffer_size);
245    let is_software = crate::gpu::is_software_adapter(info);
246    let resident_timed_dispatch_supported =
247        wgpu_resident_timed_dispatch_supported(backend.device_queue().0.features());
248    tracing::info!(
249        target: "keyhog::routing",
250        device_identity,
251        "selected WGPU peer backend acquired"
252    );
253    let backend: Arc<dyn vyre::VyreBackend> = backend;
254    Ok(AcquiredGpuPeer {
255        backend,
256        device_identity: Some(device_identity),
257        is_software,
258        resident_timed_dispatch_supported,
259    })
260}
261
262#[cfg(not(feature = "gpu"))]
263fn acquire_wgpu_peer() -> Result<AcquiredGpuPeer, String> {
264    Err("WGPU peer is not compiled in this build".to_string())
265}
266
267#[cfg(all(feature = "gpu", target_os = "linux"))]
268pub(crate) fn load_dynamic_library(name: &std::ffi::CStr) -> Result<(), String> {
269    // SAFETY: `name` is NUL-terminated for `dlopen`; a successful handle is
270    // closed before returning, and `dlerror` remains valid until the next
271    // dynamic-loader call on this thread.
272    unsafe {
273        let handle = libc::dlopen(name.as_ptr(), libc::RTLD_NOW | libc::RTLD_LOCAL);
274        if handle.is_null() {
275            let error = libc::dlerror();
276            let detail = if error.is_null() {
277                "unknown dynamic-loader error".to_owned()
278            } else {
279                std::ffi::CStr::from_ptr(error)
280                    .to_string_lossy()
281                    .into_owned()
282            };
283            return Err(format!("{}: {detail}", name.to_string_lossy()));
284        }
285        libc::dlclose(handle);
286    }
287    Ok(())
288}
289
290#[cfg(all(feature = "gpu", target_os = "linux"))]
291fn ensure_cuda_driver_library_loadable() -> Result<(), String> {
292    let first_error = match load_dynamic_library(c"libcuda.so.1") {
293        Ok(()) => return Ok(()),
294        Err(error) => error,
295    };
296    let second_error = match load_dynamic_library(c"libcuda.so") {
297        Ok(()) => return Ok(()),
298        Err(error) => error,
299    };
300    Err(format!(
301        "CUDA driver library is unavailable ({second_error}; first attempt: {first_error}). Fix: install or expose the NVIDIA driver libcuda.so before enabling this backend"
302    ))
303}
304
305#[cfg(all(feature = "gpu", target_os = "linux"))]
306pub(crate) fn probe_cuda_peer() -> Result<vyre_driver_cuda::device::CudaDeviceCaps, String> {
307    run_cuda_after_preflight(
308        ensure_cuda_driver_library_loadable,
309        || vyre_driver_cuda::device::CudaDeviceCaps::probe(0).map_err(|error| error.to_string()),
310        "device probe",
311    )
312}
313
314#[cfg(feature = "gpu")]
315pub(crate) struct GpuContext {
316    pub(super) device_queue: Arc<(wgpu::Device, wgpu::Queue)>,
317    pub(super) adapter_info: wgpu::AdapterInfo,
318    pub(super) device_limits: wgpu::Limits,
319    pub(super) artifacts: MoeArtifacts,
320}
321
322#[cfg(feature = "gpu")]
323impl GpuContext {
324    pub(crate) fn vram_mb(&self) -> Option<u64> {
325        const SANE_CAP_MB: u64 = 256 * 1024;
326        Some((self.device_limits.max_buffer_size / (1024 * 1024)).min(SANE_CAP_MB))
327    }
328
329    pub(crate) fn gpu_name(&self) -> &str {
330        &self.adapter_info.name
331    }
332
333    pub(super) fn device(&self) -> &wgpu::Device {
334        &self.device_queue.0
335    }
336
337    pub(super) fn queue(&self) -> &wgpu::Queue {
338        &self.device_queue.1
339    }
340
341    pub(super) fn artifacts(&self) -> &MoeArtifacts {
342        &self.artifacts
343    }
344}
345
346#[cfg(feature = "gpu")]
347static GPU: LazyLock<Result<Option<GpuContext>, GpuBackendError>> =
348    LazyLock::new(|| match init_moe_gpu() {
349        Ok(context) => {
350            tracing::info!("GPU MoE inference initialized (shared device)");
351            Ok(Some(context))
352        }
353        Err(error) => {
354            on_gpu_init_failed(
355                &error,
356                crate::gpu::gpu_disabled_by_policy(),
357                crate::gpu::gpu_required_by_policy(),
358            )?;
359            Ok(None)
360        }
361    });
362
363#[cfg(feature = "gpu")]
364fn init_moe_gpu() -> Result<GpuContext, GpuInitError> {
365    let vyre_backend = vyre_driver_wgpu::WgpuBackend::shared().map_err(|error| {
366        GpuInitError::no_adapter(format!("vyre WgpuBackend unavailable: {error}"))
367    })?;
368    let adapter_info = vyre_backend.adapter_info().clone();
369    if crate::gpu::is_software_adapter(&adapter_info) {
370        return Err(GpuInitError::no_adapter(format!(
371            "GPU adapter is a software fallback ({} on {:?}); refusing to use",
372            adapter_info.name, adapter_info.backend
373        )));
374    }
375    let device_limits = vyre_backend.device_limits().clone();
376    let device_queue = vyre_backend.device_queue();
377    tracing::info!(
378        gpu = %adapter_info.name,
379        backend = ?adapter_info.backend,
380        device_type = ?adapter_info.device_type,
381        driver = %adapter_info.driver,
382        "GPU MoE: reusing vyre shared device"
383    );
384    let artifacts = load_moe_artifacts(&device_queue.0, &adapter_info, &device_limits)
385        .map_err(GpuInitError::adapter_unusable)?;
386    Ok(GpuContext {
387        device_queue,
388        adapter_info,
389        device_limits,
390        artifacts,
391    })
392}
393
394#[cfg(feature = "gpu")]
395pub(crate) fn get_gpu() -> Result<Option<&'static GpuContext>, GpuBackendError> {
396    match &*GPU {
397        Ok(context) => Ok(context.as_ref()),
398        Err(error) => Err(error.clone()),
399    }
400}
401
402#[cfg(all(test, feature = "gpu"))]
403#[path = "../../../tests/unit/gpu_backend_acquisition.rs"]
404mod tests;
keyhog_scanner/gpu/backend/acquisition.rs

keyhog_scanner/gpu/backend/
acquisition.rs
