1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
use super::*;
static GPU_AC_DEGENERATE_DISABLED: std::sync::atomic::AtomicBool =
std::sync::atomic::AtomicBool::new(false);
impl CompiledScanner {
pub fn scan_coalesced_gpu_ac_phase1(&self, chunks: &[keyhog_core::Chunk]) -> GpuPhase1Output {
let Some(matcher) = self.gpu_matcher() else {
return self.gpu_degrade_done_with_reason(
chunks,
crate::hw_probe::ScanBackend::Gpu,
Some("GPU literal matcher unavailable for AC dispatch"),
);
};
let Some(program) = self.ac_gpu_program() else {
return self.gpu_degrade_done_with_reason(
chunks,
crate::hw_probe::ScanBackend::Gpu,
Some("GPU AC dispatch program unavailable"),
);
};
if self.gpu_backend.is_none() {
return self.gpu_degrade_done_with_reason(
chunks,
crate::hw_probe::ScanBackend::Gpu,
Some("GPU backend handle unavailable for AC dispatch"),
);
}
if GPU_AC_DEGENERATE_DISABLED.load(std::sync::atomic::Ordering::Relaxed) {
return self.gpu_degrade_done_with_reason(
chunks,
crate::hw_probe::ScanBackend::Gpu,
Some("GPU AC previously emitted degenerate match triples (end <= start); skipping known-corrupt Vyre dispatch"),
);
}
let (entries, mut buffer) = super::gpu_coalesce::coalesce_chunks(chunks);
// ASCII-lowercase the coalesced haystack so the AC literal automaton
// matches case-INSENSITIVELY, exactly like the SIMD Hyperscan path
// (compiled CASELESS for every pattern). Without this the GPU drops
// matches on uppercase occurrences of lowercase literal prefixes
// (PERF-07 gpu_parity: `csb_` literal vs `CSB_...` in soc21_enum.h ->
// SIMD 4, GPU 0). The literal set is lowercased to the same fold in
// `build_gpu_literals`. This buffer is the phase-1 PREFILTER only -
// phase 2 re-confirms each hit on the ORIGINAL chunk bytes with the
// caseless regex - and ASCII fold is 1-byte-to-1-byte (only A-Z), so
// the match offsets attributed back to chunks are unchanged and the
// reported credential keeps its original case.
buffer.make_ascii_lowercase();
// Same buffer 4-alignment trick as `scan_coalesced_gpu`: lets
// every shard pass `&buffer[start..end]` straight to vyre's
// u32-typed haystack input instead of running pack_haystack_u32
// (a 2x memcopy producing byte-identical output for aligned
// slices). Eliminates ~2x buffer.len() of transient allocations
// per scan. NUL padding is recall-safe (literals can't contain
// NUL).
while !buffer.len().is_multiple_of(4) {
buffer.push(0);
}
#[cfg(target_os = "linux")]
// SAFETY: same contract as scan_coalesced_gpu - `buffer` is a
// live owned Vec describing a valid range; madvise is advisory.
unsafe {
libc::madvise(
buffer.as_ptr() as *mut libc::c_void,
buffer.len(),
libc::MADV_DONTDUMP,
);
}
let workgroup_x = program.workgroup_size[0] as usize;
// WGSL workgroups-per-dim ceiling is 65 535. At workgroup_x = 64
// that's a ~4 MiB shard. The shard cap is here so we never feed
// the dispatch a workgroup count > 65 535 (validation error).
const GPU_DISPATCH_MAX_WORKGROUPS_AC: usize = 65_535;
let gpu_dispatch_max_bytes: usize = GPU_DISPATCH_MAX_WORKGROUPS_AC * workgroup_x;
let started = std::time::Instant::now();
let mut shard_ranges: Vec<(usize, usize)> = Vec::new();
let mut shard_start = 0usize;
while shard_start < buffer.len() {
let shard_end = (shard_start + gpu_dispatch_max_bytes).min(buffer.len());
shard_ranges.push((shard_start, shard_end));
shard_start = shard_end;
}
let shard_count = shard_ranges.len();
// Constants packed ONCE per process via the scanner-level
// OnceLock. Same rationale as `scan_coalesced_gpu`: AC kernel
// re-ran four `pack_u32_slice` calls on identical bytes every
// dispatch.
// The AC program's binding layout:
// 0: haystack (per shard, slice into padded buffer)
// 1: transitions
// 2: output_offsets
// 3: output_records
// 4: pattern_lengths
// 5: haystack_len (per shard, packed)
// 6: match_count (per shard, atomic counter)
// 7: matches (output, backend-allocated from BufferDecl)
let ac_packs = self
.gpu_ac_const_packs
.get_or_init(|| super::gpu_cache::AcConstPacks {
transitions: vyre_libs::scan::dispatch_io::pack_u32_slice(&matcher.dfa.transitions),
output_offsets: vyre_libs::scan::dispatch_io::pack_u32_slice(
&matcher.dfa.output_offsets,
),
output_records: vyre_libs::scan::dispatch_io::pack_u32_slice(
&matcher.dfa.output_records,
),
pattern_lengths: vyre_libs::scan::dispatch_io::pack_u32_slice(
&matcher.pattern_lengths,
),
});
struct ShardOwnedAc {
haystack_len: Vec<u8>,
atomic_count: Vec<u8>,
config: vyre::DispatchConfig,
}
let mut shard_owned: Vec<ShardOwnedAc> = Vec::with_capacity(shard_count);
for &(s_start, s_end) in &shard_ranges {
let shard_len = (s_end - s_start) as u32;
shard_owned.push(ShardOwnedAc {
haystack_len: vyre_libs::scan::dispatch_io::pack_u32_slice(&[shard_len]),
atomic_count: vec![0u8; 4],
config: vyre_libs::scan::dispatch_io::byte_scan_dispatch_config(
shard_len,
program.workgroup_size[0],
),
});
}
let shard_input_arrays: Vec<[&[u8]; 7]> = shard_owned
.iter()
.zip(shard_ranges.iter())
.map(|(s, &(start, end))| {
[
&buffer[start..end],
ac_packs.transitions.as_slice(),
ac_packs.output_offsets.as_slice(),
ac_packs.output_records.as_slice(),
ac_packs.pattern_lengths.as_slice(),
s.haystack_len.as_slice(),
s.atomic_count.as_slice(),
]
})
.collect();
// Sub-batched dispatch: dynamically scaled MAX_SHARDS_PER_GPU_BATCH
// budget based on system RAM keeps transient host-side memory
// bounded while maximizing dispatch concurrency for high-tier GPUs
// and leaving vyre's 2048-slot readback ring deeply under-subscribed.
let max_shards_per_gpu_batch: usize = {
let total_ram_mb = crate::hw_probe::probe_hardware()
.total_memory_mb
.unwrap_or(0);
if total_ram_mb >= 32 * 1024 {
256
} else if total_ram_mb >= 16 * 1024 {
128
} else {
64
}
};
let mut matches: Vec<vyre_libs::scan::LiteralMatch> = Vec::new();
for sub_start in (0..shard_count).step_by(max_shards_per_gpu_batch) {
let sub_end = (sub_start + max_shards_per_gpu_batch).min(shard_count);
let sub_inputs: Vec<&[&[u8]]> = (sub_start..sub_end)
.map(|i| &shard_input_arrays[i][..])
.collect();
let sub_configs: Vec<vyre::DispatchConfig> = (sub_start..sub_end)
.map(|i| shard_owned[i].config.clone())
.collect();
let batch_results = match self.dispatch_gpu_shards(program, &sub_inputs, &sub_configs) {
Ok(r) => r,
Err(e) => {
tracing::error!(
shards = sub_end - sub_start,
"AC GPU batched dispatch failed, falling back to CPU: {e}"
);
let reason = format!("AC GPU batched dispatch failed: {e}");
return self.gpu_degrade_done_with_reason(
chunks,
crate::hw_probe::ScanBackend::Gpu,
Some(&reason),
);
}
};
for (offset_in_sub, result) in batch_results.into_iter().enumerate() {
let i = sub_start + offset_in_sub;
let outputs = match result {
Ok(o) => o,
Err(e) => {
tracing::error!(
shard_index = i,
"AC GPU shard within batch failed, falling back to CPU: {e}"
);
let reason = format!("AC GPU shard {i} dispatch failed: {e}");
return self.gpu_degrade_done_with_reason(
chunks,
crate::hw_probe::ScanBackend::Gpu,
Some(&reason),
);
}
};
if outputs.len() < 2 {
tracing::error!(
shard_index = i,
outputs = outputs.len(),
"AC GPU shard output buffer count too small; falling back to CPU"
);
let reason = format!(
"AC GPU shard {i} returned {} output buffer(s), expected at least 2",
outputs.len()
);
return self.gpu_degrade_done_with_reason(
chunks,
crate::hw_probe::ScanBackend::Gpu,
Some(&reason),
);
}
let count_bytes = &outputs[0];
let matches_bytes = &outputs[1];
if count_bytes.len() < 4 {
tracing::error!(
shard_index = i,
"AC GPU shard count buffer truncated; falling back to CPU"
);
let reason = format!(
"AC GPU shard {i} returned truncated count buffer ({} byte(s), expected 4)",
count_bytes.len()
);
return self.gpu_degrade_done_with_reason(
chunks,
crate::hw_probe::ScanBackend::Gpu,
Some(&reason),
);
}
let count = u32::from_le_bytes([
count_bytes[0],
count_bytes[1],
count_bytes[2],
count_bytes[3],
]);
if count > super::rule_pipeline::AC_GPU_MAX_MATCHES_PER_DISPATCH {
tracing::warn!(
cap = super::rule_pipeline::AC_GPU_MAX_MATCHES_PER_DISPATCH,
count,
shard_index = i,
"AC GPU shard exceeded dense-prefix cap; rerouting batch through SIMD coalesced scan"
);
if self.has_simd_prefilter() {
if std::env::var_os("KH_PERF").is_some() {
eprintln!(
"KH_PERF gpu_ac_cap_reroute: chunks={} shard={} shard_matches={} cap={} shard_bytes={}",
chunks.len(),
i,
count,
super::rule_pipeline::AC_GPU_MAX_MATCHES_PER_DISPATCH,
shard_ranges[i].1 - shard_ranges[i].0
);
}
return GpuPhase1Output::Done(self.scan_coalesced_non_gpu(chunks));
}
let reason = format!(
"AC GPU shard {i} reported {count} matches, exceeding dense-prefix cap {} and no SIMD fallback is available",
super::rule_pipeline::AC_GPU_MAX_MATCHES_PER_DISPATCH
);
return self.gpu_degrade_done_with_reason(
chunks,
crate::hw_probe::ScanBackend::Gpu,
Some(&reason),
);
}
let shard_matches = vyre_libs::scan::dispatch_io::unpack_match_triples(
matches_bytes,
count.min(super::rule_pipeline::AC_GPU_MAX_MATCHES_PER_DISPATCH),
);
let offset = shard_ranges[i].0 as u32;
for m in &shard_matches {
matches.push(vyre_libs::scan::LiteralMatch::new(
m.pattern_id,
m.start.saturating_add(offset),
m.end.saturating_add(offset),
));
}
}
}
let elapsed_ms = started.elapsed().as_millis();
tracing::debug!(
target: "keyhog::routing",
chunks = chunks.len(),
buffer_bytes = buffer.len(),
matches = matches.len(),
shards = shard_count,
elapsed_ms,
"AC GPU batched scan completed"
);
// PERF-07c correctness guard: a sound AC kernel emits `end = i + 1`
// and `start = end - pat_len` with `pat_len >= 1`, so EVERY real match
// has `end > start`. A triple with `end <= start` (observed: a flood of
// degenerate `(pid=0, start=0, end=0)`) is impossible from correct
// output. The vyre CUDA PTX emit path currently produces such triples;
// folded to `(0,0)` they mis-attribute every PID to chunk 0 of a
// coalesced batch, silently dropping real hits in chunks > 0 - a
// fail-OPEN recall gap that only manifests on multi-file batches
// (single-file scans put the target in chunk 0 and mask it). Until the
// emitter is fixed (tracked as the vyre GPU upgrade), detect the
// corruption and degrade THIS batch to the SIMD/CPU literal path, which
// is correct and - measured on the kernel - actually faster than the
// GPU AC path here. The GPU MoE scorer still runs in phase 2. This is
// self-validating: a backend that emits sound triples (zero degenerate)
// never degrades, so the guard auto-clears once vyre's CUDA emit is
// fixed, with no keyhog change required.
if matches.iter().any(|m| m.end <= m.start) {
GPU_AC_DEGENERATE_DISABLED.store(true, std::sync::atomic::Ordering::Relaxed);
tracing::warn!(
target: "keyhog::routing",
raw_matches = matches.len(),
chunks = chunks.len(),
"GPU AC emitted degenerate match triples (end <= start); vyre CUDA \
emit bug PERF-07c. Degrading this batch to the SIMD/CPU literal \
path to preserve recall parity."
);
return self.gpu_degrade_done_with_reason(
chunks,
crate::hw_probe::ScanBackend::Gpu,
Some("GPU AC emitted degenerate match triples (end <= start); vyre CUDA emit bug PERF-07c"),
);
}
if self.has_simd_prefilter()
&& super::gpu_postprocess::gpu_phase2_hits_are_dense(
matches.len(),
buffer.len(),
chunks.len(),
)
{
tracing::warn!(
target: "keyhog::routing",
raw_matches = matches.len(),
buffer_bytes = buffer.len(),
chunks = chunks.len(),
"GPU AC prefix output is too dense for phase 2; rerouting this batch through SIMD coalesced scan",
);
if std::env::var_os("KH_PERF").is_some() {
eprintln!(
"KH_PERF gpu_ac_dense_phase2_reroute: chunks={} buffer_bytes={} raw_matches={} bytes_per_hit={:.1}",
chunks.len(),
buffer.len(),
matches.len(),
buffer.len() as f64 / matches.len().max(1) as f64
);
}
return GpuPhase1Output::Done(self.scan_coalesced_non_gpu(chunks));
}
super::gpu_postprocess::fold_overlapping_same_pid_inplace(&mut matches);
let total_patterns = self.ac_map.len() + self.fallback.len();
let per_chunk_hits = super::gpu_postprocess::attribute_matches_to_chunks(
&matches,
&entries,
total_patterns,
chunks.len(),
);
// Hand the hits back to the orchestrator so it can run phase 2
// on a separate thread (pipelined). Combined-wrapper callers
// (`scan_coalesced_gpu_ac`) call phase 2 inline immediately
// after this returns, preserving the original synchronous
// behaviour.
GpuPhase1Output::Hits(per_chunk_hits)
}
}