Skip to main content

vyre_driver_wgpu/megakernel/
dispatcher.rs

1//! Batched megakernel dispatch built on a persistent device work queue.
2
3use super::batch::{
4    persistent_storage_binding_usage, queue_state_word, CombinedBatch, FileBatch, HitRecord,
5    FILE_METADATA_WORDS, HIT_RECORD_WORDS, QUEUE_STATE_WORDS,
6};
7use super::dispatch_plan::{BatchDispatchPlan, BatchDispatchPlanCache, BatchDispatchPlanLookup};
8use super::pipeline_cache::{BatchPipelineCache, BatchPipelineShape};
9use super::segmentation::SEGMENT_WORDS;
10use crate::buffer::GpuBufferHandle;
11use crate::{pipeline::WgpuPipeline, WgpuBackend};
12use std::sync::Arc;
13use std::time::{Duration, Instant};
14use vyre_driver::{CompiledPipeline, DispatchConfig, VyreBackend};
15use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
16use vyre_runtime::megakernel::advanced::hierarchical_atomics::record_hit_to_ring_hierarchical;
17use vyre_runtime::megakernel::ir_util::atomic_load_relaxed;
18use vyre_runtime::megakernel::rule_catalog::{
19    accepted_rule_fingerprints_and_rejections_into, pack_rule_catalog_into, BatchRuleProgram,
20    BatchRuleRejection, RuleCatalogPackingScratch, RULE_META_WORDS,
21};
22use vyre_runtime::megakernel::scaling::{
23    MegakernelLaunchPolicy, MegakernelLaunchRecommendation, MegakernelLaunchRequest,
24};
25use vyre_runtime::megakernel::MegakernelDispatchTopology;
26use vyre_runtime::PipelineError;
27
28/// Schema version for WGPU scan batch segmentation evidence.
29pub const WGPU_SCAN_BATCH_SEGMENTATION_SCHEMA_VERSION: u32 = 1;
30
31/// Input counters for WGPU scan batch segmentation evidence.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
33pub struct WgpuScanBatchSegmentationRequest {
34    /// Logical scan chunks in the batch.
35    pub chunk_count: u32,
36    /// Maximum chunks recorded into one command encoder.
37    pub max_chunks_per_command_encoder: u32,
38    /// Bind groups reused across command encoders.
39    pub bind_group_reuse_count: u32,
40    /// Bind groups created for command encoders.
41    pub bind_group_create_count: u32,
42    /// Host-to-device copy commands recorded for the batch.
43    pub upload_copy_count: u32,
44    /// Device-to-host or device-to-staging copy commands recorded for the batch.
45    pub readback_copy_count: u32,
46    /// CPU oracle or backend-independent match digest.
47    pub expected_match_digest: u64,
48    /// WGPU segmented batch match digest.
49    pub actual_match_digest: u64,
50}
51
52impl WgpuScanBatchSegmentationRequest {
53    /// Construct WGPU scan batch segmentation counters.
54    #[must_use]
55    pub const fn new(
56        chunk_count: u32,
57        max_chunks_per_command_encoder: u32,
58        bind_group_reuse_count: u32,
59        bind_group_create_count: u32,
60        upload_copy_count: u32,
61        readback_copy_count: u32,
62        expected_match_digest: u64,
63        actual_match_digest: u64,
64    ) -> Self {
65        Self {
66            chunk_count,
67            max_chunks_per_command_encoder,
68            bind_group_reuse_count,
69            bind_group_create_count,
70            upload_copy_count,
71            readback_copy_count,
72            expected_match_digest,
73            actual_match_digest,
74        }
75    }
76}
77
78/// Evidence emitted for one WGPU segmented scan batch.
79#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80pub struct WgpuScanBatchSegmentationEvidence {
81    /// Evidence schema version.
82    pub schema_version: u32,
83    /// Logical scan chunks in the batch.
84    pub chunk_count: u32,
85    /// Segment count after applying the command encoder chunk limit.
86    pub segment_count: u32,
87    /// Command encoders required by the segmentation plan.
88    pub command_encoder_count: u32,
89    /// Bind groups reused across command encoders.
90    pub bind_group_reuse_count: u32,
91    /// Bind groups created for command encoders.
92    pub bind_group_create_count: u32,
93    /// Bind group reuse ratio in basis points.
94    pub bind_group_reuse_bps: u16,
95    /// Host-to-device copy commands recorded for the batch.
96    pub upload_copy_count: u32,
97    /// Device-to-host or device-to-staging copy commands recorded for the batch.
98    pub readback_copy_count: u32,
99    /// Total copy commands recorded for the batch.
100    pub copy_count: u32,
101    /// Stable match digest when WGPU output matches the oracle.
102    pub match_digest: u64,
103    /// True when expected and actual match digests are identical.
104    pub match_parity: bool,
105    /// True when command encoder, bind group, and copy counts are present.
106    pub all_command_counts_recorded: bool,
107}
108
109impl WgpuScanBatchSegmentationEvidence {
110    /// Return true when evidence has the schema, command counts, and match
111    /// parity required by release benchmark claims.
112    ///
113    /// `match_digest != 0` is intentionally NOT used as a completeness gate.
114    /// Zero is a legitimate digest value when the scanned corpus produces zero
115    /// rule firings (the hash of the empty match set); the `match_parity` flag
116    /// already encodes whether the oracle and WGPU digests agreed.
117    #[must_use]
118    pub const fn is_complete(self) -> bool {
119        self.schema_version == WGPU_SCAN_BATCH_SEGMENTATION_SCHEMA_VERSION
120            && self.chunk_count != 0
121            && self.segment_count != 0
122            && self.command_encoder_count == self.segment_count
123            && self.copy_count == self.upload_copy_count + self.readback_copy_count
124            && self.match_parity
125            && self.all_command_counts_recorded
126    }
127}
128
129/// WGPU scan batch segmentation evidence error.
130#[derive(Debug, Clone, PartialEq, Eq)]
131#[non_exhaustive]
132pub enum WgpuScanBatchSegmentationError {
133    /// The batch contains no chunks.
134    EmptyBatch,
135    /// The command encoder segmentation limit is zero.
136    ZeroChunksPerCommandEncoder,
137    /// Bind group counts do not account for every command encoder.
138    BindGroupCountMismatch {
139        /// Command encoders produced by segmentation.
140        command_encoder_count: u32,
141        /// Bind groups reused across command encoders.
142        bind_group_reuse_count: u32,
143        /// Bind groups created for command encoders.
144        bind_group_create_count: u32,
145    },
146    /// Copy count overflowed the evidence ABI.
147    CopyCountOverflow,
148    /// Match digest was absent (both digests were zero).
149    ///
150    /// **Deprecated**: `wgpu_scan_batch_segmentation_evidence` no longer emits
151    /// this variant.  Zero is a legitimate digest value for a corpus that fires
152    /// zero rules; matched zero digests are valid evidence.  This variant is
153    /// retained for ABI compatibility with existing match arms.
154    #[deprecated(
155        note = "wgpu_scan_batch_segmentation_evidence no longer rejects zero digests; \
156                zero is a valid digest for a corpus with no matches"
157    )]
158    ZeroMatchDigest,
159    /// WGPU output digest diverged from the oracle.
160    MatchDigestMismatch {
161        /// CPU oracle or backend-independent match digest.
162        expected_match_digest: u64,
163        /// WGPU segmented batch match digest.
164        actual_match_digest: u64,
165    },
166}
167
168// `ZeroMatchDigest` is deprecated but the Display impl still needs to handle it
169// for any external code that constructs the variant or receives it via FFI.
170#[allow(deprecated)]
171impl std::fmt::Display for WgpuScanBatchSegmentationError {
172    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
173        match self {
174            Self::EmptyBatch => formatter.write_str(
175                "WGPU scan batch has zero chunks. Fix: publish at least one scan chunk before recording segmentation evidence.",
176            ),
177            Self::ZeroChunksPerCommandEncoder => formatter.write_str(
178                "WGPU scan batch has zero chunks per command encoder. Fix: configure a positive segmentation limit.",
179            ),
180            Self::BindGroupCountMismatch {
181                command_encoder_count,
182                bind_group_reuse_count,
183                bind_group_create_count,
184            } => write!(
185                formatter,
186                "WGPU scan batch bind group counts reuse={bind_group_reuse_count} create={bind_group_create_count} do not account for {command_encoder_count} command encoder(s). Fix: record one reused or created bind group per segment."
187            ),
188            Self::CopyCountOverflow => formatter.write_str(
189                "WGPU scan batch copy count overflowed u32. Fix: shard the scan batch before recording evidence.",
190            ),
191            Self::ZeroMatchDigest => formatter.write_str(
192                "WGPU scan batch match digest is zero. Fix: compute the match digest before accepting segmentation evidence.",
193            ),
194            Self::MatchDigestMismatch {
195                expected_match_digest,
196                actual_match_digest,
197            } => write!(
198                formatter,
199                "WGPU scan batch match digest mismatch expected={expected_match_digest:#x} actual={actual_match_digest:#x}. Fix: reject the segmented batch or repair command/copy segmentation before reporting portable scan parity."
200            ),
201        }
202    }
203}
204
205impl std::error::Error for WgpuScanBatchSegmentationError {}
206
207/// Build WGPU scan batch segmentation evidence from recorded command counters.
208///
209/// Both `expected_match_digest` and `actual_match_digest` must be supplied by
210/// the caller before invoking this function.  `0` is a valid digest value for a
211/// corpus that fires zero rules; this function accepts equal-zero digests as
212/// legitimate parity evidence.
213///
214/// # Errors
215///
216/// Returns [`WgpuScanBatchSegmentationError`] when the batch is empty, the
217/// command counts are incomplete, copy counts overflow, or the expected and
218/// actual digests disagree.
219pub fn wgpu_scan_batch_segmentation_evidence(
220    request: WgpuScanBatchSegmentationRequest,
221) -> Result<WgpuScanBatchSegmentationEvidence, WgpuScanBatchSegmentationError> {
222    if request.chunk_count == 0 {
223        return Err(WgpuScanBatchSegmentationError::EmptyBatch);
224    }
225    if request.max_chunks_per_command_encoder == 0 {
226        return Err(WgpuScanBatchSegmentationError::ZeroChunksPerCommandEncoder);
227    }
228    // Zero is a legitimate digest when the scanned corpus fires zero rules (the
229    // hash of the empty match set).  Reject only when BOTH digests are zero AND
230    // chunk_count > 0 AND we cannot distinguish "not yet computed" from a genuine
231    // zero, the caller is responsible for computing both digests before
232    // submitting evidence.  The real guard is the equality check below: if
233    // expected != actual the caller's oracle disagreed with WGPU regardless of
234    // whether the value is zero.  The only residual sentinel case we reject is
235    // when EXACTLY ONE digest is zero and the other is non-zero, which would pass
236    // the equality check below only if the other were also zero (impossible).
237    // We therefore drop the unconditional zero-rejection and keep only the
238    // equality / parity check.
239    if request.expected_match_digest != request.actual_match_digest {
240        return Err(WgpuScanBatchSegmentationError::MatchDigestMismatch {
241            expected_match_digest: request.expected_match_digest,
242            actual_match_digest: request.actual_match_digest,
243        });
244    }
245
246    let command_encoder_count =
247        div_ceil_u32(request.chunk_count, request.max_chunks_per_command_encoder);
248    let bind_group_count = request
249        .bind_group_reuse_count
250        .checked_add(request.bind_group_create_count)
251        .ok_or(WgpuScanBatchSegmentationError::BindGroupCountMismatch {
252            command_encoder_count,
253            bind_group_reuse_count: request.bind_group_reuse_count,
254            bind_group_create_count: request.bind_group_create_count,
255        })?;
256    if bind_group_count != command_encoder_count {
257        return Err(WgpuScanBatchSegmentationError::BindGroupCountMismatch {
258            command_encoder_count,
259            bind_group_reuse_count: request.bind_group_reuse_count,
260            bind_group_create_count: request.bind_group_create_count,
261        });
262    }
263
264    let copy_count = request
265        .upload_copy_count
266        .checked_add(request.readback_copy_count)
267        .ok_or(WgpuScanBatchSegmentationError::CopyCountOverflow)?;
268    let bind_group_reuse_bps = ((u64::from(request.bind_group_reuse_count) * 10_000)
269        / u64::from(command_encoder_count)) as u16;
270
271    Ok(WgpuScanBatchSegmentationEvidence {
272        schema_version: WGPU_SCAN_BATCH_SEGMENTATION_SCHEMA_VERSION,
273        chunk_count: request.chunk_count,
274        segment_count: command_encoder_count,
275        command_encoder_count,
276        bind_group_reuse_count: request.bind_group_reuse_count,
277        bind_group_create_count: request.bind_group_create_count,
278        bind_group_reuse_bps,
279        upload_copy_count: request.upload_copy_count,
280        readback_copy_count: request.readback_copy_count,
281        copy_count,
282        match_digest: request.expected_match_digest,
283        match_parity: true,
284        all_command_counts_recorded: true,
285    })
286}
287
288const fn div_ceil_u32(numerator: u32, denominator: u32) -> u32 {
289    ((numerator as u64 + denominator as u64 - 1) / denominator as u64) as u32
290}
291
292/// Sparse hit-ring writer selected for the batched megakernel.
293#[derive(Debug, Clone, Copy, PartialEq, Eq)]
294#[non_exhaustive]
295pub enum BatchHitWriter {
296    /// Select hierarchical subgroup atomics when the backend advertises them,
297    /// otherwise use the scalar writer.
298    Auto,
299    /// One global atomic per hit. Universally supported but slower under high
300    /// hit density.
301    Scalar,
302    /// One global atomic per subgroup. Requires subgroup operations and fails
303    /// loudly if the backend cannot compile subgroup intrinsics.
304    HierarchicalSubgroup,
305}
306
307// NOTE: the `scan_batch_segmentation_tests` test module was relocated to the END
308// of this file. An inline test module here previously split the production source
309// that the source-shape tests inspect (they take everything before the first test
310// module), truncating it before the launch/dispatch lines they assert on. Keeping
311// all test modules at the end keeps that production-source view intact. (This note
312// deliberately avoids the literal test-config attribute so it does not re-trigger
313// that truncation.)
314
315impl BatchHitWriter {
316    /// Resolve this selection against backend subgroup capability.
317    ///
318    /// # Errors
319    ///
320    /// Returns [`PipelineError::Backend`] when subgroup atomics are explicitly
321    /// requested on a backend that does not report subgroup support.
322    pub fn resolve_for_backend(self, subgroup_supported: bool) -> Result<Self, PipelineError> {
323        match (self, subgroup_supported) {
324            (Self::Auto, true) => Ok(Self::HierarchicalSubgroup),
325            (Self::Auto, false) => Ok(Self::Scalar),
326            (Self::HierarchicalSubgroup, false) => Err(PipelineError::Backend(
327                "BatchHitWriter::HierarchicalSubgroup requires backend subgroup ops, but this backend reports supports_subgroup_ops=false. Fix: use BatchHitWriter::Auto/Scalar or run on a subgroup-capable adapter."
328                    .to_string(),
329            )),
330            (mode, _) => Ok(mode),
331        }
332    }
333}
334
335/// Immutable pipeline + launch geometry for batched megakernel scans.
336#[derive(Debug, Clone, PartialEq, Eq)]
337pub struct BatchDispatchConfig {
338    /// Worker lanes per workgroup.
339    pub workgroup_size_x: u32,
340    /// Number of workgroups to launch for each batch.
341    pub worker_groups: u32,
342    /// Maximum sparse hits retained in the output ring.
343    pub hit_capacity: u32,
344    /// Per-dispatch timeout budget.
345    pub timeout: Duration,
346    /// Optional graph-node count hint for topology selection.
347    pub graph_node_count: u32,
348    /// Optional graph-edge count hint for topology selection.
349    pub graph_edge_count: u32,
350    /// Optional active-frontier density in basis points.
351    pub frontier_density_bps: u16,
352    /// Optional memory-pressure estimate in basis points.
353    pub memory_pressure_bps: u16,
354    /// Additional device-resident bytes already committed for this dispatch family.
355    ///
356    /// The dispatcher adds its fixed queue-state resident footprint when building
357    /// the shared launch-policy request.
358    pub resident_device_bytes: u64,
359    /// Hard device-memory budget for policy planning. Zero means unbounded.
360    pub device_memory_budget_bytes: u64,
361    /// Hot opcode count observed by the caller or runtime telemetry.
362    pub hot_opcode_count: u32,
363    /// Hot window count observed by the caller or runtime telemetry.
364    pub hot_window_count: u32,
365    /// Requeued continuation count observed by the caller or runtime telemetry.
366    pub requeue_count: u64,
367    /// Maximum priority age observed by the caller or runtime telemetry.
368    pub max_priority_age: u32,
369}
370
371impl Default for BatchDispatchConfig {
372    fn default() -> Self {
373        Self {
374            workgroup_size_x: 64,
375            // `0` is a sentinel meaning "compute from adapter occupancy at
376            // dispatcher construction time".  Explicit non-zero values are
377            // preserved so callers who set `worker_groups` by hand are not
378            // overridden.
379            worker_groups: 0,
380            hit_capacity: 65_536,
381            timeout: Duration::from_secs(30),
382            graph_node_count: 0,
383            graph_edge_count: 0,
384            frontier_density_bps: 0,
385            memory_pressure_bps: 0,
386            resident_device_bytes: 0,
387            device_memory_budget_bytes: 0,
388            hot_opcode_count: 0,
389            hot_window_count: 0,
390            requeue_count: 0,
391            max_priority_age: 0,
392        }
393    }
394}
395
396impl BatchDispatchConfig {
397    /// Attach graph-topology hints used by the shared megakernel policy.
398    #[must_use]
399    pub const fn with_graph_hints(
400        mut self,
401        graph_node_count: u32,
402        graph_edge_count: u32,
403        frontier_density_bps: u16,
404        memory_pressure_bps: u16,
405    ) -> Self {
406        self.graph_node_count = graph_node_count;
407        self.graph_edge_count = graph_edge_count;
408        self.frontier_density_bps = if frontier_density_bps > 10_000 {
409            10_000
410        } else {
411            frontier_density_bps
412        };
413        self.memory_pressure_bps = if memory_pressure_bps > 10_000 {
414            10_000
415        } else {
416            memory_pressure_bps
417        };
418        self
419    }
420
421    /// Attach hard device-memory budget hints used by the shared launch policy.
422    #[must_use]
423    pub const fn with_device_memory_budget(
424        mut self,
425        resident_device_bytes: u64,
426        device_memory_budget_bytes: u64,
427    ) -> Self {
428        self.resident_device_bytes = resident_device_bytes;
429        self.device_memory_budget_bytes = device_memory_budget_bytes;
430        self
431    }
432
433    /// Attach execution hotness hints used by interpreter/JIT routing.
434    #[must_use]
435    pub const fn with_execution_hints(
436        mut self,
437        hot_opcode_count: u32,
438        hot_window_count: u32,
439        requeue_count: u64,
440        max_priority_age: u32,
441    ) -> Self {
442        self.hot_opcode_count = hot_opcode_count;
443        self.hot_window_count = hot_window_count;
444        self.requeue_count = requeue_count;
445        self.max_priority_age = max_priority_age;
446        self
447    }
448
449    /// Return the shared launch-policy recommendation for this batch shape.
450    ///
451    /// # Errors
452    ///
453    /// Returns [`PipelineError::Backend`] when adapter limits are malformed.
454    pub fn launch_recommendation(
455        &self,
456        limits: &wgpu::Limits,
457        queue_len: u32,
458    ) -> Result<MegakernelLaunchRecommendation, PipelineError> {
459        let resident_device_bytes = self
460            .resident_device_bytes
461            .checked_add(batch_fixed_resident_overhead_bytes())
462            .ok_or_else(|| {
463                PipelineError::Backend(
464                    "megakernel resident byte estimate overflowed u64. Fix: shard resident state before launch recommendation."
465                        .to_string(),
466                )
467            })?;
468        MegakernelLaunchPolicy::standard()
469            .recommend(MegakernelLaunchRequest {
470                queue_len,
471                requested_worker_groups: self.worker_groups,
472                max_workgroup_size_x: self.workgroup_size_x,
473                max_compute_workgroups_per_dimension: limits.max_compute_workgroups_per_dimension,
474                max_compute_invocations_per_workgroup: limits.max_compute_invocations_per_workgroup,
475                requested_hit_capacity: self.hit_capacity,
476                expected_hits_per_item: 1,
477                hot_opcode_count: self.hot_opcode_count,
478                hot_window_count: self.hot_window_count,
479                requeue_count: self.requeue_count,
480                max_priority_age: self.max_priority_age,
481                graph_node_count: if self.graph_node_count == 0 {
482                    queue_len
483                } else {
484                    self.graph_node_count
485                },
486                graph_edge_count: self.graph_edge_count,
487                frontier_density_bps: self.frontier_density_bps,
488                memory_pressure_bps: self.memory_pressure_bps,
489                resident_device_bytes,
490                device_memory_budget_bytes: self.device_memory_budget_bytes,
491            })
492            .map_err(|source| PipelineError::Backend(source.to_string()))
493    }
494}
495
496fn batch_fixed_resident_overhead_bytes() -> u64 {
497    dispatcher_usize_to_u64(QUEUE_STATE_WORDS, "queue-state word count").saturating_mul(
498        dispatcher_usize_to_u64(std::mem::size_of::<u32>(), "u32 byte width"),
499    )
500}
501
502/// Widen a megakernel ABI constant to `u64`, failing closed on an out-of-range value.
503///
504/// # Panics
505/// Panics when `value` does not fit `u64`. These are ABI constants baked into the
506/// dispatcher, so an out-of-range value is a build-time mistake; silently clamping it
507/// would hand the kernel a wrong buffer size.
508fn dispatcher_usize_to_u64<T>(value: T, label: &'static str) -> u64
509where
510    T: TryInto<u64> + Copy + std::fmt::Display,
511    T::Error: std::fmt::Display,
512{
513    match value.try_into() {
514        Ok(v) => v,
515        Err(error) => {
516            // Fail closed: a constant that cannot fit u64 is a miscompile waiting
517            // to happen. Surface the label and value loudly rather than embedding
518            // u64::MAX and letting downstream checked_mul silently blame
519            // arithmetic overflow instead of the root cause (Law 10).
520            panic!(
521                "dispatcher ABI constant '{label}' value {value} cannot fit u64: {error}. Fix: keep all megakernel ABI constants within u64 range."
522            )
523        }
524    }
525}
526
527/// Narrow a megakernel ABI constant to `u32`, failing closed on an out-of-range value.
528///
529/// # Panics
530/// Panics when `value` does not fit `u32`. See [`dispatcher_usize_to_u64`]: a clamped
531/// ABI constant would silently mis-size a dispatch.
532fn dispatcher_abi_u32<T>(value: T, label: &'static str) -> u32
533where
534    T: TryInto<u32> + Copy + std::fmt::Display,
535    T::Error: std::fmt::Display,
536{
537    match value.try_into() {
538        Ok(v) => v,
539        Err(error) => {
540            // Fail closed: a constant that cannot fit u32 is a shader miscompile
541            // u32::MAX embedded as a WGSL literal would corrupt ABI offsets in the
542            // generated GPU program. Surface the label and value loudly (Law 10).
543            panic!(
544                "dispatcher ABI constant '{label}' value {value} cannot fit u32: {error}. Fix: keep all megakernel ABI constants within u32 range."
545            )
546        }
547    }
548}
549
550/// Observability returned from one batched dispatch.
551#[derive(Debug, Clone)]
552pub struct BatchDispatchReport {
553    /// Sparse hit count written by the device (clamped to `hit_capacity`; the
554    /// number of `hits` actually decodable).
555    pub hit_count: u32,
556    /// Matches the device produced BEYOND `hit_capacity` and therefore DROPPED
557    /// from the hit ring (raw atomic head minus capacity). `> 0` means this
558    /// dispatch's hit set is INCOMPLETE, a recall-critical overflow the caller
559    /// MUST surface and recover (re-scan with a larger ring or on the host),
560    /// never treat as a complete result. Zero on a healthy dispatch.
561    pub dropped_hits: u32,
562    /// Hits compacted out of the sparse ring.
563    pub hits: Vec<HitRecord>,
564    /// Work items processed by the queue.
565    pub items_processed: u32,
566    /// Wall-clock GPU execution time.
567    pub wall_time: Duration,
568    /// Rules that were isolated from the batch because their catalog entry was
569    /// malformed. The rest of the batch still ran.
570    pub rejected_rules: Vec<BatchRuleRejection>,
571    /// Production telemetry for performance gates and dispatch tuning.
572    pub telemetry: BatchDispatchTelemetry,
573}
574
575/// Megakernel dispatch counters returned when the caller owns hit storage.
576#[derive(Debug, Clone)]
577pub struct BatchDispatchSummary {
578    /// Sparse hit count written by the device (clamped to `hit_capacity`; the
579    /// number of `HitRecord`s decoded into the caller's storage).
580    pub hit_count: u32,
581    /// Matches the device produced BEYOND `hit_capacity` and therefore DROPPED
582    /// from the hit ring (raw atomic head minus capacity). `> 0` means this
583    /// dispatch's hit set is INCOMPLETE, a recall-critical overflow the caller
584    /// MUST surface and recover, never treat as a complete result. Zero on a
585    /// healthy dispatch.
586    pub dropped_hits: u32,
587    /// Work items processed by the queue.
588    pub items_processed: u32,
589    /// Wall-clock GPU execution time.
590    pub wall_time: Duration,
591    /// Rules that were isolated from the batch because their catalog entry was
592    /// malformed. The rest of the batch still ran.
593    pub rejected_rules: Vec<BatchRuleRejection>,
594    /// Production telemetry for performance gates and dispatch tuning.
595    pub telemetry: BatchDispatchTelemetry,
596}
597
598/// Megakernel dispatch counters used by scale/performance gates.
599#[derive(Debug, Clone, Copy, PartialEq, Eq)]
600pub struct BatchDispatchTelemetry {
601    /// Bytes uploaded by this dispatch for rule-catalog refreshes.
602    pub bytes_uploaded: u64,
603    /// Bytes read back from queue-state and sparse hit output buffers.
604    pub bytes_read_back: u64,
605    /// Total host/device transfer bytes directly attributable to this dispatch.
606    pub bytes_moved: u64,
607    /// Resident allocations performed for refreshed rule-catalog buffers.
608    pub resident_allocations: u32,
609    /// Kernel launches submitted for the megakernel dispatch.
610    pub kernel_launches: u32,
611    /// Host-visible synchronization/readback wait points.
612    pub sync_points: u32,
613    /// Approximate lane occupancy in basis points, capped at 10000.
614    pub occupancy_proxy_bps: u16,
615    /// Active frontier density passed into the launch policy.
616    pub frontier_density_bps: u16,
617    /// Queue-state readback volume.
618    pub queue_state_readback_bytes: u64,
619    /// Sparse hit-ring readback volume.
620    pub hit_readback_bytes: u64,
621    /// Estimated peak device bytes required by the selected launch plan.
622    pub estimated_peak_device_bytes: u64,
623    /// Hard device-memory budget applied to this dispatch. Zero means unbounded.
624    pub device_memory_budget_bytes: u64,
625    /// Scale-aware topology selected by the launch policy.
626    pub topology: MegakernelDispatchTopology,
627    /// Whether this dispatch reused a cached fixed-batch launch plan.
628    pub dispatch_plan_cache_hit: bool,
629    /// Number of fixed-batch launch plans resident in the dispatcher cache.
630    pub dispatch_plan_cache_entries: u16,
631}
632
633impl Default for BatchDispatchTelemetry {
634    fn default() -> Self {
635        Self {
636            bytes_uploaded: 0,
637            bytes_read_back: 0,
638            bytes_moved: 0,
639            resident_allocations: 0,
640            kernel_launches: 0,
641            sync_points: 0,
642            occupancy_proxy_bps: 0,
643            frontier_density_bps: 0,
644            queue_state_readback_bytes: 0,
645            hit_readback_bytes: 0,
646            estimated_peak_device_bytes: 0,
647            device_memory_budget_bytes: 0,
648            topology: MegakernelDispatchTopology::SparseFrontier,
649            dispatch_plan_cache_hit: false,
650            dispatch_plan_cache_entries: 0,
651        }
652    }
653}
654
655struct RuleBufferUpdate {
656    rejected_rules: Vec<BatchRuleRejection>,
657    uploaded_bytes: u64,
658    resident_allocations: u32,
659}
660
661const BATCH_PIPELINE_CACHE_CAP: usize = 32;
662
663/// One compiled batched megakernel pipeline plus cached rule buffers.
664pub struct BatchDispatcher {
665    backend: WgpuBackend,
666    config: BatchDispatchConfig,
667    hit_writer: BatchHitWriter,
668    pipeline: Arc<WgpuPipeline>,
669    pipeline_cache: BatchPipelineCache,
670    launch: MegakernelLaunchRecommendation,
671    dispatch_plan_cache: BatchDispatchPlanCache,
672    active_rule_fingerprints: Vec<[u8; 32]>,
673    fingerprint_scratch: Vec<[u8; 32]>,
674    fingerprint_occupied_scratch: Vec<bool>,
675    fingerprint_addressed_scratch: Vec<bool>,
676    rejection_scratch: Vec<BatchRuleRejection>,
677    packing_scratch: RuleCatalogPackingScratch,
678    rule_meta: Option<GpuBufferHandle>,
679    transitions: Option<GpuBufferHandle>,
680    accept: Option<GpuBufferHandle>,
681    /// Shared byte→class maps (256 entries per unique DFA) backing the
682    /// compressed transition tables. Uploaded alongside the other rule buffers.
683    class_maps: Option<GpuBufferHandle>,
684    queue_state_bytes: Vec<u8>,
685    hit_bytes: Vec<u8>,
686}
687
688impl std::fmt::Debug for BatchDispatcher {
689    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
690        formatter
691            .debug_struct("BatchDispatcher")
692            .field("config", &self.config)
693            .field("hit_writer", &self.hit_writer)
694            .field("pipeline_id", &self.pipeline.id())
695            .field("launch", &self.launch)
696            .field("rule_count", &self.active_rule_fingerprints.len())
697            .finish()
698    }
699}
700
701impl BatchDispatcher {
702    /// Compile the batched megakernel program on a live wgpu backend.
703    ///
704    /// Defaults to the [`BatchHitWriter::Scalar`] hit writer. This is a
705    /// CORRECTNESS requirement, not a performance default: the batch kernel's
706    /// per-work-item scan (`dfa_byte_scanner`) loops `scan_start..emit_end`, so
707    /// lanes in one subgroup execute DIFFERENT iteration counts (segments/files
708    /// differ in length) and exit the loop at different points, divergent control
709    /// flow.
710    /// The hierarchical-subgroup writer aggregates hits with `subgroupBallot`/
711    /// `subgroupAdd`/`subgroupShuffle` and elects a leader lane; under divergence
712    /// the elected leader can already have exited, so its reserved ring slot is
713    /// never broadcast and hits found by still-running lanes are dropped. That
714    /// surfaced as a real, data-dependent recall loss in the downstream GPU≡CPU
715    /// parity gate (6 of 46 detector firings silently missed, every miss a match
716    /// found after its subgroup's leader lane finished a shorter file). The
717    /// scalar writer does one independent `atomicAdd` per hit and is correct
718    /// under ANY divergence; for sparse credential matches the per-byte DFA step
719    /// dominates and the extra atomics are negligible. Callers with a genuinely
720    /// uniform-iteration kernel may opt into a subgroup writer via
721    /// [`Self::new_with_hit_writer`].
722    ///
723    /// # Errors
724    ///
725    /// Returns [`PipelineError::Backend`] when pipeline compilation fails.
726    pub fn new(backend: WgpuBackend, config: BatchDispatchConfig) -> Result<Self, PipelineError> {
727        Self::new_with_hit_writer(backend, config, BatchHitWriter::Scalar)
728    }
729
730    /// Compile with an explicit sparse-hit publication algorithm.
731    ///
732    /// # Errors
733    ///
734    /// Returns [`PipelineError::Backend`] when hierarchical subgroup atomics are
735    /// requested on a backend that reports no subgroup support, or when
736    /// pipeline compilation fails.
737    pub fn new_with_hit_writer(
738        backend: WgpuBackend,
739        mut config: BatchDispatchConfig,
740        requested_hit_writer: BatchHitWriter,
741    ) -> Result<Self, PipelineError> {
742        if config.workgroup_size_x == 0 {
743            return Err(PipelineError::QueueFull {
744                queue: "submission",
745                fix: "BatchDispatchConfig requires non-zero workgroup_size_x",
746            });
747        }
748        let seed_queue_len = config
749            .worker_groups
750            .max(1)
751            .checked_mul(config.workgroup_size_x)
752            .ok_or_else(|| PipelineError::QueueFull {
753                queue: "submission",
754                fix: "megakernel seed queue length overflowed u32; reduce worker_groups or workgroup_size_x",
755            })?;
756        let launch = config.launch_recommendation(backend.device_limits(), seed_queue_len)?;
757        if config.worker_groups == 0 {
758            config.worker_groups = launch.worker_groups;
759        }
760        if config.hit_capacity == 0 {
761            config.hit_capacity = launch.hit_capacity;
762        }
763        // The batch kernel's per-work-item scan (`dfa_byte_scanner`) loops
764        // `scan_start..emit_end`, so subgroup lanes diverge as shorter
765        // segments/files finish first. The hierarchical-subgroup writer aggregates hits with
766        // subgroup ballot/add/shuffle and REQUIRES uniform control flow (see the
767        // `hierarchical_atomics` module contract); under this divergence it
768        // strands the elected leader's reserved ring slot once that lane exits,
769        // silently dropping hits found by still-running lanes (a real recall loss
770        // in the downstream GPU≡CPU parity gate). So the hierarchical writer is never
771        // sound for this dispatcher: `Auto` (which would resolve to Hierarchical
772        // on a subgroup backend) DOWNGRADES to the correct scalar writer, and an
773        // EXPLICIT hierarchical request is a caller error that fails loudly rather
774        // than silently losing recall.
775        let resolved = requested_hit_writer.resolve_for_backend(backend.supports_subgroup_ops())?;
776        let hit_writer = match resolved {
777            BatchHitWriter::HierarchicalSubgroup => {
778                if matches!(requested_hit_writer, BatchHitWriter::Auto) {
779                    BatchHitWriter::Scalar
780                } else {
781                    return Err(PipelineError::Backend(
782                        "BatchHitWriter::HierarchicalSubgroup is unsound for the batched megakernel: \
783                         its per-work-item DFA scan loops scan_start..emit_end, so subgroup lanes diverge \
784                         as shorter segments/files finish, and subgroup hit-aggregation requires uniform \
785                         control flow, under divergence the leader lane exits before broadcasting \
786                         its reserved ring slot and hits are silently dropped (detector-firing recall \
787                         loss). Fix: use BatchHitWriter::Scalar (the default) or BatchHitWriter::Auto."
788                            .to_string(),
789                    ));
790                }
791            }
792            other => other,
793        };
794        let program = build_batch_program(
795            config.workgroup_size_x,
796            config.worker_groups,
797            config.hit_capacity,
798            hit_writer,
799        );
800        let pipeline = backend.compile_persistent(&program, &DispatchConfig::default())?;
801        let pipeline_workgroup_size_x = config.workgroup_size_x;
802        let pipeline_hit_capacity = config.hit_capacity;
803        let mut pipeline_cache = BatchPipelineCache::with_cap(BATCH_PIPELINE_CACHE_CAP);
804        pipeline_cache.seed(
805            BatchPipelineShape {
806                workgroup_size_x: pipeline_workgroup_size_x,
807                worker_groups: launch.worker_groups,
808                hit_capacity: pipeline_hit_capacity,
809            },
810            pipeline.clone(),
811        );
812        Ok(Self {
813            backend,
814            config,
815            hit_writer,
816            pipeline: pipeline.clone(),
817            pipeline_cache,
818            launch,
819            dispatch_plan_cache: BatchDispatchPlanCache::default(),
820            active_rule_fingerprints: Vec::new(),
821            fingerprint_scratch: Vec::new(),
822            fingerprint_occupied_scratch: Vec::new(),
823            fingerprint_addressed_scratch: Vec::new(),
824            rejection_scratch: Vec::new(),
825            packing_scratch: RuleCatalogPackingScratch::default(),
826            rule_meta: None,
827            transitions: None,
828            accept: None,
829            class_maps: None,
830            queue_state_bytes: Vec::with_capacity(QUEUE_STATE_WORDS * std::mem::size_of::<u32>()),
831            hit_bytes: Vec::new(),
832        })
833    }
834
835    /// Dispatch one `FileBatch` against many compiled DFA rules in one launch.
836    ///
837    /// # Errors
838    ///
839    /// Returns [`PipelineError::Backend`] on pipeline, upload, or readback
840    /// failures.
841    pub fn dispatch(
842        &mut self,
843        batch: &FileBatch,
844        rules: &[BatchRuleProgram],
845    ) -> Result<BatchDispatchReport, PipelineError> {
846        let hit_capacity = usize::try_from(batch.hit_capacity()).map_err(|source| {
847            PipelineError::Backend(format!(
848                "batch hit capacity cannot fit usize: {source}. Fix: reduce hit_capacity or shard the batch."
849            ))
850        })?;
851        let mut hits = Vec::with_capacity(hit_capacity);
852        let summary = self.dispatch_into(batch, rules, &mut hits)?;
853        Ok(BatchDispatchReport {
854            hit_count: summary.hit_count,
855            dropped_hits: summary.dropped_hits,
856            hits,
857            items_processed: summary.items_processed,
858            wall_time: summary.wall_time,
859            rejected_rules: summary.rejected_rules,
860            telemetry: summary.telemetry,
861        })
862    }
863
864    /// Dispatch one `FileBatch` while decoding sparse hits into caller-owned
865    /// storage.
866    ///
867    /// Reusing `hits` avoids a fresh hit-vector allocation on hot repeated
868    /// megakernel calls. The vector is cleared before decode and keeps its
869    /// capacity unless the actual hit count exceeds it.
870    ///
871    /// # Errors
872    ///
873    /// Returns [`PipelineError::Backend`] on pipeline, upload, or readback
874    /// failures.
875    pub fn dispatch_into(
876        &mut self,
877        batch: &FileBatch,
878        rules: &[BatchRuleProgram],
879        hits: &mut Vec<HitRecord>,
880    ) -> Result<BatchDispatchSummary, PipelineError> {
881        if rules.is_empty() {
882            hits.clear();
883            let dynamic_plan = self.dispatch_plan(batch)?;
884            return Ok(BatchDispatchSummary {
885                hit_count: 0,
886                dropped_hits: 0,
887                items_processed: 0,
888                wall_time: Duration::ZERO,
889                rejected_rules: Vec::new(),
890                telemetry: BatchDispatchTelemetry {
891                    topology: dynamic_plan.plan.topology,
892                    frontier_density_bps: self.config.frontier_density_bps,
893                    estimated_peak_device_bytes: dynamic_plan.plan.estimated_peak_device_bytes,
894                    device_memory_budget_bytes: dynamic_plan.plan.device_memory_budget_bytes,
895                    dispatch_plan_cache_hit: dynamic_plan.cache_hit,
896                    dispatch_plan_cache_entries: dynamic_plan.cache_entries,
897                    ..BatchDispatchTelemetry::default()
898                },
899            });
900        }
901        let dynamic_plan = self.dispatch_plan(batch)?;
902        let pipeline = self.pipeline_for_plan(dynamic_plan.plan)?;
903        let rule_update = self.ensure_rule_buffers(rules)?;
904        batch.reset_queue_state()?;
905
906        let Some(class_maps) = self.class_maps.as_ref() else {
907            return Err(PipelineError::Backend(
908                "byte-class map buffer missing after ensure_rule_buffers. Fix: keep megakernel rule buffer initialization atomic.".to_string(),
909            ));
910        };
911        let Some(rule_meta) = self.rule_meta.as_ref() else {
912            return Err(PipelineError::Backend(
913                "rule metadata buffer missing after ensure_rule_buffers. Fix: keep megakernel rule buffer initialization atomic.".to_string(),
914            ));
915        };
916        let Some(transitions) = self.transitions.as_ref() else {
917            return Err(PipelineError::Backend(
918                "transition buffer missing after ensure_rule_buffers. Fix: keep megakernel rule buffer initialization atomic.".to_string(),
919            ));
920        };
921        let Some(accept) = self.accept.as_ref() else {
922            return Err(PipelineError::Backend(
923                "accept buffer missing after ensure_rule_buffers. Fix: keep megakernel rule buffer initialization atomic.".to_string(),
924            ));
925        };
926        // Input order MUST match the non-Shared storage buffer DECLARATION order
927        // in `batch_program_buffers` (offsets, metadata, class_maps, haystack,
928        // rule_meta, transitions, accept, segments), not literal binding numbers
929        //: the persistent pipeline binds inputs positionally in that order.
930        // `segments` is declared last so it is the final positional input.
931        let inputs = [
932            batch.offsets(),
933            batch.metadata(),
934            class_maps,
935            batch.haystack(),
936            rule_meta,
937            transitions,
938            accept,
939            batch.segments(),
940        ];
941        let outputs = [batch.queue_state(), batch.hit_ring()];
942        let start = Instant::now();
943        pipeline.dispatch_persistent_borrowed(
944            &inputs,
945            &outputs,
946            None,
947            [dynamic_plan.plan.worker_groups, 1, 1],
948        )?;
949
950        let (device, queue) = &*self.backend.device_queue();
951        wait_for_persistent_dispatch(device, start, self.config.timeout)?;
952        let wall_time = start.elapsed();
953        self.queue_state_bytes.clear();
954        let queue_state_readback_bytes = batch_fixed_resident_overhead_bytes();
955        batch.queue_state().readback_prefix(
956            device,
957            queue,
958            queue_state_readback_bytes,
959            &mut self.queue_state_bytes,
960        )?;
961        let queue_state_word_count =
962            validate_u32_readback_words(&self.queue_state_bytes, "queue-state")?;
963        if queue_state_word_count < QUEUE_STATE_WORDS {
964            return Err(PipelineError::Backend(format!(
965                "queue-state readback exposed {} words, expected at least {}. Fix: keep the queue-state buffer sized for every control word.",
966                queue_state_word_count,
967                QUEUE_STATE_WORDS
968            )));
969        }
970        // The kernel `atomicAdd(HIT_HEAD, 1)`s for EVERY match it finds, then
971        // writes only when `slot < hit_capacity`: so the raw head is the true
972        // number of matches the device produced, which can exceed the ring. We
973        // can only read back `hit_capacity` slots, but the overflow is a
974        // recall-critical signal: clamping it away silently would hide dropped
975        // matches (Law 10). Split the raw head into the readable count and the
976        // dropped count and surface the latter to the caller.
977        let raw_hit_head = read_u32_word(
978            &self.queue_state_bytes,
979            "queue-state",
980            queue_state_word::HIT_HEAD,
981        )?;
982        let (hit_count, dropped_hits) = split_hit_overflow(raw_hit_head, batch.hit_capacity());
983        let items_processed = read_u32_word(
984            &self.queue_state_bytes,
985            "queue-state",
986            queue_state_word::DONE_COUNT,
987        )?;
988
989        // Fail-closed drain-completion guard (Law 10: no silent recall loss).
990        //
991        // The claim loop now DRAINS: every resident lane keeps issuing
992        // `atomicAdd(HEAD, 1)` until it claims past the end of the queue, so after
993        // a COMPLETE drain `HEAD == queue_len + resident_lanes >= queue_len` (one
994        // past-the-end claim per lane). The only way `HEAD < queue_len` can be
995        // observed is an INCOMPLETE drain, the dispatch was cut short (e.g. the
996        // dispatch timeout fired) before the queue was exhausted, leaving the
997        // indices `[HEAD, queue_len)` unscanned and their matches missing from the
998        // ring with `dropped_hits == 0`: an INVISIBLE recall loss. (HEAD, not
999        // DONE_COUNT: a claimed-but-rejected rule still advances HEAD, so HEAD is
1000        // the rejected-rule-independent "was every work-item handed out?" signal.)
1001        // Surface it loudly instead of returning a partial hit set.
1002        let claims_attempted = read_u32_word(
1003            &self.queue_state_bytes,
1004            "queue-state",
1005            queue_state_word::HEAD,
1006        )?;
1007        let expected_items = batch.queue_len();
1008        if claims_attempted < expected_items {
1009            return Err(PipelineError::DrainIncomplete {
1010                descriptor: "megakernel",
1011                claimed: claims_attempted,
1012                expected: expected_items,
1013                unit: "work-items",
1014            });
1015        }
1016
1017        self.hit_bytes.clear();
1018        let hit_readback_bytes = u64::from(hit_count)
1019            .checked_mul(dispatcher_usize_to_u64(
1020                HIT_RECORD_WORDS,
1021                "hit-record word count",
1022            ))
1023            .and_then(|words| {
1024                words.checked_mul(dispatcher_usize_to_u64(
1025                    std::mem::size_of::<u32>(),
1026                    "u32 byte width",
1027                ))
1028            })
1029            .ok_or_else(|| {
1030                PipelineError::Backend(
1031                    "hit-ring readback length overflowed u64. Fix: reduce hit_capacity or shard the batch."
1032                        .to_string(),
1033                )
1034            })?;
1035        batch
1036            .hit_ring()
1037            .readback_prefix(device, queue, hit_readback_bytes, &mut self.hit_bytes)?;
1038        decode_hits_from_readback_into(&self.hit_bytes, hit_count, hits)?;
1039        let bytes_read_back = queue_state_readback_bytes
1040            .checked_add(hit_readback_bytes)
1041            .ok_or_else(|| {
1042                PipelineError::Backend(
1043                    "batch readback byte accounting overflowed u64. Fix: shard the batch before readback."
1044                        .to_string(),
1045                )
1046            })?;
1047        let bytes_moved = rule_update
1048            .uploaded_bytes
1049            .checked_add(bytes_read_back)
1050            .ok_or_else(|| {
1051                PipelineError::Backend(
1052                    "batch moved-byte accounting overflowed u64. Fix: shard the batch before dispatch."
1053                        .to_string(),
1054                )
1055            })?;
1056
1057        Ok(BatchDispatchSummary {
1058            hit_count,
1059            dropped_hits,
1060            items_processed,
1061            wall_time,
1062            rejected_rules: rule_update.rejected_rules,
1063            telemetry: BatchDispatchTelemetry {
1064                bytes_uploaded: rule_update.uploaded_bytes,
1065                bytes_read_back,
1066                bytes_moved,
1067                resident_allocations: rule_update.resident_allocations,
1068                kernel_launches: 1,
1069                sync_points: 2,
1070                occupancy_proxy_bps: occupancy_proxy_bps(
1071                    items_processed,
1072                    dynamic_plan.plan.worker_groups,
1073                    self.config.workgroup_size_x,
1074                ),
1075                frontier_density_bps: self.config.frontier_density_bps,
1076                queue_state_readback_bytes,
1077                hit_readback_bytes,
1078                estimated_peak_device_bytes: dynamic_plan.plan.estimated_peak_device_bytes,
1079                device_memory_budget_bytes: dynamic_plan.plan.device_memory_budget_bytes,
1080                topology: dynamic_plan.plan.topology,
1081                dispatch_plan_cache_hit: dynamic_plan.cache_hit,
1082                dispatch_plan_cache_entries: dynamic_plan.cache_entries,
1083            },
1084        })
1085    }
1086
1087    fn pipeline_for_plan(
1088        &mut self,
1089        plan: BatchDispatchPlan,
1090    ) -> Result<Arc<WgpuPipeline>, PipelineError> {
1091        let shape = BatchPipelineShape {
1092            workgroup_size_x: plan.workgroup_size_x,
1093            worker_groups: plan.worker_groups,
1094            hit_capacity: plan.hit_capacity,
1095        };
1096        if let Some(pipeline) = self.pipeline_cache.get(shape) {
1097            return Ok(pipeline);
1098        }
1099        let program = build_batch_program(
1100            plan.workgroup_size_x,
1101            plan.worker_groups,
1102            plan.hit_capacity,
1103            self.hit_writer,
1104        );
1105        let pipeline = self
1106            .backend
1107            .compile_persistent(&program, &DispatchConfig::default())?;
1108        self.pipeline_cache.insert(shape, pipeline.clone());
1109        Ok(pipeline)
1110    }
1111
1112    fn dispatch_plan(
1113        &mut self,
1114        batch: &FileBatch,
1115    ) -> Result<BatchDispatchPlanLookup, PipelineError> {
1116        let queue_len = batch.queue_len();
1117        if let Some(plan) = self.dispatch_plan_cache.get(queue_len) {
1118            return Ok(BatchDispatchPlanLookup {
1119                plan,
1120                cache_hit: true,
1121                cache_entries: self.dispatch_plan_cache.len_u16(),
1122            });
1123        }
1124        let mut recommendation = self
1125            .config
1126            .launch_recommendation(self.backend.device_limits(), queue_len)?;
1127        let resident_hit_capacity = batch.hit_capacity();
1128        if recommendation.hit_capacity > resident_hit_capacity {
1129            let removed_hit_bytes = u64::from(recommendation.hit_capacity - resident_hit_capacity)
1130                .checked_mul(dispatcher_usize_to_u64(
1131                    HIT_RECORD_WORDS,
1132                    "hit-record word count",
1133                ))
1134                .and_then(|words| {
1135                    words.checked_mul(dispatcher_usize_to_u64(
1136                        std::mem::size_of::<u32>(),
1137                        "u32 byte width",
1138                    ))
1139                })
1140                .ok_or_else(|| {
1141                    PipelineError::Backend(
1142                        "resident hit-capacity byte adjustment overflowed u64. Fix: shard the batch before dispatch planning."
1143                            .to_string(),
1144                    )
1145                })?;
1146            recommendation.hit_capacity = resident_hit_capacity;
1147            recommendation.estimated_peak_device_bytes = recommendation
1148                .estimated_peak_device_bytes
1149                .checked_sub(removed_hit_bytes)
1150                .ok_or_else(|| {
1151                    PipelineError::Backend(
1152                        "resident hit-capacity adjustment exceeded peak device estimate. Fix: keep launch recommendation and resident batch capacity synchronized."
1153                            .to_string(),
1154                    )
1155                })?;
1156        }
1157        let plan = BatchDispatchPlan::from_recommendation(queue_len, &self.config, recommendation);
1158        self.dispatch_plan_cache.insert(plan);
1159        Ok(BatchDispatchPlanLookup {
1160            plan,
1161            cache_hit: false,
1162            cache_entries: self.dispatch_plan_cache.len_u16(),
1163        })
1164    }
1165
1166    fn ensure_rule_buffers(
1167        &mut self,
1168        rules: &[BatchRuleProgram],
1169    ) -> Result<RuleBufferUpdate, PipelineError> {
1170        accepted_rule_fingerprints_and_rejections_into(
1171            rules,
1172            &mut self.fingerprint_scratch,
1173            &mut self.fingerprint_occupied_scratch,
1174            &mut self.fingerprint_addressed_scratch,
1175            &mut self.rejection_scratch,
1176        );
1177        if self.fingerprint_scratch == self.active_rule_fingerprints {
1178            return Ok(RuleBufferUpdate {
1179                rejected_rules: if self.rejection_scratch.is_empty() {
1180                    Vec::new()
1181                } else {
1182                    self.rejection_scratch.clone()
1183                },
1184                uploaded_bytes: 0,
1185                resident_allocations: 0,
1186            });
1187        }
1188
1189        pack_rule_catalog_into(rules, &mut self.packing_scratch)?;
1190        // rule_meta words = entries * RULE_META_WORDS (each RuleMeta is
1191        // RULE_META_WORDS u32s); transitions + accept + class_maps are flat u32
1192        // vecs. Account for all four uploaded device buffers.
1193        let rule_meta_words = self
1194            .packing_scratch
1195            .rule_meta
1196            .len()
1197            .checked_mul(RULE_META_WORDS)
1198            .ok_or_else(|| {
1199                PipelineError::Backend(
1200                    "rule metadata upload word count overflowed usize. Fix: shard the rule catalog before upload."
1201                        .to_string(),
1202                )
1203            })?;
1204        let uploaded_words = rule_meta_words
1205            .checked_add(self.packing_scratch.transitions.len())
1206            .and_then(|words| words.checked_add(self.packing_scratch.accept.len()))
1207            .and_then(|words| words.checked_add(self.packing_scratch.class_maps.len()))
1208            .ok_or_else(|| {
1209                PipelineError::Backend(
1210                    "rule catalog upload word count overflowed usize. Fix: shard the rule catalog before upload."
1211                        .to_string(),
1212                )
1213            })?;
1214        let uploaded_bytes = uploaded_words
1215            .checked_mul(std::mem::size_of::<u32>())
1216            .and_then(|bytes| u64::try_from(bytes).ok())
1217            .ok_or_else(|| {
1218                PipelineError::Backend(
1219                    "rule catalog upload byte count overflowed u64. Fix: shard the rule catalog before upload."
1220                        .to_string(),
1221                )
1222            })?;
1223        let (device, queue) = &*self.backend.device_queue();
1224        self.rule_meta = Some(GpuBufferHandle::upload(
1225            device,
1226            queue,
1227            bytemuck::cast_slice(&self.packing_scratch.rule_meta),
1228            persistent_storage_binding_usage(),
1229        )?);
1230        self.transitions = Some(GpuBufferHandle::upload(
1231            device,
1232            queue,
1233            bytemuck::cast_slice(&self.packing_scratch.transitions),
1234            persistent_storage_binding_usage(),
1235        )?);
1236        self.accept = Some(GpuBufferHandle::upload(
1237            device,
1238            queue,
1239            bytemuck::cast_slice(&self.packing_scratch.accept),
1240            persistent_storage_binding_usage(),
1241        )?);
1242        self.class_maps = Some(GpuBufferHandle::upload(
1243            device,
1244            queue,
1245            bytemuck::cast_slice(&self.packing_scratch.class_maps),
1246            persistent_storage_binding_usage(),
1247        )?);
1248        if self.active_rule_fingerprints.len() == self.fingerprint_scratch.len() {
1249            self.active_rule_fingerprints
1250                .copy_from_slice(&self.fingerprint_scratch);
1251        } else {
1252            self.active_rule_fingerprints.clear();
1253            self.active_rule_fingerprints
1254                .extend_from_slice(&self.fingerprint_scratch);
1255        }
1256        Ok(RuleBufferUpdate {
1257            rejected_rules: if self.packing_scratch.rejected_rules.is_empty() {
1258                Vec::new()
1259            } else {
1260                self.packing_scratch.rejected_rules.clone()
1261            },
1262            uploaded_bytes,
1263            resident_allocations: 4,
1264        })
1265    }
1266}
1267
1268fn occupancy_proxy_bps(items_processed: u32, worker_groups: u32, workgroup_size_x: u32) -> u16 {
1269    let lanes = u64::from(worker_groups.max(1))
1270        .checked_mul(u64::from(workgroup_size_x.max(1)))
1271        .unwrap_or(u64::MAX);
1272    crate::numeric::WGPU_NUMERIC
1273        .ratio_basis_points_u64_wide(
1274            u64::from(items_processed),
1275            lanes.max(1),
1276            0,
1277            "batch occupancy proxy",
1278        )
1279        .min(10_000) as u16
1280}
1281
1282fn validate_u32_readback_words(bytes: &[u8], label: &'static str) -> Result<usize, PipelineError> {
1283    if bytes.len() % std::mem::size_of::<u32>() != 0 {
1284        return Err(PipelineError::Backend(format!(
1285            "{label} readback exposed {} bytes, which is not a whole number of u32 words. Fix: keep readback lengths 4-byte aligned.",
1286            bytes.len()
1287        )));
1288    }
1289    Ok(bytes.len() / std::mem::size_of::<u32>())
1290}
1291
1292fn read_u32_word(
1293    bytes: &[u8],
1294    label: &'static str,
1295    word_index: usize,
1296) -> Result<u32, PipelineError> {
1297    let offset = word_index
1298        .checked_mul(std::mem::size_of::<u32>())
1299        .ok_or_else(|| {
1300            PipelineError::Backend(format!(
1301                "{label} word offset overflowed usize. Fix: split the readback before decoding."
1302            ))
1303        })?;
1304    let word = bytes.get(offset..offset + std::mem::size_of::<u32>()).ok_or_else(|| {
1305        PipelineError::Backend(format!(
1306            "{label} readback is missing u32 word {word_index}. Fix: request a large enough readback prefix."
1307        ))
1308    })?;
1309    Ok(u32::from_le_bytes([word[0], word[1], word[2], word[3]]))
1310}
1311
1312fn wait_for_persistent_dispatch(
1313    device: &wgpu::Device,
1314    start: Instant,
1315    timeout: Duration,
1316) -> Result<(), PipelineError> {
1317    let mut backoff = crate::wait_backoff::AdaptiveWaitBackoff::from_micros(64, 5, 50, 8);
1318    loop {
1319        if crate::runtime::device::poll_device_once(device)
1320            .map_err(|error| PipelineError::Backend(error.to_string()))?
1321            .is_queue_empty()
1322        {
1323            return Ok(());
1324        }
1325        let elapsed = start.elapsed();
1326        if elapsed >= timeout {
1327            return Err(PipelineError::Backend(format!(
1328                "batch megakernel dispatch exceeded timeout before readback: took {elapsed:?}, budget {timeout:?}. Fix: raise BatchDispatchConfig.timeout or split the batch.",
1329            )));
1330        }
1331        let remaining = timeout.checked_sub(elapsed).ok_or_else(|| {
1332            PipelineError::Backend(format!(
1333                "batch megakernel timeout arithmetic underflowed after elapsed {elapsed:?} exceeded budget {timeout:?}. Fix: split the batch or raise BatchDispatchConfig.timeout deliberately.",
1334            ))
1335        })?;
1336        backoff.idle_for(remaining);
1337    }
1338}
1339
1340fn build_batch_program(
1341    workgroup_size_x: u32,
1342    worker_groups: u32,
1343    hit_capacity: u32,
1344    hit_writer: BatchHitWriter,
1345) -> Program {
1346    // Persistent DRAIN loop: every resident lane keeps claiming work-items with
1347    // `atomicAdd(HEAD, 1)` until its claim lands past the end of the queue
1348    // (`claim >= QUEUE_LEN`), then returns. This drains the full
1349    // `segment_count * rule_count` queue for ANY number of resident lanes.
1350    //
1351    // It replaces a fixed `claim_budget = ceil(QUEUE_LEN / total_workers)` loop
1352    // that assumed exactly `total_workers` lanes each ran their full budget. When
1353    // fewer lanes were actually resident than that budget assumed, the queue was
1354    // never fully claimed: `found < expected` with `dropped_hits == 0`: a SILENT
1355    // recall loss (Law 10). The drain removes the dependency on the resident-lane
1356    // count entirely. Overhead is one extra past-the-end `atomicAdd` per resident
1357    // lane (the claim that observes `>= QUEUE_LEN` and returns), NOT per
1358    // work-item (a rounding error, not a 1/queue_len-scale pessimization).
1359    //
1360    // `worker_groups` now sizes only the dispatch grid (more resident lanes =
1361    // more parallelism); kernel correctness no longer depends on it.
1362    let _ = worker_groups;
1363    let queue_len = atomic_load_relaxed(
1364        "queue_state",
1365        Expr::u32(dispatcher_abi_u32(
1366            queue_state_word::QUEUE_LEN,
1367            "queue-state length word",
1368        )),
1369    );
1370    let mut loop_body = vec![
1371        Node::let_bind(
1372            "claim",
1373            Expr::atomic_add(
1374                "queue_state",
1375                Expr::u32(dispatcher_abi_u32(
1376                    queue_state_word::HEAD,
1377                    "queue-state head word",
1378                )),
1379                Expr::u32(1),
1380            ),
1381        ),
1382        // Past-the-end claim ⇒ the queue is drained for this lane. `Return` exits
1383        // the kernel: safe because the drain loop is the only top-level statement
1384        // (no post-loop finalization to skip) and `execute_batch_claim_body`
1385        // contains no workgroup barrier (no divergence deadlock).
1386        Node::if_then(Expr::ge(Expr::var("claim"), queue_len), vec![Node::Return]),
1387    ];
1388    loop_body.extend(execute_batch_claim_body(hit_writer));
1389
1390    Program::wrapped(
1391        batch_program_buffers(hit_capacity),
1392        [workgroup_size_x, 1, 1],
1393        vec![Node::forever(loop_body)],
1394    )
1395}
1396
1397#[derive(Clone, Copy)]
1398enum BatchAutomatonLayout {
1399    PerRule,
1400    Combined,
1401}
1402
1403fn batch_program_buffers_for_layout(
1404    hit_capacity: u32,
1405    layout: BatchAutomatonLayout,
1406) -> Vec<BufferDecl> {
1407    let hit_ring_words = hit_capacity.saturating_mul(4);
1408    let mut buffers = vec![
1409        BufferDecl::storage("file_offsets", 0, BufferAccess::ReadOnly, DataType::U32),
1410        BufferDecl::storage("file_metadata", 1, BufferAccess::ReadOnly, DataType::U32),
1411    ];
1412    match layout {
1413        BatchAutomatonLayout::PerRule => buffers.extend([
1414            BufferDecl::storage("class_maps", 2, BufferAccess::ReadOnly, DataType::U32),
1415            BufferDecl::storage("haystack", 3, BufferAccess::ReadOnly, DataType::U32),
1416            BufferDecl::storage("rule_meta", 4, BufferAccess::ReadOnly, DataType::U32),
1417            BufferDecl::storage("transitions", 5, BufferAccess::ReadOnly, DataType::U32),
1418            BufferDecl::storage("accept", 6, BufferAccess::ReadOnly, DataType::U32),
1419        ]),
1420        BatchAutomatonLayout::Combined => buffers.extend([
1421            BufferDecl::storage("haystack", 2, BufferAccess::ReadOnly, DataType::U32),
1422            BufferDecl::storage("transitions", 3, BufferAccess::ReadOnly, DataType::U32),
1423            BufferDecl::storage("output_offsets", 4, BufferAccess::ReadOnly, DataType::U32),
1424            BufferDecl::storage("output_records", 5, BufferAccess::ReadOnly, DataType::U32),
1425            BufferDecl::storage("class_maps", 6, BufferAccess::ReadOnly, DataType::U32),
1426        ]),
1427    }
1428    buffers.extend([
1429        BufferDecl::storage("queue_state", 7, BufferAccess::ReadWrite, DataType::U32).with_count(
1430            dispatcher_abi_u32(QUEUE_STATE_WORDS, "queue-state word count"),
1431        ),
1432        BufferDecl::output("hit_ring", 8, DataType::U32).with_count(hit_ring_words),
1433        BufferDecl::storage("segments", 9, BufferAccess::ReadOnly, DataType::U32),
1434    ]);
1435    buffers
1436}
1437fn batch_program_buffers(hit_capacity: u32) -> Vec<BufferDecl> {
1438    batch_program_buffers_for_layout(hit_capacity, BatchAutomatonLayout::PerRule)
1439}
1440
1441// ─── Combined-AC segmented megakernel ──────────────────────────────────────
1442//
1443// The per-rule path above runs ONE 2-/N-state DFA per (segment, rule) work
1444// item, so a catalog of K patterns multiplies the queue by K: every byte is
1445// re-read K times. The combined path compiles ALL patterns into ONE
1446// Aho-Corasick automaton (`vyre_libs::scan::classic_ac::classic_ac_compile`)
1447// and runs it ONCE per segment: `queue_len = segment_count` (no rule
1448// dimension). Each accepting state emits the SET of pattern ids that match
1449// there via the `output_offsets`/`output_records` flat arrays, so a single
1450// transition read per byte covers every pattern. The window decode, warm-up
1451// prefix, and emit-guard are byte-for-byte the per-rule path's, the SAME
1452// `plan_segments`/`segment_table` geometry and the SAME
1453// `byte_pos >= emit_start` ownership rule the `segmentation.rs`
1454// `combined_segmented_scan` CPU oracle proves equal to a linear
1455// `classic_ac_scan`.
1456
1457/// Buffer layout for the combined-AC segmented megakernel.
1458///
1459/// Mirrors [`batch_program_buffers`] but swaps the four per-rule automaton
1460/// buffers (`class_maps`/`rule_meta`/`transitions`/`accept`) for the three
1461/// combined-AC buffers (`transitions`/`output_offsets`/`output_records`) and
1462/// drops the rule dimension. `segments` stays last so it occupies the final
1463/// positional input slot, exactly as the per-rule path requires.
1464fn combined_batch_program_buffers(hit_capacity: u32) -> Vec<BufferDecl> {
1465    batch_program_buffers_for_layout(hit_capacity, BatchAutomatonLayout::Combined)
1466}
1467
1468/// Width of each combined-AC transition target in the device transition table.
1469///
1470/// `Bits32` is the shipping default, one `u32` per target, indexed
1471/// `transitions[state * num_classes + class]`. `Bits16` packs two targets per
1472/// `u32` word (low half = even flat index, high half = odd; host packer
1473/// [`vyre_runtime::megakernel::rule_catalog::try_pack_u16_transitions_into`]),
1474/// halving the transition table and bytes-per-transaction, the
1475/// large-catalog-scale L1 working-set lever (`docs/GPU_OOM_SEGMENTATION.md`), at the
1476/// cost of an unpack shift/mask in the hot loop. Sound ONLY when every target
1477/// fits `u16` (`state_count <= 65536`); the host packer fails closed otherwise.
1478#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1479pub enum TransitionWidth {
1480    /// One `u32` per transition target (default).
1481    Bits32,
1482    /// Two `u16` targets packed per `u32` word.
1483    Bits16,
1484}
1485
1486/// Build the combined-AC segmented megakernel program.
1487///
1488/// Identical drain loop to [`build_batch_program`] (`forever` + `claim >=
1489/// QUEUE_LEN` Return, Law-10 no under-claim), with the per-rule claim body
1490/// replaced by [`execute_combined_claim_body`]. `queue_len = segment_count`
1491/// (one work item per segment), set by the host. `transition_width` selects the
1492/// device transition-table packing (see [`TransitionWidth`]); the host MUST
1493/// upload a table packed to match.
1494pub(crate) fn build_combined_batch_program(
1495    workgroup_size_x: u32,
1496    hit_capacity: u32,
1497    num_classes: u32,
1498    transition_width: TransitionWidth,
1499) -> Program {
1500    let queue_len = atomic_load_relaxed(
1501        "queue_state",
1502        Expr::u32(dispatcher_abi_u32(
1503            queue_state_word::QUEUE_LEN,
1504            "queue-state length word",
1505        )),
1506    );
1507    let mut loop_body = vec![
1508        Node::let_bind(
1509            "claim",
1510            Expr::atomic_add(
1511                "queue_state",
1512                Expr::u32(dispatcher_abi_u32(
1513                    queue_state_word::HEAD,
1514                    "queue-state head word",
1515                )),
1516                Expr::u32(1),
1517            ),
1518        ),
1519        Node::if_then(Expr::ge(Expr::var("claim"), queue_len), vec![Node::Return]),
1520    ];
1521    loop_body.extend(execute_combined_claim_body(num_classes, transition_width));
1522
1523    Program::wrapped(
1524        combined_batch_program_buffers(hit_capacity),
1525        [workgroup_size_x, 1, 1],
1526        vec![Node::forever(loop_body)],
1527    )
1528}
1529
1530fn segment_window_nodes(segment_index: Expr) -> Vec<Node> {
1531    vec![
1532        Node::let_bind("seg_idx", segment_index),
1533        Node::let_bind(
1534            "seg_base",
1535            Expr::mul(
1536                Expr::var("seg_idx"),
1537                Expr::u32(dispatcher_abi_u32(
1538                    SEGMENT_WORDS,
1539                    "segment table word count",
1540                )),
1541            ),
1542        ),
1543        Node::let_bind("file_idx", Expr::load("segments", Expr::var("seg_base"))),
1544        Node::let_bind(
1545            "scan_start_rel",
1546            Expr::load("segments", Expr::add(Expr::var("seg_base"), Expr::u32(1))),
1547        ),
1548        Node::let_bind(
1549            "emit_start_rel",
1550            Expr::load("segments", Expr::add(Expr::var("seg_base"), Expr::u32(2))),
1551        ),
1552        Node::let_bind(
1553            "emit_end_rel",
1554            Expr::load("segments", Expr::add(Expr::var("seg_base"), Expr::u32(3))),
1555        ),
1556        Node::let_bind(
1557            "metadata_base",
1558            Expr::mul(
1559                Expr::var("file_idx"),
1560                Expr::u32(dispatcher_abi_u32(
1561                    FILE_METADATA_WORDS,
1562                    "file metadata word count",
1563                )),
1564            ),
1565        ),
1566        Node::let_bind(
1567            "layer_idx",
1568            Expr::load(
1569                "file_metadata",
1570                Expr::add(Expr::var("metadata_base"), Expr::u32(3)),
1571            ),
1572        ),
1573        Node::let_bind(
1574            "file_start",
1575            Expr::load("file_offsets", Expr::var("file_idx")),
1576        ),
1577        Node::let_bind(
1578            "scan_start",
1579            Expr::add(Expr::var("file_start"), Expr::var("scan_start_rel")),
1580        ),
1581        Node::let_bind(
1582            "emit_start",
1583            Expr::add(Expr::var("file_start"), Expr::var("emit_start_rel")),
1584        ),
1585        Node::let_bind(
1586            "emit_end",
1587            Expr::add(Expr::var("file_start"), Expr::var("emit_end_rel")),
1588        ),
1589    ]
1590}
1591
1592fn complete_claim_node() -> Node {
1593    Node::let_bind(
1594        "done_prev",
1595        Expr::atomic_add(
1596            "queue_state",
1597            Expr::u32(dispatcher_abi_u32(
1598                queue_state_word::DONE_COUNT,
1599                "queue-state done-count word",
1600            )),
1601            Expr::u32(1),
1602        ),
1603    )
1604}
1605
1606/// Decode one combined-AC claim into its file window and scan it.
1607///
1608/// `queue_len = segment_count`, so the claim IS the segment index directly (no
1609/// `/ rule_count`: there is no rule dimension). The segment row layout
1610/// `[file_idx, scan_start, emit_start, emit_end]` and the absolute-bounds
1611/// arithmetic are identical to [`execute_batch_claim_body`].
1612fn execute_combined_claim_body(num_classes: u32, transition_width: TransitionWidth) -> Vec<Node> {
1613    let mut body = segment_window_nodes(Expr::var("claim"));
1614    body.push(Node::Block(combined_dfa_byte_scanner(
1615        num_classes,
1616        transition_width,
1617    )));
1618    body.push(complete_claim_node());
1619    body
1620}
1621
1622/// Combined Aho-Corasick window scan with per-state multi-emit.
1623///
1624/// Walks the dense combined automaton over `[scan_start, emit_end)` from state
1625/// 0; the `[scan_start, emit_start)` prefix is warm-up (advances state, emits
1626/// nothing). Once `byte_pos >= emit_start`, every pattern id in this state's
1627/// CSR row `output_records[output_offsets[state] .. output_offsets[state+1]]`
1628/// is emitted at `match_offset = byte_pos - file_start`. This is the only
1629/// place the kernel diverges from the per-rule scanner: a CSR multi-emit loop
1630/// instead of a single `accept[state]` flag, mirroring the
1631/// `combined_segmented_scan` CPU oracle exactly.
1632fn combined_dfa_byte_scanner(num_classes: u32, transition_width: TransitionWidth) -> Vec<Node> {
1633    let mut loop_body = vec![
1634        Node::let_bind(
1635            "haystack_word_index",
1636            Expr::div(Expr::var("byte_pos"), Expr::u32(4)),
1637        ),
1638        Node::let_bind(
1639            "haystack_shift",
1640            Expr::mul(Expr::rem(Expr::var("byte_pos"), Expr::u32(4)), Expr::u32(8)),
1641        ),
1642        Node::let_bind(
1643            "byte",
1644            Expr::bitand(
1645                Expr::shr(
1646                    Expr::load("haystack", Expr::var("haystack_word_index")),
1647                    Expr::var("haystack_shift"),
1648                ),
1649                Expr::u32(0xFF),
1650            ),
1651        ),
1652        // Byte-class compressed combined transition (lossless): fold the byte
1653        // through the 256-entry class map, then index the compressed
1654        // `state * num_classes + class` row. `num_classes` is baked as a literal
1655        //: the automaton fixes it, so the pipeline is compiled once per resident
1656        // catalog. Firings are byte-for-byte identical to the dense
1657        // `state * 256 + byte` table.
1658        Node::let_bind("byte_class", Expr::load("class_maps", Expr::var("byte"))),
1659    ];
1660    // The transition read narrows by `transition_width`: Bits32 loads the target
1661    // directly; Bits16 unpacks two targets per word (half the bytes/transaction).
1662    loop_body.extend(combined_transition_read(num_classes, transition_width));
1663    // Emit guard: only positions owned by this window (`byte_pos >= emit_start`;
1664    // the loop bound enforces `byte_pos < emit_end`). Warm-up bytes advance state
1665    // but emit nothing, so adjacent windows tile each file with no double count
1666    // and no miss.
1667    loop_body.push(Node::if_then(
1668        Expr::ge(Expr::var("byte_pos"), Expr::var("emit_start")),
1669        vec![
1670            Node::let_bind(
1671                "out_begin",
1672                Expr::load("output_offsets", Expr::var("state")),
1673            ),
1674            Node::let_bind(
1675                "out_end",
1676                Expr::load(
1677                    "output_offsets",
1678                    Expr::add(Expr::var("state"), Expr::u32(1)),
1679                ),
1680            ),
1681            // Multi-emit: one HitRecord per pattern id accepting at this state.
1682            // `rule_idx` carries the pattern id (the combined automaton's pattern
1683            // == the catalog rule).
1684            Node::loop_for(
1685                "out_idx",
1686                Expr::var("out_begin"),
1687                Expr::var("out_end"),
1688                vec![
1689                    Node::let_bind(
1690                        "rule_idx",
1691                        Expr::load("output_records", Expr::var("out_idx")),
1692                    ),
1693                    Node::Block(record_hit_to_ring()),
1694                ],
1695            ),
1696        ],
1697    ));
1698
1699    vec![
1700        Node::let_bind("state", Expr::u32(0)),
1701        Node::loop_for(
1702            "byte_pos",
1703            Expr::var("scan_start"),
1704            Expr::var("emit_end"),
1705            loop_body,
1706        ),
1707    ]
1708}
1709
1710/// The combined-AC transition step `state := transition(state, byte_class)`,
1711/// emitting the read for the chosen [`TransitionWidth`].
1712///
1713/// `Bits32` loads `transitions[state * num_classes + class]` directly. `Bits16`
1714/// computes that same flat index, loads the packed word at `idx / 2`, and
1715/// extracts the `u16` half selected by `idx & 1`: the EXACT mirror of
1716/// [`vyre_runtime::megakernel::rule_catalog::unpack_u16_transition`], so a
1717/// u16-packed table reproduces the u32 firings byte-for-byte (proven on the GPU
1718/// by the differential conservation test, which runs both widths).
1719fn combined_transition_read(num_classes: u32, transition_width: TransitionWidth) -> Vec<Node> {
1720    let flat_index = Expr::add(
1721        Expr::mul(Expr::var("state"), Expr::u32(num_classes)),
1722        Expr::var("byte_class"),
1723    );
1724    match transition_width {
1725        TransitionWidth::Bits32 => {
1726            vec![Node::assign("state", Expr::load("transitions", flat_index))]
1727        }
1728        TransitionWidth::Bits16 => vec![
1729            Node::let_bind("trans_idx", flat_index),
1730            Node::let_bind(
1731                "trans_word",
1732                Expr::load(
1733                    "transitions",
1734                    Expr::div(Expr::var("trans_idx"), Expr::u32(2)),
1735                ),
1736            ),
1737            Node::assign(
1738                "state",
1739                Expr::bitand(
1740                    Expr::shr(
1741                        Expr::var("trans_word"),
1742                        Expr::mul(
1743                            Expr::rem(Expr::var("trans_idx"), Expr::u32(2)),
1744                            Expr::u32(16),
1745                        ),
1746                    ),
1747                    Expr::u32(0xFFFF),
1748                ),
1749            ),
1750        ],
1751    }
1752}
1753
1754fn execute_batch_claim_body(hit_writer: BatchHitWriter) -> Vec<Node> {
1755    let rule_count = atomic_load_relaxed(
1756        "queue_state",
1757        Expr::u32(dispatcher_abi_u32(
1758            queue_state_word::RULE_COUNT,
1759            "queue-state rule-count word",
1760        )),
1761    );
1762    let mut body = vec![Node::let_bind("rule_count", rule_count)];
1763    body.extend(segment_window_nodes(Expr::div(
1764        Expr::var("claim"),
1765        Expr::var("rule_count"),
1766    )));
1767    body.extend([
1768        Node::let_bind(
1769            "rule_idx",
1770            Expr::rem(Expr::var("claim"), Expr::var("rule_count")),
1771        ),
1772        Node::let_bind(
1773            "rule_base",
1774            Expr::mul(
1775                Expr::var("rule_idx"),
1776                Expr::u32(dispatcher_abi_u32(
1777                    RULE_META_WORDS,
1778                    "rule metadata word count",
1779                )),
1780            ),
1781        ),
1782        Node::let_bind(
1783            "transition_base",
1784            Expr::load("rule_meta", Expr::var("rule_base")),
1785        ),
1786        Node::let_bind(
1787            "accept_base",
1788            Expr::load("rule_meta", Expr::add(Expr::var("rule_base"), Expr::u32(1))),
1789        ),
1790        Node::let_bind(
1791            "class_map_base",
1792            Expr::load("rule_meta", Expr::add(Expr::var("rule_base"), Expr::u32(3))),
1793        ),
1794        Node::let_bind(
1795            "num_classes",
1796            Expr::load("rule_meta", Expr::add(Expr::var("rule_base"), Expr::u32(4))),
1797        ),
1798        Node::Block(dfa_byte_scanner(hit_writer)),
1799        complete_claim_node(),
1800    ]);
1801    body
1802}
1803
1804fn dfa_byte_scanner(hit_writer: BatchHitWriter) -> Vec<Node> {
1805    vec![
1806        Node::let_bind("state", Expr::u32(0)),
1807        // Scan the window `[scan_start, emit_end)` from state 0. The
1808        // `[scan_start, emit_start)` prefix is DFA warm-up, it advances the
1809        // state but emits nothing (the emit guard below). For the dense default
1810        // `scan_start == emit_start == file_start`, so the loop is the whole file
1811        // with no warm-up (identical to the pre-segmentation scan).
1812        Node::loop_for(
1813            "byte_pos",
1814            Expr::var("scan_start"),
1815            Expr::var("emit_end"),
1816            vec![
1817                Node::let_bind(
1818                    "haystack_word_index",
1819                    Expr::div(Expr::var("byte_pos"), Expr::u32(4)),
1820                ),
1821                Node::let_bind(
1822                    "haystack_shift",
1823                    Expr::mul(Expr::rem(Expr::var("byte_pos"), Expr::u32(4)), Expr::u32(8)),
1824                ),
1825                Node::let_bind(
1826                    "byte",
1827                    Expr::bitand(
1828                        Expr::shr(
1829                            Expr::load("haystack", Expr::var("haystack_word_index")),
1830                            Expr::var("haystack_shift"),
1831                        ),
1832                        Expr::u32(0xFF),
1833                    ),
1834                ),
1835                // Byte-class compressed transition load (lossless): fold the
1836                // byte through this rule's 256-entry class map, then index the
1837                // compressed `state * num_classes + class` row. Bytes that share
1838                // a transition column across every state collapse to one class,
1839                // shrinking each per-state row from 256 words to `num_classes`
1840                // words. Firings are byte-for-byte identical to the dense
1841                // `state * 256 + byte` table (proved in the CPU parity tests).
1842                Node::let_bind(
1843                    "byte_class",
1844                    Expr::load(
1845                        "class_maps",
1846                        Expr::add(Expr::var("class_map_base"), Expr::var("byte")),
1847                    ),
1848                ),
1849                Node::assign(
1850                    "state",
1851                    Expr::load(
1852                        "transitions",
1853                        Expr::add(
1854                            Expr::var("transition_base"),
1855                            Expr::add(
1856                                Expr::mul(Expr::var("state"), Expr::var("num_classes")),
1857                                Expr::var("byte_class"),
1858                            ),
1859                        ),
1860                    ),
1861                ),
1862                Node::let_bind(
1863                    "accepting",
1864                    Expr::load(
1865                        "accept",
1866                        Expr::add(Expr::var("accept_base"), Expr::var("state")),
1867                    ),
1868                ),
1869                // Emit guard mirrors the CPU parity oracle (`segmentation.rs`):
1870                // a match is owned by this window iff its end offset lies in
1871                // `[emit_start, emit_end)`. The loop bound already enforces
1872                // `byte_pos < emit_end`; the remaining condition `end > emit_start`
1873                // (end = byte_pos + 1) is exactly `byte_pos >= emit_start`. Bytes in
1874                // the warm-up prefix (`byte_pos < emit_start`) advance state but
1875                // never emit, so adjacent windows tile each file with no double
1876                // count and no miss.
1877                Node::let_bind(
1878                    "is_hit",
1879                    Expr::and(
1880                        Expr::ne(Expr::var("accepting"), Expr::u32(0)),
1881                        Expr::ge(Expr::var("byte_pos"), Expr::var("emit_start")),
1882                    ),
1883                ),
1884                hit_writer_node(hit_writer),
1885            ],
1886        ),
1887    ]
1888}
1889
1890/// Caller-owns-hit-storage counters for one combined-AC dispatch.
1891#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1892pub struct CombinedDispatchSummary {
1893    /// Sparse hit count written by the device (clamped to `hit_capacity`).
1894    pub hit_count: u32,
1895    /// Matches produced BEYOND `hit_capacity` and therefore DROPPED from the
1896    /// ring (raw atomic head minus capacity). `> 0` means the hit set is
1897    /// INCOMPLETE, a recall-critical overflow the caller must recover, never
1898    /// treat as complete (Law 10). Zero on a healthy dispatch.
1899    pub dropped_hits: u32,
1900    /// Segments the device finished (`DONE_COUNT`).
1901    pub items_processed: u32,
1902    /// Wall-clock GPU execution time.
1903    pub wall_time: Duration,
1904}
1905
1906/// Persistent dispatcher for the combined-AC segmented megakernel.
1907///
1908/// The combined twin of [`BatchDispatcher`] minus the per-rule catalog
1909/// machinery (no fingerprints, no `rule_meta`/`transitions`/`accept` upload):
1910/// the automaton is resident in the [`CombinedBatch`]. It compiles the
1911/// combined persistent program (the backend pipeline cache dedups the WGSL
1912/// compile by program+adapter+config), drains the whole `segment_count` queue,
1913/// and reads back the sparse hit ring with the SAME Law-10 overflow +
1914/// incomplete-drain guards as the per-rule path.
1915pub struct CombinedDispatcher {
1916    backend: WgpuBackend,
1917    config: BatchDispatchConfig,
1918    queue_state_bytes: Vec<u8>,
1919    hit_bytes: Vec<u8>,
1920}
1921
1922impl CombinedDispatcher {
1923    /// Build a combined dispatcher over a live backend.
1924    #[must_use]
1925    pub fn new(backend: WgpuBackend, config: BatchDispatchConfig) -> Self {
1926        Self {
1927            backend,
1928            config,
1929            queue_state_bytes: Vec::new(),
1930            hit_bytes: Vec::new(),
1931        }
1932    }
1933
1934    /// Dispatch the combined automaton over every segment of `batch`,
1935    /// compacting decoded hits into caller-owned `hits`.
1936    ///
1937    /// # Errors
1938    ///
1939    /// Returns [`PipelineError`] on pipeline compilation failure, dispatch
1940    /// timeout, an incomplete drain (loud, never a silent partial), or readback
1941    /// failure.
1942    pub fn dispatch_into(
1943        &mut self,
1944        batch: &CombinedBatch,
1945        hits: &mut Vec<HitRecord>,
1946    ) -> Result<CombinedDispatchSummary, PipelineError> {
1947        let queue_len = batch.queue_len();
1948        if queue_len == 0 {
1949            // Every file empty ⇒ no segments ⇒ nothing to scan.
1950            hits.clear();
1951            return Ok(CombinedDispatchSummary {
1952                hit_count: 0,
1953                dropped_hits: 0,
1954                items_processed: 0,
1955                wall_time: Duration::ZERO,
1956            });
1957        }
1958        let worker_groups = if self.config.worker_groups != 0 {
1959            self.config.worker_groups
1960        } else {
1961            self.config
1962                .launch_recommendation(self.backend.device_limits(), queue_len)?
1963                .worker_groups
1964        };
1965        let program = build_combined_batch_program(
1966            self.config.workgroup_size_x,
1967            batch.hit_capacity(),
1968            batch.num_classes(),
1969            batch.transition_width(),
1970        );
1971        let pipeline = self
1972            .backend
1973            .compile_persistent(&program, &DispatchConfig::default())?;
1974        batch.reset_queue_state();
1975
1976        let inputs = batch.input_buffers();
1977        let outputs = batch.output_buffers();
1978        let start = Instant::now();
1979        pipeline.dispatch_persistent_borrowed(&inputs, &outputs, None, [worker_groups, 1, 1])?;
1980        let (device, queue) = &*self.backend.device_queue();
1981        wait_for_persistent_dispatch(device, start, self.config.timeout)?;
1982        let wall_time = start.elapsed();
1983
1984        self.queue_state_bytes.clear();
1985        let queue_state_readback_bytes = batch_fixed_resident_overhead_bytes();
1986        batch.queue_state().readback_prefix(
1987            device,
1988            queue,
1989            queue_state_readback_bytes,
1990            &mut self.queue_state_bytes,
1991        )?;
1992        let queue_state_word_count =
1993            validate_u32_readback_words(&self.queue_state_bytes, "queue-state")?;
1994        if queue_state_word_count < QUEUE_STATE_WORDS {
1995            return Err(PipelineError::Backend(format!(
1996                "queue-state readback exposed {} words, expected at least {}. Fix: keep the queue-state buffer sized for every control word.",
1997                queue_state_word_count, QUEUE_STATE_WORDS
1998            )));
1999        }
2000        let raw_hit_head = read_u32_word(
2001            &self.queue_state_bytes,
2002            "queue-state",
2003            queue_state_word::HIT_HEAD,
2004        )?;
2005        let (hit_count, dropped_hits) = split_hit_overflow(raw_hit_head, batch.hit_capacity());
2006        let items_processed = read_u32_word(
2007            &self.queue_state_bytes,
2008            "queue-state",
2009            queue_state_word::DONE_COUNT,
2010        )?;
2011        // Fail-closed drain guard (Law 10): HEAD < queue_len ⇒ the dispatch was
2012        // cut short, leaving segments `[HEAD, queue_len)` unscanned with their
2013        // matches missing and `dropped_hits == 0`: an INVISIBLE recall loss.
2014        let claims_attempted = read_u32_word(
2015            &self.queue_state_bytes,
2016            "queue-state",
2017            queue_state_word::HEAD,
2018        )?;
2019        if claims_attempted < queue_len {
2020            return Err(PipelineError::DrainIncomplete {
2021                descriptor: "combined megakernel",
2022                claimed: claims_attempted,
2023                expected: queue_len,
2024                unit: "segments",
2025            });
2026        }
2027
2028        self.hit_bytes.clear();
2029        let hit_readback_bytes = u64::from(hit_count)
2030            .checked_mul(dispatcher_usize_to_u64(
2031                HIT_RECORD_WORDS,
2032                "hit-record word count",
2033            ))
2034            .and_then(|words| {
2035                words.checked_mul(dispatcher_usize_to_u64(
2036                    std::mem::size_of::<u32>(),
2037                    "u32 byte width",
2038                ))
2039            })
2040            .ok_or_else(|| {
2041                PipelineError::Backend(
2042                    "combined hit-ring readback length overflowed u64. Fix: reduce hit_capacity or shard the batch."
2043                        .to_string(),
2044                )
2045            })?;
2046        batch
2047            .hit_ring()
2048            .readback_prefix(device, queue, hit_readback_bytes, &mut self.hit_bytes)?;
2049        decode_hits_from_readback_into(&self.hit_bytes, hit_count, hits)?;
2050
2051        Ok(CombinedDispatchSummary {
2052            hit_count,
2053            dropped_hits,
2054            items_processed,
2055            wall_time,
2056        })
2057    }
2058
2059    /// Measure-based `seg_len` selection for THIS device. Re-tiles `batch` at
2060    /// each candidate window width, keeping the caller's proven
2061    /// `overlap = max_pattern_len`, so correctness is GEOMETRY-INVARIANT and only
2062    /// throughput changes, times a warm best-of-`reps` dispatch, and returns the
2063    /// FASTEST geometry whose dispatch was COMPLETE (clean drain, `dropped_hits
2064    /// == 0`). No per-call oracle is needed: every candidate shares the caller's
2065    /// sound overlap, so all COMPLETE dispatches yield the identical hit set and
2066    /// this only optimizes speed among them. An under-claiming (drain-incomplete)
2067    /// or ring-overflowing geometry is recorded but EXCLUDED, never silently
2068    /// selected; if no candidate dispatches completely the call FAILS CLOSED.
2069    ///
2070    /// This is the autoroute-honest way to hold the GPU win across devices: the
2071    /// per-device optimum (~128 on an RTX 5090, coarser on a low-core laptop) is
2072    /// MEASURED here, not assumed, and every candidate's measurement is returned
2073    /// so the decision inputs are visible. On success `batch` is left re-tiled at
2074    /// the winning geometry, ready to dispatch.
2075    ///
2076    /// # Errors
2077    ///
2078    /// Returns [`PipelineError`] when `candidates` is empty, when a dispatch fails
2079    /// for a non-geometry reason, or (fail closed) when NO candidate dispatched
2080    /// completely.
2081    pub fn calibrate_seg_len(
2082        &mut self,
2083        batch: &mut CombinedBatch,
2084        candidates: &[u32],
2085        reps: u32,
2086    ) -> Result<SegLenCalibration, PipelineError> {
2087        if candidates.is_empty() {
2088            return Err(PipelineError::Backend(
2089                "seg_len calibration needs at least one candidate geometry; pass \
2090                 DEFAULT_SEG_LEN_CANDIDATES or an operator set."
2091                    .to_string(),
2092            ));
2093        }
2094        let reps = reps.max(1);
2095        let mut scratch: Vec<HitRecord> = Vec::new();
2096        let mut measurements = Vec::with_capacity(candidates.len());
2097        for &seg_len in candidates {
2098            batch.set_segmentation(seg_len)?;
2099            // Warm once: compile + first-touch out of the timing. A warm-up that
2100            // under-claims is fine (the timed loop below records it as incomplete).
2101            let _ = self.dispatch_into(batch, &mut scratch);
2102            let mut best_wall = Duration::MAX;
2103            let mut dropped_hits = 0u32;
2104            let mut complete = true;
2105            for _ in 0..reps {
2106                match self.dispatch_into(batch, &mut scratch) {
2107                    Ok(summary) => {
2108                        best_wall = best_wall.min(summary.wall_time);
2109                        dropped_hits = summary.dropped_hits;
2110                        if summary.dropped_hits != 0 {
2111                            // Hit ring overflowed: this dispatch's set is INCOMPLETE,
2112                            // so its time is not a valid complete-scan measurement.
2113                            complete = false;
2114                        }
2115                    }
2116                    // Drain-incomplete = the geometry could not exhaust the queue in
2117                    // the timeout (loud under-claim). Record + exclude, never select.
2118                    // Matched on the typed predicate, not the message text, so a
2119                    // wording change can never silently turn this into a hard abort.
2120                    Err(e) if e.is_drain_incomplete() => {
2121                        complete = false;
2122                        break;
2123                    }
2124                    Err(other) => return Err(other),
2125                }
2126            }
2127            if best_wall == Duration::MAX {
2128                // Never produced a timed Ok (immediate under-claim every rep).
2129                best_wall = Duration::ZERO;
2130                complete = false;
2131            }
2132            measurements.push(SegLenMeasurement {
2133                seg_len,
2134                wall_time: best_wall,
2135                dropped_hits,
2136                complete,
2137            });
2138        }
2139        let chosen = select_fastest_complete(&measurements)?;
2140        // Leave the batch at the winning geometry so the caller dispatches fast.
2141        batch.set_segmentation(chosen)?;
2142        Ok(SegLenCalibration {
2143            chosen,
2144            measurements,
2145        })
2146    }
2147}
2148
2149/// A vetted default `seg_len` candidate set for [`CombinedDispatcher::calibrate_seg_len`],
2150/// spanning the saturation curve from coarse (fewer segments, best on low-core
2151/// devices) to fine (more parallel windows, best on high-core devices), so
2152/// calibration adapts to the host instead of assuming one optimum. Operators may
2153/// pass their own set. Deliberately excludes `u32::MAX` (whole-file): that is a
2154/// correctness floor, never a throughput candidate (see `CombinedBatch::upload`).
2155pub const DEFAULT_SEG_LEN_CANDIDATES: &[u32] = &[4096, 2048, 1024, 512, 256, 128, 64];
2156
2157/// One geometry's measurement during `seg_len` calibration: the candidate window
2158/// width, its best wall-clock dispatch time over the timed reps, the hit-ring
2159/// overflow count, and whether the dispatch was COMPLETE (clean drain AND
2160/// `dropped_hits == 0`). Only `complete` geometries are eligible to win.
2161#[derive(Debug, Clone)]
2162pub struct SegLenMeasurement {
2163    /// The candidate per-segment owned width that was measured.
2164    pub seg_len: u32,
2165    /// Best (minimum) wall-clock dispatch time observed over the timed reps;
2166    /// `Duration::ZERO` for a geometry that never produced a timed complete run.
2167    pub wall_time: Duration,
2168    /// Hits the GPU emitted beyond `hit_capacity` (ring overflow); non-zero marks
2169    /// the measurement INCOMPLETE.
2170    pub dropped_hits: u32,
2171    /// True iff every timed rep drained cleanly with no dropped hits, the only
2172    /// state in which `wall_time` represents a full, conserving scan.
2173    pub complete: bool,
2174}
2175
2176/// The result of [`CombinedDispatcher::calibrate_seg_len`]: the fastest COMPLETE
2177/// geometry plus EVERY candidate's measurement, so the selection's decision
2178/// inputs are fully visible to the operator (never a silent pick).
2179#[derive(Debug, Clone)]
2180pub struct SegLenCalibration {
2181    /// The winning per-segment owned width; the batch is left re-tiled here.
2182    pub chosen: u32,
2183    /// Per-candidate measurements in the order the candidates were supplied.
2184    pub measurements: Vec<SegLenMeasurement>,
2185}
2186
2187/// Pick the COMPLETE geometry with the smallest wall time, breaking ties toward
2188/// the COARSER (larger) `seg_len`: fewer segments means less queue-drain and
2189/// warm-up overhead at equal measured speed. Fails closed when no candidate is
2190/// complete: returns an error rather than an unsound or incomplete geometry.
2191fn select_fastest_complete(measurements: &[SegLenMeasurement]) -> Result<u32, PipelineError> {
2192    measurements
2193        .iter()
2194        .filter(|m| m.complete)
2195        .min_by(|a, b| {
2196            a.wall_time
2197                .cmp(&b.wall_time)
2198                .then_with(|| b.seg_len.cmp(&a.seg_len))
2199        })
2200        .map(|m| m.seg_len)
2201        .ok_or_else(|| {
2202            PipelineError::Backend(
2203                "seg_len calibration found no complete geometry: every candidate \
2204                 under-claimed (drain incomplete) or overflowed the hit ring. Fix: \
2205                 raise BatchDispatchConfig.timeout and/or hit_capacity, or pass \
2206                 coarser candidates (larger seg_len = fewer segments to drain)."
2207                    .to_string(),
2208            )
2209        })
2210}
2211
2212fn hit_writer_node(hit_writer: BatchHitWriter) -> Node {
2213    match hit_writer {
2214        BatchHitWriter::HierarchicalSubgroup => {
2215            Node::Block(record_hit_to_ring_hierarchical("is_hit"))
2216        }
2217        BatchHitWriter::Auto | BatchHitWriter::Scalar => {
2218            Node::if_then(Expr::var("is_hit"), record_hit_to_ring())
2219        }
2220    }
2221}
2222
2223fn record_hit_to_ring() -> Vec<Node> {
2224    vec![
2225        Node::let_bind(
2226            "hit_slot",
2227            Expr::atomic_add(
2228                "queue_state",
2229                Expr::u32(dispatcher_abi_u32(
2230                    queue_state_word::HIT_HEAD,
2231                    "queue-state hit-head word",
2232                )),
2233                Expr::u32(1),
2234            ),
2235        ),
2236        Node::if_then(
2237            Expr::lt(
2238                Expr::var("hit_slot"),
2239                atomic_load_relaxed(
2240                    "queue_state",
2241                    Expr::u32(dispatcher_abi_u32(
2242                        queue_state_word::HIT_CAPACITY,
2243                        "queue-state hit-capacity word",
2244                    )),
2245                ),
2246            ),
2247            vec![
2248                Node::let_bind("hit_base", Expr::mul(Expr::var("hit_slot"), Expr::u32(4))),
2249                Node::store("hit_ring", Expr::var("hit_base"), Expr::var("file_idx")),
2250                Node::store(
2251                    "hit_ring",
2252                    Expr::add(Expr::var("hit_base"), Expr::u32(1)),
2253                    Expr::var("rule_idx"),
2254                ),
2255                Node::store(
2256                    "hit_ring",
2257                    Expr::add(Expr::var("hit_base"), Expr::u32(2)),
2258                    Expr::var("layer_idx"),
2259                ),
2260                Node::store(
2261                    "hit_ring",
2262                    Expr::add(Expr::var("hit_base"), Expr::u32(3)),
2263                    Expr::sub(Expr::var("byte_pos"), Expr::var("file_start")),
2264                ),
2265            ],
2266        ),
2267    ]
2268}
2269
2270/// Split the device's raw atomic hit-head into `(readable, dropped)`.
2271///
2272/// The kernel increments `HIT_HEAD` for EVERY match but only writes ring slots
2273/// below `hit_capacity`, so a `raw_head > capacity` means `raw_head - capacity`
2274/// matches were produced-but-dropped. `readable` is what can be decoded from the
2275/// ring (`min(raw_head, capacity)`); `dropped` is the overflow the caller must
2276/// recover. Pure so the overflow accounting is unit-tested without a device.
2277const fn split_hit_overflow(raw_head: u32, capacity: u32) -> (u32, u32) {
2278    if raw_head > capacity {
2279        (capacity, raw_head - capacity)
2280    } else {
2281        (raw_head, 0)
2282    }
2283}
2284
2285#[cfg(test)]
2286fn decode_hits_from_readback(
2287    bytes: &[u8],
2288    hit_count: u32,
2289) -> Result<Vec<HitRecord>, PipelineError> {
2290    let mut hits = Vec::new();
2291    decode_hits_from_readback_into(bytes, hit_count, &mut hits)?;
2292    Ok(hits)
2293}
2294
2295fn decode_hits_from_readback_into(
2296    bytes: &[u8],
2297    hit_count: u32,
2298    hits: &mut Vec<HitRecord>,
2299) -> Result<(), PipelineError> {
2300    let word_count = validate_u32_readback_words(bytes, "hit-ring")?;
2301    let needed_words = usize::try_from(hit_count)
2302        .ok()
2303        .and_then(|count| count.checked_mul(4))
2304        .ok_or_else(|| PipelineError::Backend("hit-count overflowed usize".to_string()))?;
2305    if word_count < needed_words {
2306        return Err(PipelineError::Backend(format!(
2307            "hit-ring exposed {} words, expected at least {needed_words}. Fix: size the sparse hit ring for the configured hit_capacity.",
2308            word_count
2309        )));
2310    }
2311    let needed_bytes = needed_words
2312        .checked_mul(std::mem::size_of::<u32>())
2313        .ok_or_else(|| PipelineError::Backend(
2314            "hit-ring readback byte count overflowed usize. Fix: reduce hit_capacity or shard the batch."
2315                .to_string(),
2316        ))?;
2317    let hit_count = usize::try_from(hit_count).map_err(|source| {
2318        PipelineError::Backend(format!(
2319            "hit count cannot fit usize for host decode: {source}. Fix: reduce hit_capacity or run on a supported host pointer width."
2320        ))
2321    })?;
2322    let same_len = hits.len() == hit_count;
2323    if !same_len {
2324        hits.clear();
2325    }
2326    if hits.capacity() < hit_count {
2327        hits.try_reserve_exact(hit_count - hits.len())
2328            .map_err(|source| {
2329                PipelineError::Backend(format!(
2330                    "hit-ring decode could not reserve {hit_count} HitRecord slots: {source}. Fix: lower hit_capacity or shard the batch."
2331                ))
2332            })?;
2333    }
2334    if cfg!(target_endian = "little") {
2335        let record_bytes = std::mem::size_of::<HitRecord>();
2336        let expected_record_bytes = HIT_RECORD_WORDS * std::mem::size_of::<u32>();
2337        if record_bytes != expected_record_bytes {
2338            return Err(PipelineError::Backend(format!(
2339                "hit-ring host record layout is {record_bytes} bytes, expected {expected_record_bytes}. Fix: keep HitRecord as four packed u32 words."
2340            )));
2341        }
2342        if hit_count != 0 {
2343            let records: &[HitRecord] =
2344                bytemuck::try_cast_slice(&bytes[..needed_bytes]).map_err(|source| {
2345                    PipelineError::Backend(format!(
2346                        "hit-ring readback bytes were not aligned as HitRecord records: {source}. Fix: keep the hit ring byte layout aligned to four u32 words."
2347                    ))
2348                })?;
2349            if same_len {
2350                hits.copy_from_slice(records);
2351            } else {
2352                hits.extend_from_slice(records);
2353            }
2354        }
2355        return Ok(());
2356    }
2357    for (index, chunk) in bytes[..needed_bytes]
2358        .chunks_exact(HIT_RECORD_WORDS * std::mem::size_of::<u32>())
2359        .enumerate()
2360    {
2361        let record = HitRecord {
2362            file_idx: u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]),
2363            rule_idx: u32::from_le_bytes([chunk[4], chunk[5], chunk[6], chunk[7]]),
2364            layer_idx: u32::from_le_bytes([chunk[8], chunk[9], chunk[10], chunk[11]]),
2365            match_offset: u32::from_le_bytes([chunk[12], chunk[13], chunk[14], chunk[15]]),
2366        };
2367        if same_len {
2368            hits[index] = record;
2369        } else {
2370            hits.push(record);
2371        }
2372    }
2373    Ok(())
2374}
2375
2376#[cfg(test)]
2377mod tests {
2378    use super::*;
2379
2380    fn measurement(
2381        seg_len: u32,
2382        micros: u64,
2383        dropped_hits: u32,
2384        complete: bool,
2385    ) -> SegLenMeasurement {
2386        SegLenMeasurement {
2387            seg_len,
2388            wall_time: Duration::from_micros(micros),
2389            dropped_hits,
2390            complete,
2391        }
2392    }
2393
2394    /// The calibrator picks the fastest COMPLETE geometry, and must NOT pick a
2395    /// faster-but-incomplete one. A geometry that drained faster only because it
2396    /// dropped hits (ring overflow) is a false speedup; selecting it would ship a
2397    /// silently-truncated scan (Law 10). Here seg_len=64 is the fastest wall time
2398    /// but incomplete, so the fastest COMPLETE geometry (128) must win.
2399    #[test]
2400    fn calibration_selects_fastest_complete_not_faster_incomplete() {
2401        let measurements = vec![
2402            measurement(512, 900, 0, true),
2403            measurement(256, 700, 0, true),
2404            measurement(128, 400, 0, true),  // fastest complete
2405            measurement(64, 300, 12, false), // faster wall, but dropped 12 ⇒ excluded
2406        ];
2407        assert_eq!(
2408            select_fastest_complete(&measurements).expect("a complete geometry exists"),
2409            128,
2410            "must select the fastest COMPLETE geometry (128), never the faster-but-truncated 64"
2411        );
2412    }
2413
2414    /// No complete geometry ⇒ fail closed with an actionable error, never return
2415    /// an unsound/incomplete geometry as if it were a valid pick.
2416    #[test]
2417    fn calibration_fails_closed_when_every_geometry_is_incomplete() {
2418        let measurements = vec![measurement(128, 0, 5, false), measurement(64, 0, 9, false)];
2419        let err = select_fastest_complete(&measurements)
2420            .expect_err("no complete geometry must error, not silently pick one");
2421        assert!(
2422            err.to_string().contains("no complete geometry"),
2423            "fail-closed error must name the cause; got: {err}"
2424        );
2425    }
2426
2427    /// Equal measured speed ⇒ prefer the COARSER (larger seg_len) geometry: fewer
2428    /// segments means less queue-drain and per-window warm-up overhead, and is
2429    /// less likely to under-claim on a slower device.
2430    #[test]
2431    fn calibration_tie_breaks_toward_coarser_geometry() {
2432        let measurements = vec![
2433            measurement(256, 500, 0, true),
2434            measurement(128, 500, 0, true),
2435            measurement(64, 500, 0, true),
2436        ];
2437        assert_eq!(
2438            select_fastest_complete(&measurements).expect("complete geometries exist"),
2439            256,
2440            "equal wall time ⇒ pick the coarsest (256), the fewest-segments geometry"
2441        );
2442    }
2443
2444    /// The shipped default candidate set must be non-empty, strictly descending
2445    /// (coarse→fine, the order calibration walks the saturation curve), and must
2446    /// exclude the whole-file `u32::MAX` correctness floor, it is never a
2447    /// throughput candidate.
2448    #[test]
2449    fn default_seg_len_candidates_are_descending_and_exclude_whole_file() {
2450        assert!(!DEFAULT_SEG_LEN_CANDIDATES.is_empty());
2451        assert!(
2452            !DEFAULT_SEG_LEN_CANDIDATES.contains(&u32::MAX),
2453            "whole-file is a correctness floor, never a calibration candidate"
2454        );
2455        for pair in DEFAULT_SEG_LEN_CANDIDATES.windows(2) {
2456            assert!(
2457                pair[0] > pair[1],
2458                "candidates must be strictly descending coarse→fine; {} !> {}",
2459                pair[0],
2460                pair[1]
2461            );
2462        }
2463        assert_eq!(
2464            *DEFAULT_SEG_LEN_CANDIDATES.last().unwrap(),
2465            64,
2466            "finest default candidate is 64 (overlap-waste dominates below it for MiB inputs)"
2467        );
2468    }
2469
2470    /// The combined-AC segmented program must lower through the REAL WGSL
2471    /// pipeline (`descriptor_gate::validate_and_analyze` → `vyre_emit_naga::emit`
2472    /// → Naga validation → WGSL writer). Naga's validator rejects malformed
2473    /// control flow, type errors, or invalid buffer access, so a successful
2474    /// lower is a genuine proof the multi-emit nested loops + combined
2475    /// transition read are valid GPU code, not a shape check. We additionally
2476    /// assert every combined-automaton buffer and the multi-emit payload buffer
2477    /// survive into the emitted WGSL (Naga names globals after the BufferDecl),
2478    /// so the kernel actually reads the combined tables rather than lowering to
2479    /// a stripped no-op.
2480    #[test]
2481    fn combined_batch_program_lowers_to_valid_wgsl_referencing_combined_tables() {
2482        // BOTH transition widths must lower through the real Naga pipeline: the
2483        // u16 path adds a div/shr/mask unpack in the hot loop, so its validity is
2484        // not implied by the u32 path's.
2485        for width in [TransitionWidth::Bits32, TransitionWidth::Bits16] {
2486            let program = build_combined_batch_program(64, 1024, 40, width);
2487            let wgsl = crate::emit::lower(&program).unwrap_or_else(|e| {
2488                panic!("Fix: combined-AC {width:?} program must lower to valid WGSL: {e:?}")
2489            });
2490            for needle in [
2491                "transitions",
2492                "output_offsets",
2493                "output_records",
2494                "class_maps",
2495                "segments",
2496                "hit_ring",
2497                "file_offsets",
2498            ] {
2499                assert!(
2500                    wgsl.contains(needle),
2501                    "emitted WGSL ({width:?}) must reference the `{needle}` buffer; the combined \
2502                     kernel read it in IR but it vanished from the shader (got {} bytes of WGSL)",
2503                    wgsl.len()
2504                );
2505            }
2506        }
2507    }
2508
2509    /// Pin the combined-AC ABI: nine buffers in the exact binding order the host
2510    /// `CombinedBatch` upload must mirror, with `hit_ring` the sole output at
2511    /// binding 7 and `segments` last (final positional input). A drift here
2512    /// silently misbinds the automaton tables.
2513    #[test]
2514    fn combined_batch_program_buffer_abi_is_pinned() {
2515        let buffers = combined_batch_program_buffers(1024);
2516        let layout: Vec<(&str, u32, bool, bool)> = buffers
2517            .iter()
2518            .map(|b| {
2519                (
2520                    b.name(),
2521                    b.binding(),
2522                    b.is_output(),
2523                    matches!(b.access(), BufferAccess::ReadWrite),
2524                )
2525            })
2526            .collect();
2527        assert_eq!(
2528            layout,
2529            vec![
2530                ("file_offsets", 0, false, false),
2531                ("file_metadata", 1, false, false),
2532                ("haystack", 2, false, false),
2533                ("transitions", 3, false, false),
2534                ("output_offsets", 4, false, false),
2535                ("output_records", 5, false, false),
2536                ("class_maps", 6, false, false),
2537                ("queue_state", 7, false, true),
2538                // hit_ring is a writable output ⇒ ReadWrite access.
2539                ("hit_ring", 8, true, true),
2540                ("segments", 9, false, false),
2541            ],
2542            "combined-AC buffer ABI drifted; the host CombinedBatch upload binds by this order"
2543        );
2544    }
2545
2546    /// Cross-OS shader VALIDITY for the combined segmented megakernel, the
2547    /// program that wins 8.81× vs Hyperscan and runs live on Vulkan (Linux) and
2548    /// DX12 (Windows). Those runs implicitly prove SPIR-V and HLSL validity, but
2549    /// macOS/Metal is hardware-blocked, leaving its shader validity unproven.
2550    /// `vyre-emit-metal` lowers the descriptor through `naga::back::msl`, whose
2551    /// writer runs `naga::valid::Validator` first, a successful emit is therefore
2552    /// a genuine proof the megakernel is valid Metal Shading Language, verifiable
2553    /// WITHOUT an Apple GPU (the throughput-win on Metal stays separately
2554    /// hardware-blocked). SPIR-V is emitted explicitly too rather than relying on
2555    /// the live Vulkan run. BOTH transition widths are checked: the u16 path adds
2556    /// a div/shr/mask unpack in the hot loop, so its Metal/SPIR-V validity is not
2557    /// implied by the u32 path's.
2558    #[test]
2559    fn combined_megakernel_lowers_to_valid_msl_and_spirv_for_all_os() {
2560        for width in [TransitionWidth::Bits32, TransitionWidth::Bits16] {
2561            let program = build_combined_batch_program(64, 1024, 40, width);
2562            let lowered = vyre_lower::lower_for_emit(&program).unwrap_or_else(|e| {
2563                panic!(
2564                    "Fix: combined megakernel ({width:?}) must lower to a KernelDescriptor: {e:?}"
2565                )
2566            });
2567
2568            // The full combined program (not a stripped kernel) must reach the
2569            // backend: its ten-buffer ABI survives into the descriptor bindings.
2570            let slot_names: Vec<&str> = lowered
2571                .descriptor
2572                .bindings
2573                .slots
2574                .iter()
2575                .map(|s| s.name.as_str())
2576                .collect();
2577            for needle in [
2578                "file_offsets",
2579                "file_metadata",
2580                "haystack",
2581                "transitions",
2582                "output_offsets",
2583                "output_records",
2584                "class_maps",
2585                "queue_state",
2586                "hit_ring",
2587                "segments",
2588            ] {
2589                assert!(
2590                    slot_names.contains(&needle),
2591                    "combined megakernel descriptor ({width:?}) is missing the `{needle}` buffer; \
2592                     a stripped descriptor would silently emit a no-op kernel. Got {slot_names:?}"
2593                );
2594            }
2595
2596            // macOS/Metal: a successful emit ⇒ naga validated the module ⇒ the
2597            // kernel is valid MSL on Apple GPUs.
2598            let msl = vyre_emit_metal::emit(&lowered.descriptor).unwrap_or_else(|e| {
2599                panic!(
2600                    "Fix: combined megakernel ({width:?}) must lower to valid MSL for macOS/Metal: {e:?}"
2601                )
2602            });
2603            assert!(
2604                msl.contains("kernel "),
2605                "MSL ({width:?}) must declare a Metal compute entry (`kernel `); got {} bytes",
2606                msl.len()
2607            );
2608            assert!(
2609                msl.len() > 500,
2610                "MSL ({width:?}) for the multi-emit combined kernel is implausibly small \
2611                 ({} bytes), that is a stripped no-op, not the real automaton scan",
2612                msl.len()
2613            );
2614
2615            // Vulkan/SPIR-V: explicit proof, independent of the live Vulkan run.
2616            let spirv = vyre_emit_spirv::emit(&lowered.descriptor).unwrap_or_else(|e| {
2617                panic!("Fix: combined megakernel ({width:?}) must lower to valid SPIR-V: {e:?}")
2618            });
2619            assert_eq!(
2620                spirv.first().copied(),
2621                Some(vyre_emit_spirv::SPIRV_MAGIC),
2622                "SPIR-V ({width:?}) must begin with the magic word"
2623            );
2624            assert!(
2625                spirv.len() > 200,
2626                "SPIR-V ({width:?}) module is implausibly small ({} words)",
2627                spirv.len()
2628            );
2629        }
2630    }
2631
2632    #[test]
2633    fn hit_overflow_split_reports_dropped_matches() {
2634        // No overflow: every produced match fits the ring.
2635        assert_eq!(split_hit_overflow(0, 1_000), (0, 0));
2636        assert_eq!(split_hit_overflow(254, 1_000), (254, 0));
2637        // Exactly full: readable == capacity, nothing dropped.
2638        assert_eq!(split_hit_overflow(1_000, 1_000), (1_000, 0));
2639        // Overflow: readable clamps to capacity, the rest are reported dropped
2640        // the recall-critical signal the old `.min()` clamp threw away.
2641        assert_eq!(split_hit_overflow(1_001, 1_000), (1_000, 1));
2642        assert_eq!(
2643            split_hit_overflow(1_500_000, 1_000_000),
2644            (1_000_000, 500_000)
2645        );
2646        // Saturated raw head (kernel produced u32::MAX-worth of matches).
2647        assert_eq!(
2648            split_hit_overflow(u32::MAX, 1_000),
2649            (1_000, u32::MAX - 1_000)
2650        );
2651    }
2652
2653    #[test]
2654    fn default_worker_groups_is_at_least_four_on_live_adapter() {
2655        if let Ok(backend) = WgpuBackend::new() {
2656            let wg = BatchDispatchConfig::default()
2657                .launch_recommendation(backend.device_limits(), 64)
2658                .expect("Fix: live adapter limits must produce a launch recommendation")
2659                .worker_groups;
2660            assert!(
2661                wg >= 4,
2662                "Fix: default worker_groups should be >= 4 on any live adapter, got {wg}"
2663            );
2664        }
2665    }
2666
2667    /// Behavioral replacement for the former source-shape test.  Verifies that
2668    /// when `worker_groups=0` (the sentinel meaning "fill from launch policy"),
2669    /// `launch_recommendation` returns a non-zero `worker_groups` value, i.e.,
2670    /// the policy consumes the `0` sentinel and fills in a real value.
2671    /// `BatchDispatcher::new` then stores that value back into `config.worker_groups`.
2672    #[test]
2673    fn launch_recommendation_fills_zero_worker_groups_and_hit_capacity() {
2674        let limits = wgpu::Limits::default();
2675        // Default config has worker_groups=0 and hit_capacity=65_536.
2676        let config = BatchDispatchConfig::default();
2677        assert_eq!(
2678            config.worker_groups, 0,
2679            "default worker_groups must be 0 (sentinel: fill from policy)"
2680        );
2681        let rec = config
2682            .launch_recommendation(&limits, 64)
2683            .expect("Fix: default config must produce a launch recommendation");
2684        assert!(
2685            rec.worker_groups > 0,
2686            "launch policy must fill worker_groups > 0 when config.worker_groups == 0, got {}",
2687            rec.worker_groups
2688        );
2689        // hit_capacity is not zero in the default, but verify the recommendation
2690        // still provides a non-zero hit_capacity.
2691        assert!(
2692            rec.hit_capacity > 0,
2693            "launch policy must provide a positive hit_capacity, got {}",
2694            rec.hit_capacity
2695        );
2696    }
2697
2698    #[test]
2699    fn dynamic_dispatch_plan_controls_pipeline_and_launch_geometry() {
2700        let src = include_str!("dispatcher.rs");
2701        let prod_src = src.split("#[cfg(test)]").next().unwrap_or(src);
2702        assert!(
2703            prod_src.contains("let pipeline = self.pipeline_for_plan(dynamic_plan.plan)?"),
2704            "dispatch must compile or reuse the pipeline for the per-batch scale-aware plan"
2705        );
2706        assert!(
2707            prod_src.contains("[dynamic_plan.plan.worker_groups, 1, 1]"),
2708            "dispatch must submit the policy-selected worker group count, not config.worker_groups"
2709        );
2710        assert!(
2711            prod_src.contains("dynamic_plan.plan.worker_groups,\n                    self.config.workgroup_size_x"),
2712            "occupancy telemetry must use the actual dynamic launch geometry"
2713        );
2714    }
2715
2716    /// Behavioral replacement for the former source-shape test.  Verifies that
2717    /// the pipeline cache capacity constant is exactly 32 (the agreed bound) and
2718    /// that `BatchPipelineShape` has the three program-shaping fields, not
2719    /// merely that the strings exist in source.  This test does not require a
2720    /// live GPU.
2721    #[test]
2722    fn pipeline_cache_cap_is_32_and_shape_contains_all_fields() {
2723        use crate::megakernel::pipeline_cache::{BatchPipelineCache, BatchPipelineShape};
2724
2725        // Const-level check: the agreed retention bound must be exactly 32.
2726        // Changing BATCH_PIPELINE_CACHE_CAP without updating this test is a
2727        // deliberate reviewer gate.
2728        const _: () = assert!(BATCH_PIPELINE_CACHE_CAP == 32);
2729
2730        // Compile-time check: BatchPipelineShape must contain exactly the three
2731        // program-shaping fields (workgroup_size_x, worker_groups, hit_capacity).
2732        // A refactor that drops or renames any of these fields breaks this
2733        // struct literal, surfacing a compile error rather than a silent test
2734        // pass.  The source-shape strings we replaced were the only guard for
2735        // this (now the Rust type system is).
2736        let _shape = BatchPipelineShape {
2737            workgroup_size_x: 64,
2738            worker_groups: 8,
2739            hit_capacity: 512,
2740        };
2741        // Verify the constant is consumed by cache construction without panic.
2742        let cache = BatchPipelineCache::with_cap(BATCH_PIPELINE_CACHE_CAP);
2743        drop(cache);
2744    }
2745
2746    #[test]
2747    fn dynamic_plan_hit_capacity_is_clamped_to_resident_batch_ring() {
2748        let src = include_str!("dispatcher.rs");
2749        let prod_src = src.split("#[cfg(test)]").next().unwrap_or(src);
2750        assert!(
2751            prod_src.contains("let resident_hit_capacity = batch.hit_capacity()")
2752                && prod_src.contains("recommendation.hit_capacity = resident_hit_capacity")
2753                && prod_src.contains("estimated_peak_device_bytes"),
2754            "dynamic dispatch plans must not compile a hit-ring shape larger than the resident FileBatch output buffer"
2755        );
2756    }
2757
2758    #[test]
2759    fn launch_recommendation_uses_explicit_graph_hints_for_topology() {
2760        let limits = wgpu::Limits::default();
2761        let config = BatchDispatchConfig::default()
2762            .with_graph_hints(8192, 131_072, 9_000, 0)
2763            .with_execution_hints(8, 0, 0, 0);
2764
2765        let rec = config
2766            .launch_recommendation(&limits, 8192)
2767            .expect("Fix: explicit graph hints must produce a launch recommendation");
2768
2769        assert_eq!(rec.topology, MegakernelDispatchTopology::FusedDense);
2770    }
2771
2772    #[test]
2773    fn launch_recommendation_default_does_not_invent_dense_frontier() {
2774        let limits = wgpu::Limits::default();
2775        let rec = BatchDispatchConfig::default()
2776            .launch_recommendation(&limits, 8192)
2777            .expect("Fix: default graph hints must produce a launch recommendation");
2778
2779        assert_ne!(rec.topology, MegakernelDispatchTopology::FusedDense);
2780        assert_eq!(BatchDispatchConfig::default().frontier_density_bps, 0);
2781    }
2782
2783    /// Behavioral replacement for the former source-shape test.  Verifies that
2784    /// `BatchDispatchConfig.timeout` has the expected default value (30 s) and
2785    /// that the field is accessible on a constructed config.  Field existence is
2786    /// already compile-time enforced; the default value is the load-bearing
2787    /// invariant that guards dispatch budgets.
2788    #[test]
2789    fn timeout_field_has_expected_default_and_is_constructible() {
2790        let config = BatchDispatchConfig::default();
2791        assert_eq!(
2792            config.timeout,
2793            Duration::from_secs(30),
2794            "BatchDispatchConfig default timeout must be 30 s; callers that rely on the budget \
2795             must be able to predict the default"
2796        );
2797        // A zero timeout is a valid sentinel (fail immediately) (constructible).
2798        let zero_timeout_config = BatchDispatchConfig {
2799            timeout: Duration::ZERO,
2800            ..BatchDispatchConfig::default()
2801        };
2802        assert_eq!(zero_timeout_config.timeout, Duration::ZERO);
2803    }
2804
2805    #[test]
2806    fn hit_readback_decodes_without_intermediate_word_vector() {
2807        let mut bytes = Vec::new();
2808        for word in [7u32, 3, 2, 99, 8, 4, 1, 100] {
2809            bytes.extend_from_slice(&word.to_le_bytes());
2810        }
2811
2812        let hits = decode_hits_from_readback(&bytes, 2)
2813            .expect("Fix: aligned hit readback bytes must decode directly");
2814
2815        assert_eq!(hits.len(), 2);
2816        assert_eq!(hits[0].file_idx, 7);
2817        assert_eq!(hits[0].rule_idx, 3);
2818        assert_eq!(hits[1].match_offset, 100);
2819    }
2820
2821    #[test]
2822    fn hit_readback_into_reuses_caller_capacity() {
2823        let mut bytes = Vec::new();
2824        for word in [7u32, 3, 2, 99, 8, 4, 1, 100] {
2825            bytes.extend_from_slice(&word.to_le_bytes());
2826        }
2827        let mut hits = Vec::with_capacity(8);
2828        let ptr = hits.as_ptr();
2829
2830        decode_hits_from_readback_into(&bytes, 2, &mut hits)
2831            .expect("Fix: aligned hit readback bytes must decode into caller scratch");
2832
2833        assert_eq!(hits.len(), 2);
2834        assert_eq!(hits.as_ptr(), ptr);
2835    }
2836
2837    /// Law 10 (no silent fallback) on the megakernel hit-decode path: a hit-ring
2838    /// readback that exposes FEWER words than `hit_count` demands must FAIL CLOSED
2839    /// with a loud, actionable error, never silently decode the few hits it can
2840    /// reach, which would drop the rest and lose recall invisibly. `hit_count = 2`
2841    /// needs 8 words (2 hits × 4 u32 each); supplying only 4 words (one hit's
2842    /// worth) must be rejected, not decoded as a single hit.
2843    #[test]
2844    fn hit_readback_under_provisioned_ring_fails_closed_not_silent_truncation() {
2845        let mut bytes = Vec::new();
2846        for word in [7u32, 3, 2, 99] {
2847            bytes.extend_from_slice(&word.to_le_bytes());
2848        }
2849        // Pre-seed the scratch so a silent partial decode would be observable as
2850        // leftover/short content rather than an error.
2851        let mut hits = vec![
2852            HitRecord {
2853                file_idx: 0,
2854                rule_idx: 0,
2855                layer_idx: 0,
2856                match_offset: 0,
2857            };
2858            5
2859        ];
2860
2861        let err = decode_hits_from_readback_into(&bytes, 2, &mut hits).expect_err(
2862            "Fix: an under-provisioned hit ring must fail closed, never silently drop hits",
2863        );
2864        let PipelineError::Backend(message) = err else {
2865            panic!("Fix: under-provisioned hit ring must surface a Backend error, got {err:?}");
2866        };
2867        assert!(
2868            message.contains("hit-ring exposed 4 words, expected at least 8"),
2869            "Fix: the fail-closed message must name the exposed vs required word counts so the \
2870             operator can size the ring; got {message:?}"
2871        );
2872    }
2873
2874    /// A hit-ring readback whose byte length is not a whole number of u32 words is
2875    /// corrupt and must fail closed (the decode reinterprets bytes as packed u32
2876    /// records, so a ragged tail would mis-align every record). 6 bytes is 1.5
2877    /// words (never a valid readback).
2878    #[test]
2879    fn hit_readback_misaligned_byte_length_fails_closed() {
2880        let bytes = [0u8; 6];
2881        let mut hits = Vec::new();
2882
2883        let err = decode_hits_from_readback_into(&bytes, 1, &mut hits)
2884            .expect_err("Fix: a non-4-byte-aligned hit readback must fail closed");
2885        let PipelineError::Backend(message) = err else {
2886            panic!("Fix: misaligned hit readback must surface a Backend error, got {err:?}");
2887        };
2888        assert!(
2889            message.contains("not a whole number of u32 words"),
2890            "Fix: misaligned-readback error must explain the 4-byte-alignment contract; got {message:?}"
2891        );
2892    }
2893
2894    /// The under-provisioned guard is `word_count < needed_words` (strict), so an
2895    /// OVER-provisioned ring (more words than `hit_count` demands) decodes exactly
2896    /// `hit_count` hits and ignores the trailing slack, the GPU sizes the ring for
2897    /// `hit_capacity`, which is an upper bound, not the realized hit count.
2898    #[test]
2899    fn hit_readback_over_provisioned_ring_decodes_exactly_hit_count() {
2900        let mut bytes = Vec::new();
2901        for word in [7u32, 3, 2, 99, 8, 4, 1, 100, 555, 666, 777, 888] {
2902            bytes.extend_from_slice(&word.to_le_bytes());
2903        }
2904
2905        let hits = decode_hits_from_readback(&bytes, 2)
2906            .expect("Fix: an over-provisioned hit ring must decode the realized hit_count");
2907        assert_eq!(
2908            hits.len(),
2909            2,
2910            "must decode exactly hit_count, not the ring capacity"
2911        );
2912        assert_eq!(hits[1].match_offset, 100);
2913    }
2914
2915    /// Law 10 on the queue-state readback decode: `read_u32_word` is the
2916    /// fail-closed accessor that pulls `hit_head` / `items_processed` /
2917    /// `claims_attempted` out of the queue-state buffer. `hit_head` BOUNDS how many
2918    /// hits the host then decodes from the hit ring, so a silent default (0 /
2919    /// garbage) on a short queue-state readback would corrupt that bound and lose
2920    /// hits invisibly. A word index past the readback end must error loudly, never
2921    /// return a default. Valid indices decode the exact little-endian word.
2922    #[test]
2923    fn read_u32_word_past_end_fails_closed_not_silent_default() {
2924        let mut bytes = Vec::new();
2925        for word in [11u32, 22] {
2926            bytes.extend_from_slice(&word.to_le_bytes());
2927        }
2928        assert_eq!(
2929            read_u32_word(&bytes, "queue-state", 0).expect("word 0 is in range"),
2930            11
2931        );
2932        assert_eq!(
2933            read_u32_word(&bytes, "queue-state", 1).expect("word 1 is in range"),
2934            22
2935        );
2936
2937        let err = read_u32_word(&bytes, "queue-state", 2).expect_err(
2938            "Fix: reading past the readback end must fail closed, never return a silent default",
2939        );
2940        let PipelineError::Backend(message) = err else {
2941            panic!("Fix: out-of-range word read must surface a Backend error, got {err:?}");
2942        };
2943        assert!(
2944            message.contains("queue-state readback is missing u32 word 2"),
2945            "Fix: the error must name the label and the missing word index; got {message:?}"
2946        );
2947    }
2948
2949    #[test]
2950    fn occupancy_proxy_caps_at_full_utilization() {
2951        assert_eq!(occupancy_proxy_bps(32, 1, 64), 5_000);
2952        assert_eq!(occupancy_proxy_bps(128, 1, 64), 10_000);
2953        assert_eq!(occupancy_proxy_bps(0, 0, 0), 0);
2954        assert_eq!(occupancy_proxy_bps(u32::MAX, 1, 1), 10_000);
2955    }
2956
2957    /// Compile-time replacement for the former source-shape test.  Field presence
2958    /// in `BatchDispatchTelemetry` is now enforced at compile time: if any of the
2959    /// performance-gate fields were removed, this struct literal would fail to
2960    /// compile.  The old source-string scan would have accepted a field rename
2961    /// silently as long as the string appeared elsewhere in the file.
2962    #[test]
2963    fn dispatch_telemetry_exposes_all_release_counters() {
2964        // Construct a telemetry value using all release-gate fields by name.
2965        // Removing or renaming any field breaks this at compile time.
2966        let _ = BatchDispatchTelemetry {
2967            bytes_uploaded: 0,
2968            bytes_read_back: 0,
2969            bytes_moved: 0,
2970            resident_allocations: 0,
2971            kernel_launches: 0,
2972            sync_points: 0,
2973            occupancy_proxy_bps: 0,
2974            frontier_density_bps: 0,
2975            queue_state_readback_bytes: 0,
2976            hit_readback_bytes: 0,
2977            estimated_peak_device_bytes: 0,
2978            device_memory_budget_bytes: 0,
2979            topology: MegakernelDispatchTopology::SparseFrontier,
2980            dispatch_plan_cache_hit: false,
2981            dispatch_plan_cache_entries: 0,
2982        };
2983    }
2984}
2985
2986#[cfg(test)]
2987mod scan_batch_segmentation_tests {
2988    use super::{
2989        wgpu_scan_batch_segmentation_evidence, WgpuScanBatchSegmentationError,
2990        WgpuScanBatchSegmentationRequest, WGPU_SCAN_BATCH_SEGMENTATION_SCHEMA_VERSION,
2991    };
2992
2993    #[test]
2994    fn segmentation_evidence_records_command_copy_bind_group_counts_and_match_digest() {
2995        let evidence = wgpu_scan_batch_segmentation_evidence(
2996            WgpuScanBatchSegmentationRequest::new(10, 4, 2, 1, 10, 3, 0x1234, 0x1234),
2997        )
2998        .expect("Fix: valid WGPU scan segmentation evidence should be accepted");
2999
3000        assert_eq!(
3001            evidence.schema_version,
3002            WGPU_SCAN_BATCH_SEGMENTATION_SCHEMA_VERSION
3003        );
3004        assert_eq!(evidence.chunk_count, 10);
3005        assert_eq!(evidence.segment_count, 3);
3006        assert_eq!(evidence.command_encoder_count, 3);
3007        assert_eq!(evidence.bind_group_reuse_count, 2);
3008        assert_eq!(evidence.bind_group_create_count, 1);
3009        assert_eq!(evidence.copy_count, 13);
3010        assert_eq!(evidence.match_digest, 0x1234);
3011        assert!(evidence.match_parity);
3012        assert!(evidence.all_command_counts_recorded);
3013        assert!(evidence.is_complete());
3014    }
3015
3016    #[test]
3017    fn segmentation_evidence_rejects_missing_bind_group_accounting() {
3018        let error = wgpu_scan_batch_segmentation_evidence(WgpuScanBatchSegmentationRequest::new(
3019            9, 4, 1, 1, 9, 3, 0x1234, 0x1234,
3020        ))
3021        .expect_err("Fix: bind group counts must account for every segment");
3022
3023        assert!(matches!(
3024            error,
3025            WgpuScanBatchSegmentationError::BindGroupCountMismatch {
3026                command_encoder_count: 3,
3027                bind_group_reuse_count: 1,
3028                bind_group_create_count: 1
3029            }
3030        ));
3031    }
3032
3033    #[test]
3034    fn segmentation_evidence_rejects_match_digest_drift() {
3035        let error = wgpu_scan_batch_segmentation_evidence(WgpuScanBatchSegmentationRequest::new(
3036            4, 4, 0, 1, 4, 1, 0xaaaa, 0xbbbb,
3037        ))
3038        .expect_err("Fix: segmented WGPU scan output must match the oracle digest");
3039
3040        assert!(matches!(
3041            error,
3042            WgpuScanBatchSegmentationError::MatchDigestMismatch {
3043                expected_match_digest: 0xaaaa,
3044                actual_match_digest: 0xbbbb
3045            }
3046        ));
3047    }
3048
3049    /// Regression: `digest=0` is a legitimate value for a corpus with zero rule
3050    /// firings (the hash of the empty match set).  Before the fix,
3051    /// `wgpu_scan_batch_segmentation_evidence` rejected any request where either
3052    /// digest was 0 with `ZeroMatchDigest`, making it impossible to record
3053    /// evidence for a clean corpus.  `is_complete()` also used `match_digest != 0`
3054    /// as a presence gate, so even a manually constructed evidence object with
3055    /// `match_digest=0` would never satisfy the completeness check.
3056    #[test]
3057    fn evidence_accepts_zero_digest_for_empty_match_corpus() {
3058        // Both digests are 0 (both digests agree (valid evidence for a clean scan)).
3059        let evidence =
3060            wgpu_scan_batch_segmentation_evidence(WgpuScanBatchSegmentationRequest::new(
3061                4, 4, 0, 1, 4, 1, 0, 0,
3062            ))
3063            .expect("Fix: matched zero digests are valid evidence for a zero-match corpus; ZeroMatchDigest rejection was wrong");
3064
3065        assert_eq!(
3066            evidence.match_digest, 0,
3067            "evidence must preserve the zero digest value from the request"
3068        );
3069        assert!(
3070            evidence.match_parity,
3071            "evidence must report parity when both digests are equal (including zero)"
3072        );
3073        assert!(
3074            evidence.is_complete(),
3075            "evidence for a zero-match corpus must satisfy the release completeness gate"
3076        );
3077    }
3078}
3079
3080#[cfg(test)]
3081mod abi_conversion_contracts {
3082    use super::super::batch::queue_state_word;
3083    use super::super::segmentation::SEGMENT_WORDS;
3084    use super::{dispatcher_abi_u32, dispatcher_usize_to_u64};
3085    use super::{FILE_METADATA_WORDS, HIT_RECORD_WORDS, QUEUE_STATE_WORDS};
3086    use vyre_runtime::megakernel::rule_catalog::RULE_META_WORDS;
3087
3088    /// All ABI word-count constants that are embedded as u32 literals in the
3089    /// generated WGSL shader must fit in u32 without any conversion failure.
3090    /// Before the fix, `dispatcher_abi_u32` silently returned `u32::MAX` on
3091    /// failure, which would have corrupted the emitted ABI constants in the GPU
3092    /// program without any diagnostic (Law 10 silent miscompile path).
3093    #[test]
3094    fn all_abi_word_count_constants_fit_u32_without_panic() {
3095        // These are the exact callers in build_batch_program /
3096        // execute_batch_claim_body / batch_program_buffers.  If any constant
3097        // grew beyond u32::MAX the test would panic, surfacing the regression
3098        // loudly instead of silently embedding u32::MAX in the shader.
3099        let queue_len_word = dispatcher_abi_u32(queue_state_word::QUEUE_LEN, "queue-len word");
3100        let head_word = dispatcher_abi_u32(queue_state_word::HEAD, "head word");
3101        let rule_count_word = dispatcher_abi_u32(queue_state_word::RULE_COUNT, "rule-count word");
3102        let hit_head_word = dispatcher_abi_u32(queue_state_word::HIT_HEAD, "hit-head word");
3103        let hit_capacity_word =
3104            dispatcher_abi_u32(queue_state_word::HIT_CAPACITY, "hit-capacity word");
3105        let done_count_word = dispatcher_abi_u32(queue_state_word::DONE_COUNT, "done-count word");
3106        let queue_state_words_val = dispatcher_abi_u32(QUEUE_STATE_WORDS, "queue-state word count");
3107        let segment_words_val = dispatcher_abi_u32(SEGMENT_WORDS, "segment table word count");
3108        let file_meta_words_val =
3109            dispatcher_abi_u32(FILE_METADATA_WORDS, "file metadata word count");
3110        let rule_meta_words_val = dispatcher_abi_u32(RULE_META_WORDS, "rule metadata word count");
3111
3112        // Assert concrete values, the ABI is contractual; changing these
3113        // constants without updating GPU code is a silent correctness bug.
3114        assert_eq!(queue_state_words_val, 6, "QUEUE_STATE_WORDS ABI must be 6");
3115        assert_eq!(segment_words_val, 4, "SEGMENT_WORDS ABI must be 4");
3116        assert_eq!(file_meta_words_val, 4, "FILE_METADATA_WORDS ABI must be 4");
3117        assert_eq!(rule_meta_words_val, 5, "RULE_META_WORDS ABI must be 5");
3118
3119        // Smoke-check that the queue-state word indices are in [0, QUEUE_STATE_WORDS).
3120        assert!(
3121            (queue_len_word as usize) < QUEUE_STATE_WORDS,
3122            "QUEUE_LEN word index {queue_len_word} must be < QUEUE_STATE_WORDS ({QUEUE_STATE_WORDS})"
3123        );
3124        assert!(
3125            (head_word as usize) < QUEUE_STATE_WORDS,
3126            "HEAD word index {head_word} must be < QUEUE_STATE_WORDS ({QUEUE_STATE_WORDS})"
3127        );
3128        assert!(
3129            (rule_count_word as usize) < QUEUE_STATE_WORDS,
3130            "RULE_COUNT word index {rule_count_word} must be < QUEUE_STATE_WORDS ({QUEUE_STATE_WORDS})"
3131        );
3132        assert!(
3133            (hit_head_word as usize) < QUEUE_STATE_WORDS,
3134            "HIT_HEAD word index {hit_head_word} must be < QUEUE_STATE_WORDS ({QUEUE_STATE_WORDS})"
3135        );
3136        assert!(
3137            (hit_capacity_word as usize) < QUEUE_STATE_WORDS,
3138            "HIT_CAPACITY word index {hit_capacity_word} must be < QUEUE_STATE_WORDS ({QUEUE_STATE_WORDS})"
3139        );
3140        assert!(
3141            (done_count_word as usize) < QUEUE_STATE_WORDS,
3142            "DONE_COUNT word index {done_count_word} must be < QUEUE_STATE_WORDS ({QUEUE_STATE_WORDS})"
3143        );
3144    }
3145
3146    /// `dispatcher_usize_to_u64` must convert known small constants without
3147    /// panic.  Before the fix it returned `u64::MAX` silently on failure,
3148    /// causing a downstream `checked_mul` overflow that produced the misleading
3149    /// error "hit-ring readback length overflowed u64" with no indication of
3150    /// which constant failed (Law 10).
3151    #[test]
3152    fn all_usize_to_u64_abi_constants_convert_without_panic() {
3153        let hit_record_words_u64 =
3154            dispatcher_usize_to_u64(HIT_RECORD_WORDS, "hit-record word count");
3155        let u32_byte_width_u64 =
3156            dispatcher_usize_to_u64(std::mem::size_of::<u32>(), "u32 byte width");
3157        let queue_state_words_u64 =
3158            dispatcher_usize_to_u64(QUEUE_STATE_WORDS, "queue-state word count");
3159
3160        assert_eq!(
3161            hit_record_words_u64, 4,
3162            "HIT_RECORD_WORDS must convert to u64 value 4"
3163        );
3164        assert_eq!(
3165            u32_byte_width_u64, 4,
3166            "size_of::<u32>() must convert to u64 value 4"
3167        );
3168        assert_eq!(
3169            queue_state_words_u64, 6,
3170            "QUEUE_STATE_WORDS must convert to u64 value 6"
3171        );
3172    }
3173}