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 = WGPU_NUMERIC.usize_to_u64(input.len().max(F32_BYTES), "parity probe input size")?;
154    let output_size_u64 = WGPU_NUMERIC.usize_to_u64(output_size.max(F32_BYTES), "parity probe output size")?;
155    let input_buffer = device.create_buffer(&wgpu::BufferDescriptor {
156        label: Some("vyre parity probe input"),
157        size: input_size,
158        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
159        mapped_at_creation: false,
160    });
161    queue.write_buffer(&input_buffer, 0, input);
162    let output_buffer = device.create_buffer(&wgpu::BufferDescriptor {
163        label: Some("vyre parity probe output"),
164        size: output_size_u64,
165        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
166        mapped_at_creation: false,
167    });
168    let readback_buffer = device.create_buffer(&wgpu::BufferDescriptor {
169        label: Some("vyre parity probe readback"),
170        size: output_size_u64,
171        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
172        mapped_at_creation: false,
173    });
174    let input_len = u32::try_from(input.len()).map_err(|_| {
175        vyre_driver::BackendError::new(
176            "parity probe input length exceeds u32::MAX bytes. Fix: split the probe input before dispatch.",
177        )
178    })?;
179    let params = [input_len, output_words, 0_u32, 0_u32];
180    let params_buffer = device.create_buffer(&wgpu::BufferDescriptor {
181        label: Some("vyre parity probe params"),
182        size: WGPU_NUMERIC.usize_to_u64(
183            params.len()
184                .checked_mul(std::mem::size_of::<u32>())
185                .ok_or_else(|| {
186                    vyre_driver::BackendError::new(
187                        "parity probe params byte length overflowed usize. Fix: reduce parameter words before probe dispatch.",
188                    )
189                })?,
190            "parity probe params size",
191        )?,
192        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
193        mapped_at_creation: false,
194    });
195    queue.write_buffer(&params_buffer, 0, bytemuck::cast_slice(&params));
196
197    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
198        label: Some("vyre parity probe shader"),
199        source: wgpu::ShaderSource::Wgsl(wgsl.into()),
200    });
201    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
202        label: Some("vyre parity probe pipeline"),
203        layout: None,
204        module: &module,
205        entry_point: Some("main"),
206        compilation_options: wgpu::PipelineCompilationOptions::default(),
207        cache: None,
208    });
209    let group0_layout = pipeline.get_bind_group_layout(0);
210    let group1_layout = pipeline.get_bind_group_layout(1);
211    let group0 = device.create_bind_group(&wgpu::BindGroupDescriptor {
212        label: Some("vyre parity probe storage bind group"),
213        layout: &group0_layout,
214        entries: &[
215            wgpu::BindGroupEntry {
216                binding: 0,
217                resource: input_buffer.as_entire_binding(),
218            },
219            wgpu::BindGroupEntry {
220                binding: 1,
221                resource: output_buffer.as_entire_binding(),
222            },
223        ],
224    });
225    let group1 = device.create_bind_group(&wgpu::BindGroupDescriptor {
226        label: Some("vyre parity probe uniform bind group"),
227        layout: &group1_layout,
228        entries: &[wgpu::BindGroupEntry {
229            binding: 2,
230            resource: params_buffer.as_entire_binding(),
231        }],
232    });
233
234    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
235        label: Some("vyre parity probe encoder"),
236    });
237    {
238        let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
239            label: Some("vyre parity probe compute"),
240            timestamp_writes: None,
241        });
242        pass.set_pipeline(&pipeline);
243        pass.set_bind_group(0, &group0, &[]);
244        pass.set_bind_group(1, &group1, &[]);
245        pass.dispatch_workgroups(output_words.div_ceil(BATCH_WORKGROUP_SIZE), 1, 1);
246    }
247    encoder.copy_buffer_to_buffer(&output_buffer, 0, &readback_buffer, 0, output_size_u64);
248    queue.submit(std::iter::once(encoder.finish()));
249
250    let slice = readback_buffer.slice(0..output_size_u64);
251    let (sender, receiver) = crossbeam_channel::bounded(1);
252    let ready = Arc::new(AtomicBool::new(false));
253    let ready_cb = Arc::clone(&ready);
254    slice.map_async(wgpu::MapMode::Read, move |result| {
255        if let Err(error) = sender.send(result) {
256            tracing::warn!("wgpu parity probe readback notification failed: {error}");
257        }
258        ready_cb.store(true, Ordering::Release);
259    });
260
261    let deadline = Instant::now() + PROBE_TIMEOUT;
262    let mut backoff = crate::wait_backoff::AdaptiveWaitBackoff::from_micros(64, 2, 50, 5);
263    while !ready.load(Ordering::Acquire) {
264        crate::runtime::device::poll_device_once(device)?;
265        let now = Instant::now();
266        if now >= deadline {
267            return Err(vyre_driver::BackendError::new(
268                "parity probe readback did not complete within 30s. Fix: inspect wgpu device polling and direct-dispatch readback liveness.",
269            ));
270        }
271        backoff.idle_for(deadline.duration_since(now));
272    }
273    let remaining = deadline
274        .checked_duration_since(Instant::now())
275        .ok_or_else(|| {
276            vyre_driver::BackendError::new(
277                "parity probe readback became ready after its receive deadline. Fix: inspect wgpu callback scheduling latency and raise the probe timeout deliberately.",
278            )
279        })?;
280    receiver
281        .recv_timeout(remaining)
282        .map_err(|error| match error {
283            RecvTimeoutError::Timeout => vyre_driver::BackendError::new(
284                "parity probe readback callback timed out after readiness. Fix: keep callback receiver and readiness flag synchronized.",
285            ),
286            RecvTimeoutError::Disconnected => vyre_driver::BackendError::new(
287                "parity probe readback callback disconnected. Fix: keep map_async callback sender alive until collection.",
288            ),
289        })?
290        .map_err(|error| {
291            vyre_driver::BackendError::new(format!(
292                "parity probe readback mapping failed: {error:?}. Fix: verify readback buffer MAP_READ/COPY_DST usage."
293            ))
294        })?;
295    let mapped = slice.get_mapped_range();
296    let mut bytes = Vec::new();
297    reserve_backend_vec(&mut bytes, output_size, "parity probe readback staging")?;
298    bytes.extend_from_slice(&mapped[..output_size]);
299    drop(mapped);
300    readback_buffer.unmap();
301    Ok(bytes)
302}
303
304fn decode_f32_batch(
305    output: &[u8],
306    expected_words: u32,
307) -> Result<Vec<f32>, vyre_driver::BackendError> {
308    let expected_words = usize::try_from(expected_words).map_err(|_| {
309        vyre_driver::BackendError::new(
310            "parity probe expected word count cannot fit host usize. Fix: split probe readback into smaller batches.",
311        )
312    })?;
313    let expected_bytes = expected_words
314        .checked_mul(F32_BYTES)
315        .ok_or_else(|| {
316            vyre_driver::BackendError::new(
317                "parity probe expected byte count overflowed usize. Fix: split probe readback into smaller batches.",
318            )
319        })?;
320    if output.len() != expected_bytes {
321        return Err(vyre_driver::BackendError::new(format!(
322            "batch probe returned {} bytes for {expected_words} f32 samples. Fix: keep dispatch_wgsl readback size synchronized with probe batch length.",
323            output.len()
324        )));
325    }
326    let mut values = Vec::new();
327    reserve_backend_vec(
328        &mut values,
329        expected_words,
330        "parity probe decoded f32 staging",
331    )?;
332    for chunk in output.chunks_exact(F32_BYTES) {
333        let mut raw = [0_u8; F32_BYTES];
334        raw.copy_from_slice(chunk);
335        values.push(f32::from_bits(u32::from_le_bytes(raw)));
336    }
337    Ok(values)
338}
339
340#[cfg(test)]
341mod tests {
342    use super::*;
343
344    #[test]
345    fn batch_probe_matches_singleton_probe() {
346        let backend = WgpuBackend::acquire()
347            .expect("Fix: parity probe tests require the local GPU-backed wgpu backend");
348        let inputs = [0.0_f32, 0.5, 1.0, -2.0];
349        let batch = backend
350            .probe_op_many(vyre_foundation::ir::UnOp::Cos, &inputs)
351            .expect("Fix: batched parity probe must dispatch successfully");
352        assert_eq!(batch.len(), inputs.len());
353
354        for (index, input) in inputs.iter().enumerate() {
355            let singleton = backend
356                .probe_op(
357                    vyre_foundation::ir::UnOp::Cos,
358                    &input.to_bits().to_le_bytes(),
359                )
360                .expect("Fix: singleton parity probe must dispatch successfully");
361            let mut raw = [0_u8; F32_BYTES];
362            raw.copy_from_slice(&singleton);
363            assert_eq!(
364                batch[index].to_bits(),
365                f32::from_bits(u32::from_le_bytes(raw)).to_bits(),
366                "Fix: batched probe result at lane {index} must match singleton probe"
367            );
368        }
369    }
370
371    #[test]
372    fn parity_probe_uses_fallible_staging() {
373        let src =
374            std::fs::read_to_string(concat!(env!("CARGO_MANIFEST_DIR"), "/src/parity_probe.rs"))
375                .expect("Fix: parity probe source must be readable");
376        let production = src
377            .split("#[cfg(test)]")
378            .next()
379            .expect("Fix: meta-test scans production sources; update fixture path if module moved - production section must exist");
380        assert!(
381            !production.contains("Vec::with_capacity("),
382            "parity probe staging must use reserve_backend_vec"
383        );
384        assert!(production.contains("reserve_backend_vec"));
385        assert!(production.contains("f32_batch_bytes(inputs)?"));
386    }
387}
vyre_driver_wgpu/parity_probe.rs

vyre_driver_wgpu/
parity_probe.rs
