wde_renderer/
compute.rs

1//! On-demand compute dispatch for callers outside the render world's own per-frame systems (e.g.
2//! background tasks). Unlike [`crate::core::PipelineManager`], which builds pipelines
3//! asynchronously across frames and is only usable from render-app ECS systems, this builds
4//! synchronously and caches by label, so it's safe to call from any thread and cheap to call
5//! repeatedly.
6
7use std::collections::HashMap;
8use std::sync::{Arc, Mutex};
9
10use wde_wgpu::bind_group::{BindGroup, BindGroupBuilder, BindGroupLayout, WgpuBindGroupLayout};
11use wde_wgpu::buffer::{Buffer, BufferBindingType, BufferUsage};
12use wde_wgpu::command_buffer::CommandBuffer;
13use wde_wgpu::compute_pipeline::ComputePipeline;
14use wde_wgpu::instance::RenderInstanceData;
15use wde_wgpu::render_pipeline::ShaderStages;
16
17use crate::core::RenderInstance;
18
19struct CachedPipeline {
20    pipeline: ComputePipeline,
21    layout: WgpuBindGroupLayout
22}
23
24/// One dispatch's GPU-side working set, checked out of the pool for the call's duration.
25struct BufferSet {
26    params: Buffer,
27    output: Buffer,
28    inputs: Vec<Buffer>,
29    staging: Buffer,
30    bind_group: BindGroup
31}
32
33/// Identifies interchangeable buffer sets: same label, params size, output/input lengths.
34#[derive(PartialEq, Eq, Hash, Clone)]
35struct BufferShape {
36    label: &'static str,
37    params_size: usize,
38    output_len: usize,
39    input_lens: Vec<usize>
40}
41
42/// Caches a compiled pipeline per label and pools GPU buffers (+ bind group) by shape, so a
43/// repeat dispatch only re-uploads buffer contents. Blocks until the result is read back.
44pub struct ComputeDispatcher {
45    instance: RenderInstance,
46    pipelines: Mutex<HashMap<&'static str, Arc<CachedPipeline>>>,
47    buffer_pool: Mutex<HashMap<BufferShape, Vec<BufferSet>>>
48}
49impl ComputeDispatcher {
50    pub fn new(instance: &RenderInstance) -> Self {
51        Self {
52            instance: RenderInstance(instance.0.clone()),
53            pipelines: Mutex::new(HashMap::new()),
54            buffer_pool: Mutex::new(HashMap::new())
55        }
56    }
57
58    /// Dispatches `shader`'s `main` over `output_len` `f32`s (binding 1), with `params` as a
59    /// uniform buffer (binding 0) and `inputs` as read-only `f32` storage buffers at bindings
60    /// `2, 3, ...`. `inputs.len()` must stay the same across calls sharing `label`, since the
61    /// bind group layout is built once and cached under it.
62    pub fn dispatch_f32<P: bytemuck::NoUninit>(
63        &self,
64        label: &'static str,
65        shader: &str,
66        params: &P,
67        inputs: &[&[f32]],
68        output_len: usize,
69        workgroups: (u32, u32, u32)
70    ) -> Result<Vec<f32>, String> {
71        let data = self
72            .instance
73            .0
74            .read()
75            .map_err(|_| "render instance lock poisoned".to_string())?;
76
77        let pipeline = self.pipeline_for(label, shader, inputs.len(), &data)?;
78        let shape = BufferShape {
79            label,
80            params_size: std::mem::size_of::<P>(),
81            output_len,
82            input_lens: inputs.iter().map(|slice| slice.len()).collect()
83        };
84        let buffers = self.checkout_buffers(shape.clone(), &pipeline, &data)?;
85
86        buffers.params.write(&data, bytemuck::bytes_of(params), 0);
87        for (buf, input) in buffers.inputs.iter().zip(inputs) {
88            buf.write(&data, bytemuck::cast_slice(input), 0);
89        }
90
91        let mut cmd = CommandBuffer::new(&data, label);
92        {
93            let mut pass = cmd.create_compute_pass(label);
94            pass.set_pipeline(&pipeline.pipeline)
95                .map_err(|e| format!("{label}: {e:?}"))?
96                .set_bind_group(0, &buffers.bind_group);
97            pass.dispatch(workgroups.0, workgroups.1, workgroups.2)
98                .map_err(|e| format!("{label}: {e:?}"))?;
99        }
100        cmd.submit(&data);
101
102        buffers.staging.copy_from_buffer(&data, &buffers.output);
103        let mut result = vec![0.0f32; output_len];
104        buffers.staging.map_read(&data, |view| {
105            result.copy_from_slice(bytemuck::cast_slice(&view));
106        });
107
108        self.buffer_pool
109            .lock()
110            .unwrap()
111            .entry(shape)
112            .or_default()
113            .push(buffers);
114        Ok(result)
115    }
116
117    fn pipeline_for(
118        &self,
119        label: &'static str,
120        shader: &str,
121        input_count: usize,
122        data: &RenderInstanceData
123    ) -> Result<Arc<CachedPipeline>, String> {
124        let mut cache = self.pipelines.lock().unwrap();
125        if let Some(cached) = cache.get(label) {
126            return Ok(cached.clone());
127        }
128
129        let layout = BindGroupLayout::new(label, |b| {
130            b.add_buffer(0, ShaderStages::COMPUTE, BufferBindingType::Uniform);
131            b.add_buffer(
132                1,
133                ShaderStages::COMPUTE,
134                BufferBindingType::Storage { read_only: false }
135            );
136            for i in 0..input_count {
137                b.add_buffer(
138                    2 + i as u32,
139                    ShaderStages::COMPUTE,
140                    BufferBindingType::Storage { read_only: true }
141                );
142            }
143        });
144        let wgpu_layout = layout.build(data).map_err(|e| format!("{label}: {e:?}"))?;
145        let mut pipeline = ComputePipeline::new(label);
146        pipeline
147            .set_shader(shader)
148            .set_bind_groups(vec![wgpu_layout.clone()]);
149        pipeline.init(data).map_err(|e| format!("{label}: {e:?}"))?;
150
151        let cached = Arc::new(CachedPipeline {
152            pipeline,
153            layout: wgpu_layout
154        });
155        cache.insert(label, cached.clone());
156        Ok(cached)
157    }
158
159    /// Pops an idle buffer set matching `shape`, or builds a fresh one.
160    fn checkout_buffers(
161        &self,
162        shape: BufferShape,
163        pipeline: &CachedPipeline,
164        data: &RenderInstanceData
165    ) -> Result<BufferSet, String> {
166        if let Some(set) = self
167            .buffer_pool
168            .lock()
169            .unwrap()
170            .get_mut(&shape)
171            .and_then(Vec::pop)
172        {
173            return Ok(set);
174        }
175
176        let label = shape.label;
177        let params = Buffer::new(
178            data,
179            &format!("{label}-params"),
180            shape.params_size,
181            BufferUsage::UNIFORM | BufferUsage::COPY_DST,
182            None
183        );
184        let output_size = shape.output_len * std::mem::size_of::<f32>();
185        let output = Buffer::new(
186            data,
187            &format!("{label}-output"),
188            output_size,
189            BufferUsage::STORAGE | BufferUsage::COPY_SRC,
190            None
191        );
192        let inputs: Vec<Buffer> = shape
193            .input_lens
194            .iter()
195            .enumerate()
196            .map(|(i, &len)| {
197                Buffer::new(
198                    data,
199                    &format!("{label}-input{i}"),
200                    len * std::mem::size_of::<f32>(),
201                    BufferUsage::STORAGE | BufferUsage::COPY_DST,
202                    None
203                )
204            })
205            .collect();
206        let staging = Buffer::new(
207            data,
208            &format!("{label}-staging"),
209            output_size,
210            BufferUsage::MAP_READ | BufferUsage::COPY_DST,
211            None
212        );
213
214        let mut entries = vec![
215            BindGroupBuilder::buffer(0, &params),
216            BindGroupBuilder::buffer(1, &output),
217        ];
218        for (i, buf) in inputs.iter().enumerate() {
219            entries.push(BindGroupBuilder::buffer(2 + i as u32, buf));
220        }
221        let bind_group = BindGroupBuilder::build(label, data, &pipeline.layout, &entries)
222            .map_err(|e| format!("{label}: {e:?}"))?;
223
224        Ok(BufferSet {
225            params,
226            output,
227            inputs,
228            staging,
229            bind_group
230        })
231    }
232}
wde_renderer/compute.rs

wde_renderer/
compute.rs
