vyre_driver_wgpu/
ext.rs

1//! WGSL-specific dispatch helpers for the wgpu backend.
2//!
3//! Raw WGSL is a property of the wgpu implementation, not the substrate-neutral
4//! [`vyre_driver::VyreBackend`] contract.
5
6use std::sync::Arc;
7
8use dashmap::mapref::entry::Entry;
9
10use crate::engine::record_and_readback::{record_and_readback, DispatchLabels, RecordAndReadback};
11use crate::pipeline::{BufferBindingInfo, OutputBindingLayout, OutputLayout};
12use crate::WgpuBackend;
13use vyre_emit_naga::program::bind_group_for;
14use vyre_foundation::ir::{BufferAccess, DataType};
15
16impl WgpuBackend {
17    /// Dispatch a raw WGSL compute shader.
18    ///
19    /// # Errors
20    ///
21    /// Returns an actionable error when shader compilation, staging-buffer
22    /// creation, command submission, or readback fails.
23    pub fn dispatch_wgsl(
24        &self,
25        wgsl: &str,
26        input: &[u8],
27        output_size: usize,
28        workgroup_size: u32,
29    ) -> Result<Vec<u8>, String> {
30        if workgroup_size == 0 {
31            return Err("Fix: dispatch_wgsl workgroup_size must be greater than zero.".to_string());
32        }
33        let device_queue = self.current_device_queue();
34        let (device, _queue) = &*device_queue;
35
36        let cache_key = dispatch_wgsl_pipeline_cache_key(wgsl, "main")?;
37        let pipeline = if let Some(hit) = self.wgsl_dispatch_pipeline_cache.get(&cache_key) {
38            Arc::clone(hit.value())
39        } else {
40            let compiled = Arc::new(
41                crate::runtime::compile_compute_pipeline(
42                    device,
43                    "vyre backend dispatch_wgsl",
44                    wgsl,
45                    "main",
46                )
47                .map_err(|error| error.to_string())?,
48            );
49            match self.wgsl_dispatch_pipeline_cache.entry(cache_key) {
50                Entry::Occupied(hit) => Arc::clone(hit.get()),
51                Entry::Vacant(slot) => {
52                    slot.insert(Arc::clone(&compiled));
53                    compiled
54                }
55            }
56        };
57
58        let output_word_count = output_size
59            .checked_add(3)
60            .and_then(|n| n.checked_div(4))
61            .ok_or_else(|| {
62                format!(
63                    "Fix: output_size {output_size} overflows WGSL dispatch word-count calculation; split the dispatch into smaller chunks."
64                )
65            })?
66            .max(1);
67        let output_bytes = output_word_count.checked_mul(4).ok_or_else(|| {
68            format!(
69                "Fix: output_word_count {output_word_count} overflows usize bytes; reduce output_size"
70            )
71        })?;
72        let input_len_u32 = u32::try_from(input.len()).map_err(|_| {
73            format!(
74                "Fix: input length {} exceeds u32 capacity; split the dispatch into u32-sized chunks",
75                input.len()
76            )
77        })?;
78        let output_len_u32 = u32::try_from(output_word_count).map_err(|_| {
79            format!(
80                "Fix: output_word_count {output_word_count} exceeds u32 capacity; reduce output_size"
81            )
82        })?;
83        let params = [input_len_u32, output_len_u32, 0u32, 0u32];
84        let params_bytes = bytemuck::try_cast_slice(&params).map_err(|error| {
85            vyre_driver::BackendError::new(format!(
86                "WGSL dispatch params could not be viewed as bytes: {error}. Fix: keep dispatch parameter buffers aligned to u32."
87            ))
88            .into_message()
89        })?;
90
91        let workgroup_count = u32::try_from(
92            output_word_count
93            .div_ceil(usize::try_from(workgroup_size).map_err(|error| {
94                format!(
95                    "Fix: WGSL workgroup_size {workgroup_size} cannot fit usize: {error}; reduce workgroup size."
96                )
97            })?)
98            .max(1),
99        )
100        .map_err(|_| {
101            format!(
102                "Fix: WGSL dispatch requires more than u32::MAX workgroups for {output_word_count} output words and workgroup size {workgroup_size}; split the dispatch."
103            )
104        })?;
105        let input_word_count = input.len().div_ceil(4).max(1);
106        let input_word_count_u32 = u32::try_from(input_word_count).map_err(|_| {
107            format!(
108                "Fix: input word count {input_word_count} exceeds u32 capacity; split the dispatch into u32-sized chunks."
109            )
110        })?;
111        let buffer_bindings = [
112            BufferBindingInfo {
113                internal_trap: false,
114                group: bind_group_for(vyre_foundation::ir::MemoryKind::Readonly),
115                binding: 0,
116                name: Arc::from("input"),
117                access: BufferAccess::ReadOnly,
118                kind: vyre_foundation::ir::MemoryKind::Readonly,
119                hints: vyre_foundation::ir::MemoryHints::default(),
120                element: DataType::U32,
121                count: input_word_count_u32,
122                is_output: false,
123                preserve_input_contents: false,
124            },
125            BufferBindingInfo {
126                internal_trap: false,
127                group: bind_group_for(vyre_foundation::ir::MemoryKind::Global),
128                binding: 1,
129                name: Arc::from("output"),
130                access: BufferAccess::ReadWrite,
131                kind: vyre_foundation::ir::MemoryKind::Global,
132                hints: vyre_foundation::ir::MemoryHints::default(),
133                element: DataType::U32,
134                count: output_len_u32,
135                is_output: true,
136                preserve_input_contents: false,
137            },
138            BufferBindingInfo {
139                internal_trap: false,
140                group: bind_group_for(vyre_foundation::ir::MemoryKind::Uniform),
141                binding: 2,
142                name: Arc::from("params"),
143                access: BufferAccess::Uniform,
144                kind: vyre_foundation::ir::MemoryKind::Uniform,
145                hints: vyre_foundation::ir::MemoryHints::default(),
146                element: DataType::U32,
147                count: 4,
148                is_output: false,
149                preserve_input_contents: false,
150            },
151        ];
152        let max_group: u32 = buffer_bindings.iter().map(|b| b.group).max().unwrap_or(0);
153        let bind_group_count = max_group.checked_add(1).ok_or_else(|| {
154            "raw WGSL bind-group count overflowed u32. Fix: lower buffer group indices before dispatch."
155                .to_string()
156        })?;
157        let bind_group_capacity = usize::try_from(bind_group_count).map_err(|error| {
158            format!(
159                "raw WGSL bind-group count {bind_group_count} cannot fit usize: {error}. Fix: lower buffer group indices before dispatch."
160            )
161        })?;
162        let mut bind_group_layouts = Vec::with_capacity(bind_group_capacity);
163        bind_group_layouts
164            .extend((0..=max_group).map(|g| Arc::new(pipeline.get_bind_group_layout(g))));
165        let inputs = [input, params_bytes];
166        let output_bindings: Arc<[OutputBindingLayout]> = Arc::from([OutputBindingLayout {
167            binding: 1,
168            name: Arc::from("output"),
169            layout: OutputLayout {
170                full_size: output_bytes,
171                read_size: output_size,
172                copy_offset: 0,
173                copy_size: output_bytes,
174                trim_start: 0,
175            },
176            word_count: output_word_count,
177        }]);
178        let dispatch_arena = self.dispatch_arena_snapshot();
179        let outputs = record_and_readback(RecordAndReadback {
180            device_queue: &device_queue,
181            pool: dispatch_arena.pool(),
182            readback_rings: None,
183            pipeline: pipeline.as_ref(),
184            bind_group_layouts: &bind_group_layouts,
185            bind_group_cache: None,
186            buffer_bindings: &buffer_bindings,
187            inputs: &inputs,
188            output_bindings: &output_bindings,
189            trap_tags: &[],
190            workgroup_count: [workgroup_count, 1, 1],
191            indirect: None,
192            labels: DispatchLabels {
193                bind_group: "vyre backend dispatch_wgsl bind group",
194                encoder: "vyre backend dispatch_wgsl",
195                compute: "vyre backend dispatch_wgsl compute",
196            },
197            iterations: 1,
198            timestamp_profile: false,
199        })
200        .map_err(|error| error.into_message())?;
201
202        outputs
203            .into_iter()
204            .next()
205            .ok_or_else(|| "WGSL dispatch produced no output. Fix: declare binding(1) as the output storage buffer.".to_string())
206    }
207}
208
209fn dispatch_wgsl_pipeline_cache_key(wgsl: &str, entry_point: &str) -> Result<[u8; 32], String> {
210    let mut hasher = blake3::Hasher::new();
211    hasher.update(b"vyre-wgpu.dispatch_wgsl.pipeline.v1");
212    hasher.update(
213        &crate::numeric::usize_to_u64(entry_point.len(), "dispatch_wgsl entry point length")
214            .map_err(|error| error.into_message())?
215            .to_le_bytes(),
216    );
217    hasher.update(entry_point.as_bytes());
218    hasher.update(
219        &crate::numeric::usize_to_u64(wgsl.len(), "dispatch_wgsl WGSL length")
220            .map_err(|error| error.into_message())?
221            .to_le_bytes(),
222    );
223    hasher.update(wgsl.as_bytes());
224    Ok(*hasher.finalize().as_bytes())
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230
231    #[test]
232    fn dispatch_wgsl_reuses_backend_pipeline_cache() {
233        let Ok(backend) = WgpuBackend::acquire() else {
234            panic!("Fix: WGPU dispatch_wgsl cache test requires a live GPU adapter");
235        };
236        let wgsl = r#"
237@group(0) @binding(0) var<storage, read> input: array<u32>;
238@group(0) @binding(1) var<storage, read_write> output: array<u32>;
239@group(1) @binding(2) var<uniform> params: vec4<u32>;
240
241@compute @workgroup_size(64)
242fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
243    if (gid.x >= params.y) {
244        return;
245    }
246    output[gid.x] = input[gid.x] + 1u;
247}
248"#;
249        let input: Vec<u8> = [41_u32, 99_u32]
250            .into_iter()
251            .flat_map(u32::to_le_bytes)
252            .collect();
253
254        let first = backend
255            .dispatch_wgsl(wgsl, &input, 8, 64)
256            .expect("Fix: first raw WGSL dispatch must compile and run");
257        let second = backend
258            .dispatch_wgsl(wgsl, &input, 8, 64)
259            .expect("Fix: second raw WGSL dispatch must reuse the cached pipeline and run");
260
261        assert_eq!(first, second);
262        assert_eq!(first, [42_u32, 100_u32].as_slice().as_bytes());
263        assert_eq!(
264            backend.wgsl_dispatch_pipeline_cache.len(),
265            1,
266            "Fix: identical dispatch_wgsl source must compile once per backend instance"
267        );
268    }
269
270    trait U32SliceBytes {
271        fn as_bytes(&self) -> &[u8];
272    }
273
274    impl U32SliceBytes for [u32] {
275        fn as_bytes(&self) -> &[u8] {
276            bytemuck::cast_slice(self)
277        }
278    }
279}
vyre_driver_wgpu/ext.rs

vyre_driver_wgpu/
ext.rs
