vyre_driver_wgpu/
parity_probe.rs

1//! Parity-test-only raw WGSL probes.
2//!
3//! These helpers deliberately bypass vyre IR validation so conformance tests can
4//! measure the backend's native WGSL transcendental behavior. They are compiled
5//! only with the `parity-testing` feature and are not part of the production
6//! dispatch path.
7
8use crate::numeric::WGPU_NUMERIC;
9use crate::staging_reserve::reserve_backend_vec;
10use crate::WgpuBackend;
11use crossbeam_channel::RecvTimeoutError;
12use std::sync::{
13    atomic::{AtomicBool, Ordering},
14    Arc,
15};
16use std::time::{Duration, Instant};
17
18const F32_BYTES: usize = std::mem::size_of::<f32>();
19const BATCH_WORKGROUP_SIZE: u32 = 64;
20const PROBE_TIMEOUT: Duration = Duration::from_secs(30);
21
22impl WgpuBackend {
23    /// Dispatch a canonical one-op f32 unary probe and return raw output bytes.
24    ///
25    /// # Errors
26    ///
27    /// Returns a backend error when `input` is not one f32, the op is not a
28    /// supported f32 unary probe, or the WGSL dispatch/readback fails.
29    pub fn probe_op(
30        &self,
31        op: vyre_foundation::ir::UnOp,
32        input: &[u8],
33    ) -> Result<Vec<u8>, vyre_driver::BackendError> {
34        if input.len() != F32_BYTES {
35            return Err(vyre_driver::BackendError::new(format!(
36                "probe_op expects exactly 4 input bytes for one f32, got {}. Fix: pass f32::to_bits().to_le_bytes().",
37                input.len()
38            )));
39        }
40        let mut raw = [0_u8; F32_BYTES];
41        raw.copy_from_slice(input);
42        let output = self.probe_op_many(op, &[f32::from_bits(u32::from_le_bytes(raw))])?;
43        let Some(sample) = output.into_iter().next() else {
44            return Err(vyre_driver::BackendError::new(
45                "probe_op produced no output for one input. Fix: inspect parity probe dispatch/readback sizing.",
46            ));
47        };
48        Ok(sample.to_bits().to_le_bytes().to_vec())
49    }
50
51    /// Dispatch a canonical f32 unary probe over a batch of inputs.
52    ///
53    /// This keeps parity tests from paying one GPU submission and readback per
54    /// scalar sample. The generated WGSL is keyed by operation, so the backend
55    /// pipeline cache reuses it across calls.
56    ///
57    /// # Errors
58    ///
59    /// Returns a backend error when the operation is unsupported, the batch is
60    /// too large for WebGPU dispatch dimensions, or dispatch/readback fails.
61    pub fn probe_op_many(
62        &self,
63        op: vyre_foundation::ir::UnOp,
64        inputs: &[f32],
65    ) -> Result<Vec<f32>, vyre_driver::BackendError> {
66        if inputs.is_empty() {
67            return Ok(Vec::new());
68        }
69        let input_words: u32 = inputs.len().try_into().map_err(|_| {
70            vyre_driver::BackendError::new(format!(
71                "probe_op_many received {} f32 samples, exceeding u32 dispatch dimensions. Fix: split the parity probe batch.",
72                inputs.len()
73            ))
74        })?;
75        let output_size = inputs.len().checked_mul(F32_BYTES).ok_or_else(|| {
76            vyre_driver::BackendError::new(format!(
77                "probe_op_many output size overflow for {} samples. Fix: split the parity probe batch.",
78                inputs.len()
79            ))
80        })?;
81        let input_bytes = f32_batch_bytes(inputs)?;
82        let output = dispatch_probe_wgsl(
83            &self.current_device_queue(),
84            &probe_wgsl(op, BATCH_WORKGROUP_SIZE)?,
85            &input_bytes,
86            output_size,
87            input_words,
88        )?;
89        decode_f32_batch(&output, input_words)
90    }
91}
92
93fn probe_wgsl(
94    op: vyre_foundation::ir::UnOp,
95    workgroup_size: u32,
96) -> Result<String, vyre_driver::BackendError> {
97    let wgsl_body = match op {
98        vyre_foundation::ir::UnOp::Sin => "sin(x)",
99        vyre_foundation::ir::UnOp::Cos => "cos(x)",
100        vyre_foundation::ir::UnOp::Sqrt => "sqrt(x)",
101        vyre_foundation::ir::UnOp::Reciprocal => "1.0 / x",
102        vyre_foundation::ir::UnOp::Exp => "exp(x)",
103        vyre_foundation::ir::UnOp::Log => "log(x)",
104        other => {
105            return Err(vyre_driver::BackendError::new(format!(
106                "unsupported f32 probe op {other:?}. Fix: use Sin, Cos, Sqrt, Reciprocal, Exp, or Log."
107            )));
108        }
109    };
110
111    Ok(format!(
112        r#"
113@group(0) @binding(0) var<storage, read> input: array<u32>;
114@group(0) @binding(1) var<storage, read_write> output: array<u32>;
115@group(1) @binding(2) var<uniform> params: vec4<u32>;
116
117@compute @workgroup_size({workgroup_size})
118fn main(@builtin(global_invocation_id) gid: vec3<u32>) {{
119    if (gid.x >= params.y) {{
120        return;
121    }}
122    let x = bitcast<f32>(input[gid.x]);
123    let y = {wgsl_body};
124    output[gid.x] = bitcast<u32>(y);
125}}
126"#
127    ))
128}
129
130fn f32_batch_bytes(inputs: &[f32]) -> Result<Vec<u8>, vyre_driver::BackendError> {
131    let byte_len = inputs.len().checked_mul(F32_BYTES).ok_or_else(|| {
132        vyre_driver::BackendError::new(format!(
133            "parity probe f32 input byte length overflow for {} samples. Fix: split the parity probe batch.",
134            inputs.len()
135        ))
136    })?;
137    let mut bytes = Vec::new();
138    reserve_backend_vec(&mut bytes, byte_len, "parity probe f32 input staging")?;
139    for value in inputs {
140        bytes.extend_from_slice(&value.to_bits().to_le_bytes());
141    }
142    Ok(bytes)
143}
144
145fn dispatch_probe_wgsl(
146    device_queue: &Arc<(wgpu::Device, wgpu::Queue)>,
147    wgsl: &str,
148    input: &[u8],
149    output_size: usize,
150    output_words: u32,
151) -> Result<Vec<u8>, vyre_driver::BackendError> {
152    let (device, queue) = &**device_queue;
153    let input_size =
154        WGPU_NUMERIC.usize_to_u64(input.len().max(F32_BYTES), "parity probe input size")?;
155    let output_size_u64 =
156        WGPU_NUMERIC.usize_to_u64(output_size.max(F32_BYTES), "parity probe output size")?;
157    let input_buffer = device.create_buffer(&wgpu::BufferDescriptor {
158        label: Some("vyre parity probe input"),
159        size: input_size,
160        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
161        mapped_at_creation: false,
162    });
163    queue.write_buffer(&input_buffer, 0, input);
164    let output_buffer = device.create_buffer(&wgpu::BufferDescriptor {
165        label: Some("vyre parity probe output"),
166        size: output_size_u64,
167        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
168        mapped_at_creation: false,
169    });
170    let readback_buffer = device.create_buffer(&wgpu::BufferDescriptor {
171        label: Some("vyre parity probe readback"),
172        size: output_size_u64,
173        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
174        mapped_at_creation: false,
175    });
176    let input_len = u32::try_from(input.len()).map_err(|_| {
177        vyre_driver::BackendError::new(
178            "parity probe input length exceeds u32::MAX bytes. Fix: split the probe input before dispatch.",
179        )
180    })?;
181    let params = [input_len, output_words, 0_u32, 0_u32];
182    let params_buffer = device.create_buffer(&wgpu::BufferDescriptor {
183        label: Some("vyre parity probe params"),
184        size: WGPU_NUMERIC.usize_to_u64(
185            params.len()
186                .checked_mul(std::mem::size_of::<u32>())
187                .ok_or_else(|| {
188                    vyre_driver::BackendError::new(
189                        "parity probe params byte length overflowed usize. Fix: reduce parameter words before probe dispatch.",
190                    )
191                })?,
192            "parity probe params size",
193        )?,
194        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
195        mapped_at_creation: false,
196    });
197    queue.write_buffer(&params_buffer, 0, bytemuck::cast_slice(&params));
198
199    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
200        label: Some("vyre parity probe shader"),
201        source: wgpu::ShaderSource::Wgsl(wgsl.into()),
202    });
203    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
204        label: Some("vyre parity probe pipeline"),
205        layout: None,
206        module: &module,
207        entry_point: Some("main"),
208        compilation_options: wgpu::PipelineCompilationOptions::default(),
209        cache: None,
210    });
211    let group0_layout = pipeline.get_bind_group_layout(0);
212    let group1_layout = pipeline.get_bind_group_layout(1);
213    let group0 = device.create_bind_group(&wgpu::BindGroupDescriptor {
214        label: Some("vyre parity probe storage bind group"),
215        layout: &group0_layout,
216        entries: &[
217            wgpu::BindGroupEntry {
218                binding: 0,
219                resource: input_buffer.as_entire_binding(),
220            },
221            wgpu::BindGroupEntry {
222                binding: 1,
223                resource: output_buffer.as_entire_binding(),
224            },
225        ],
226    });
227    let group1 = device.create_bind_group(&wgpu::BindGroupDescriptor {
228        label: Some("vyre parity probe uniform bind group"),
229        layout: &group1_layout,
230        entries: &[wgpu::BindGroupEntry {
231            binding: 2,
232            resource: params_buffer.as_entire_binding(),
233        }],
234    });
235
236    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
237        label: Some("vyre parity probe encoder"),
238    });
239    {
240        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
241            label: Some("vyre parity probe compute"),
242            timestamp_writes: None,
243        });
244        pass.set_pipeline(&pipeline);
245        pass.set_bind_group(0, &group0, &[]);
246        pass.set_bind_group(1, &group1, &[]);
247        pass.dispatch_workgroups(output_words.div_ceil(BATCH_WORKGROUP_SIZE), 1, 1);
248    }
249    encoder.copy_buffer_to_buffer(&output_buffer, 0, &readback_buffer, 0, output_size_u64);
250    queue.submit(std::iter::once(encoder.finish()));
251
252    let slice = readback_buffer.slice(0..output_size_u64);
253    let (sender, receiver) = crossbeam_channel::bounded(1);
254    let ready = Arc::new(AtomicBool::new(false));
255    let ready_cb = Arc::clone(&ready);
256    slice.map_async(wgpu::MapMode::Read, move |result| {
257        if let Err(error) = sender.send(result) {
258            tracing::warn!("wgpu parity probe readback notification failed: {error}");
259        }
260        ready_cb.store(true, Ordering::Release);
261    });
262
263    let deadline = Instant::now() + PROBE_TIMEOUT;
264    let mut backoff = crate::wait_backoff::AdaptiveWaitBackoff::from_micros(64, 2, 50, 5);
265    while !ready.load(Ordering::Acquire) {
266        crate::runtime::device::poll_device_once(device)?;
267        let now = Instant::now();
268        if now >= deadline {
269            return Err(vyre_driver::BackendError::new(
270                "parity probe readback did not complete within 30s. Fix: inspect wgpu device polling and direct-dispatch readback liveness.",
271            ));
272        }
273        backoff.idle_for(deadline.duration_since(now));
274    }
275    let remaining = deadline
276        .checked_duration_since(Instant::now())
277        .ok_or_else(|| {
278            vyre_driver::BackendError::new(
279                "parity probe readback became ready after its receive deadline. Fix: inspect wgpu callback scheduling latency and raise the probe timeout deliberately.",
280            )
281        })?;
282    receiver
283        .recv_timeout(remaining)
284        .map_err(|error| match error {
285            RecvTimeoutError::Timeout => vyre_driver::BackendError::new(
286                "parity probe readback callback timed out after readiness. Fix: keep callback receiver and readiness flag synchronized.",
287            ),
288            RecvTimeoutError::Disconnected => vyre_driver::BackendError::new(
289                "parity probe readback callback disconnected. Fix: keep map_async callback sender alive until collection.",
290            ),
291        })?
292        .map_err(|error| {
293            vyre_driver::BackendError::new(format!(
294                "parity probe readback mapping failed: {error:?}. Fix: verify readback buffer MAP_READ/COPY_DST usage."
295            ))
296        })?;
297    let mapped = slice.get_mapped_range();
298    let mut bytes = Vec::new();
299    reserve_backend_vec(&mut bytes, output_size, "parity probe readback staging")?;
300    bytes.extend_from_slice(&mapped[..output_size]);
301    drop(mapped);
302    readback_buffer.unmap();
303    Ok(bytes)
304}
305
306fn decode_f32_batch(
307    output: &[u8],
308    expected_words: u32,
309) -> Result<Vec<f32>, vyre_driver::BackendError> {
310    let expected_words = usize::try_from(expected_words).map_err(|_| {
311        vyre_driver::BackendError::new(
312            "parity probe expected word count cannot fit host usize. Fix: split probe readback into smaller batches.",
313        )
314    })?;
315    let expected_bytes = expected_words
316        .checked_mul(F32_BYTES)
317        .ok_or_else(|| {
318            vyre_driver::BackendError::new(
319                "parity probe expected byte count overflowed usize. Fix: split probe readback into smaller batches.",
320            )
321        })?;
322    if output.len() != expected_bytes {
323        return Err(vyre_driver::BackendError::new(format!(
324            "batch probe returned {} bytes for {expected_words} f32 samples. Fix: keep dispatch_wgsl readback size synchronized with probe batch length.",
325            output.len()
326        )));
327    }
328    let mut values = Vec::new();
329    reserve_backend_vec(
330        &mut values,
331        expected_words,
332        "parity probe decoded f32 staging",
333    )?;
334    for chunk in output.chunks_exact(F32_BYTES) {
335        let mut raw = [0_u8; F32_BYTES];
336        raw.copy_from_slice(chunk);
337        values.push(f32::from_bits(u32::from_le_bytes(raw)));
338    }
339    Ok(values)
340}
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345
346    #[test]
347    fn batch_probe_matches_singleton_probe() {
348        let backend = WgpuBackend::acquire()
349            .expect("Fix: parity probe tests require the local GPU-backed wgpu backend");
350        let inputs = [0.0_f32, 0.5, 1.0, -2.0];
351        let batch = backend
352            .probe_op_many(vyre_foundation::ir::UnOp::Cos, &inputs)
353            .expect("Fix: batched parity probe must dispatch successfully");
354        assert_eq!(batch.len(), inputs.len());
355
356        for (index, input) in inputs.iter().enumerate() {
357            let singleton = backend
358                .probe_op(
359                    vyre_foundation::ir::UnOp::Cos,
360                    &input.to_bits().to_le_bytes(),
361                )
362                .expect("Fix: singleton parity probe must dispatch successfully");
363            let mut raw = [0_u8; F32_BYTES];
364            raw.copy_from_slice(&singleton);
365            assert_eq!(
366                batch[index].to_bits(),
367                f32::from_bits(u32::from_le_bytes(raw)).to_bits(),
368                "Fix: batched probe result at lane {index} must match singleton probe"
369            );
370        }
371    }
372
373    #[test]
374    fn parity_probe_uses_fallible_staging() {
375        let src =
376            std::fs::read_to_string(concat!(env!("CARGO_MANIFEST_DIR"), "/src/parity_probe.rs"))
377                .expect("Fix: parity probe source must be readable");
378        let production = src
379            .split("#[cfg(test)]")
380            .next()
381            .expect("Fix: meta-test scans production sources; update fixture path if module moved - production section must exist");
382        assert!(
383            !production.contains("Vec::with_capacity("),
384            "parity probe staging must use reserve_backend_vec"
385        );
386        assert!(production.contains("reserve_backend_vec"));
387        assert!(production.contains("f32_batch_bytes(inputs)?"));
388    }
389}
vyre_driver_wgpu/parity_probe.rs

vyre_driver_wgpu/
parity_probe.rs
