1#[derive(Debug, Clone, PartialEq, Eq)]
3pub struct GpuSelfTest {
4 pub adapter_name: String,
6 pub vram_mb: Option<u64>,
8 pub scores: usize,
10}
11
12#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct VyreGpuSelfTest {
15 pub direct_matches: usize,
17 pub coalesced_matches: usize,
19}
20
21#[cfg(feature = "gpu")]
22static GPU_SELF_TEST_CACHE: std::sync::OnceLock<std::result::Result<GpuSelfTest, String>> =
23 std::sync::OnceLock::new();
24
25pub fn gpu_self_test() -> Result<GpuSelfTest, String> {
31 #[cfg(not(feature = "gpu"))]
32 {
33 Err(
34 "GPU support not compiled in (lean ci build). Rebuild with `--features gpu` \
35 (or the default profile) to exercise the wgpu/CUDA path."
36 .to_string(),
37 )
38 }
39 #[cfg(feature = "gpu")]
40 {
41 GPU_SELF_TEST_CACHE
42 .get_or_init(|| {
43 let gpu = super::backend::get_gpu().ok_or_else(|| {
44 "GPU adapter unavailable; install or enable a non-software GPU adapter and driver"
45 .to_string()
46 })?;
47
48 let max_abs = super::backend::gpu_moe_parity_max_divergence(
61 std::time::Duration::from_millis(
62 crate::scanner_config::ScannerTuningConfig::GPU_MOE_TIMEOUT_MS_DEFAULT,
63 ),
64 )?;
65 if max_abs > super::backend::GPU_MOE_PARITY_TOLERANCE {
66 return Err(format!(
67 "GPU MoE compute shader diverges from the CPU MoE reference by {max_abs:.4} \
68 (tolerance {:.4}); the GPU would score findings differently from the \
69 CPU/SIMD path. Indicates a shader miscompile, weights-packing mismatch, \
70 or driver bug. Scans record the GPU MoE degrade and use the CPU MoE \
71 (correct + deterministic), \
72 so detection is unaffected, but GPU ML acceleration is OFF on this host.",
73 super::backend::GPU_MOE_PARITY_TOLERANCE
74 ));
75 }
76
77 Ok(GpuSelfTest {
78 adapter_name: gpu.gpu_name().to_string(),
79 vram_mb: gpu.vram_mb(),
80 scores: crate::ml_scorer::GPU_BATCH_THRESHOLD,
81 })
82 })
83 .clone()
84 }
85}
86
87pub fn vyre_gpu_self_test() -> Result<VyreGpuSelfTest, String> {
92 #[cfg(not(feature = "gpu"))]
93 {
94 Err(
95 "VYRE GPU self-test not available in the lean CI build (no WGPU driver compiled in). \
96 Rebuild with `--features gpu`."
97 .to_string(),
98 )
99 }
100 #[cfg(feature = "gpu")]
101 {
102 vyre_gpu_self_test_impl()
103 }
104}
105
106#[cfg(feature = "gpu")]
107fn vyre_gpu_self_test_impl() -> Result<VyreGpuSelfTest, String> {
108 use vyre_driver_wgpu::WgpuBackend;
109 use vyre_libs::scan::GpuLiteralSet;
110
111 let patterns: Vec<Vec<u8>> = vec![b"needle".to_vec()];
112 let pattern_refs: Vec<&[u8]> = patterns.iter().map(Vec::as_slice).collect();
113
114 let backend = WgpuBackend::shared().map_err(|e| format!("failed to init wgpu backend: {e}"))?;
115 let scanner = GpuLiteralSet::compile(&pattern_refs);
116
117 let direct = scanner
118 .scan(backend.as_ref(), b"needle", 100)
119 .map_err(|error| format!("vyre direct GPU scan failed: {error}"))?;
120 if direct.len() != 1 || direct[0].pattern_id != 0 || direct[0].start != 0 {
121 return Err(format!(
122 "vyre direct GPU scan returned unexpected matches: {direct:?}"
123 ));
124 }
125
126 const COALESCED_ITEMS: usize = 100;
127 let items: Vec<Vec<u8>> = (0..COALESCED_ITEMS)
128 .map(|index| format!("id-{index:03}-needle").into_bytes())
129 .collect();
130 let mut buffer = Vec::with_capacity(items.iter().map(Vec::len).sum());
131 for item in &items {
132 buffer.extend_from_slice(item);
133 }
134
135 let coalesced = scanner
136 .scan(backend.as_ref(), &buffer, 10_000)
137 .map_err(|error| format!("vyre coalesced GPU scan failed: {error}"))?;
138 if coalesced.len() != COALESCED_ITEMS {
139 return Err(format!(
140 "vyre coalesced GPU scan returned {} matches, expected {COALESCED_ITEMS}",
141 coalesced.len()
142 ));
143 }
144
145 Ok(VyreGpuSelfTest {
146 direct_matches: direct.len(),
147 coalesced_matches: coalesced.len(),
148 })
149}
150
151#[derive(Debug, Clone, PartialEq, Eq)]
153pub struct GpuRegionPresencePeerSelfTest {
154 pub backend: crate::hw_probe::ScanBackend,
156 pub backend_id: &'static str,
158 pub matches: usize,
160}
161
162#[derive(Debug, Clone, PartialEq, Eq)]
164pub struct GpuRegionPresenceSelfTest {
165 pub peers: Vec<GpuRegionPresencePeerSelfTest>,
167}
168
169#[derive(Debug, Clone, PartialEq, Eq)]
171pub struct GpuRegionPresenceSelfTestFailure {
172 pub acquired_backends: Vec<crate::hw_probe::ScanBackend>,
174 pub message: String,
176}
177
178impl std::fmt::Display for GpuRegionPresenceSelfTestFailure {
179 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
180 formatter.write_str(&self.message)
181 }
182}
183
184impl std::error::Error for GpuRegionPresenceSelfTestFailure {}
185
186pub fn gpu_region_presence_self_test(
190) -> Result<GpuRegionPresenceSelfTest, GpuRegionPresenceSelfTestFailure> {
191 #[cfg(not(feature = "gpu"))]
192 {
193 Err(GpuRegionPresenceSelfTestFailure {
194 acquired_backends: Vec::new(),
195 message: "GPU region-presence self-test not available in the lean ci build. Rebuild with `--features gpu` to exercise the production GPU trigger path.".to_string(),
196 })
197 }
198 #[cfg(feature = "gpu")]
199 {
200 gpu_region_presence_self_test_impl()
201 }
202}
203
204#[cfg(feature = "gpu")]
205fn gpu_region_presence_self_test_impl(
206) -> Result<GpuRegionPresenceSelfTest, GpuRegionPresenceSelfTestFailure> {
207 use crate::engine::CompiledScanner;
208 use crate::hw_probe::ScanBackend;
209 use keyhog_core::{Chunk, ChunkMetadata, DetectorFile};
210
211 const PLANTED: &str = "KHGPUSELFTEST_A1b2C3d4E5f6";
221 let detector =
222 toml::from_str::<DetectorFile>(include_str!("../../data/gpu-self-test-detector.toml"))
223 .map(|file| file.detector)
224 .map_err(|error| GpuRegionPresenceSelfTestFailure {
225 acquired_backends: Vec::new(),
226 message: format!("bundled GPU self-test detector TOML is invalid: {error}"),
227 })?;
228
229 let scanner = CompiledScanner::compile(vec![detector]).map_err(|error| {
230 GpuRegionPresenceSelfTestFailure {
231 acquired_backends: Vec::new(),
232 message: format!("CompiledScanner::compile failed during self-test: {error}"),
233 }
234 })?;
235
236 let candidates = scanner.gpu_backend_candidates();
237 let acquired_backends: Vec<_> = candidates
238 .iter()
239 .filter(|candidate| candidate.is_eligible())
240 .map(|candidate| candidate.backend)
241 .collect();
242 if acquired_backends.is_empty() {
243 let diagnostics = candidates
244 .iter()
245 .map(|candidate| {
246 let diagnostic = match candidate.acquisition_error.as_deref() {
247 Some(reason) => reason,
248 None => "driver was not acquired and returned no diagnostic",
249 };
250 format!("{}: {diagnostic}", candidate.backend.label())
251 })
252 .collect::<Vec<_>>()
253 .join("; ");
254 return Err(GpuRegionPresenceSelfTestFailure {
255 acquired_backends,
256 message: format!("no GPU region-presence peer was acquired ({diagnostics})"),
257 });
258 }
259
260 let make_chunk = || Chunk {
261 data: format!("gpu_secret = {PLANTED}").into(),
262 metadata: ChunkMetadata::default(),
263 };
264
265 let cpu_results = scanner.scan_chunks_with_backend(&[make_chunk()], ScanBackend::CpuFallback);
269 let cpu_total: usize = cpu_results.iter().map(Vec::len).sum();
270 if cpu_total == 0 {
271 return Err(GpuRegionPresenceSelfTestFailure {
272 acquired_backends,
273 message: "GPU self-test probe matched on no backend (CPU baseline is zero); fix the self-test probe so it survives suppression.".to_string(),
274 });
275 }
276
277 let mut peers = Vec::with_capacity(acquired_backends.len());
278 let mut failures = Vec::new();
279 for candidate in candidates
280 .into_iter()
281 .filter(|candidate| candidate.is_eligible())
282 {
283 let route = candidate.backend;
284 let Some(backend_id) = candidate.driver_id else {
285 failures.push(format!(
286 "{}: acquired driver returned no identity",
287 route.label()
288 ));
289 continue;
290 };
291 let degrade_before = scanner.runtime_status().gpu_degrade_count;
292 let results = match scanner.try_scan_coalesced_gpu_region_presence(
293 &[make_chunk()],
294 route,
295 scanner.execution_route_for_backend(route),
296 ) {
297 Ok(results) => results,
298 Err(error) => {
299 failures.push(format!(
300 "{} ({backend_id}): dispatch failed: {error}",
301 route.label()
302 ));
303 continue;
304 }
305 };
306 if scanner.runtime_status().gpu_degrade_count > degrade_before {
307 let diagnostic = match scanner.last_gpu_degrade_reason() {
308 Some(reason) => reason,
309 None => "runtime degrade recorded without a diagnostic".to_owned(),
310 };
311 failures.push(format!("{} ({backend_id}): {diagnostic}", route.label()));
312 continue;
313 }
314 let total: usize = results.iter().map(Vec::len).sum();
315 if total != cpu_total {
316 failures.push(format!(
317 "{} ({backend_id}): found {total} match(es), CPU found {cpu_total}",
318 route.label()
319 ));
320 continue;
321 }
322 peers.push(GpuRegionPresencePeerSelfTest {
323 backend: route,
324 backend_id,
325 matches: total,
326 });
327 }
328 if !failures.is_empty() {
329 let passed = peers
330 .iter()
331 .map(|peer| format!("{} ({})", peer.backend.label(), peer.backend_id))
332 .collect::<Vec<_>>()
333 .join(", ");
334 let passed = if passed.is_empty() {
335 "none".to_string()
336 } else {
337 passed
338 };
339 return Err(GpuRegionPresenceSelfTestFailure {
340 acquired_backends,
341 message: format!(
342 "GPU region-presence peer parity failed: {}; passed peers: {passed}",
343 failures.join("; ")
344 ),
345 });
346 }
347 Ok(GpuRegionPresenceSelfTest { peers })
348}