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