wgpu_primitives/scan/
scanner.rs

1use super::pipeline::{ScanDispatch, ScanPipeline};
2use crate::{
3    Error, common,
4    context::Context,
5    profiling::{GpuProfile, TimestampRecorder},
6};
7
8#[derive(Clone, Copy)]
9enum ScanMode {
10    Inclusive,
11    Exclusive,
12}
13
14struct ScanRecording<'a> {
15    input: &'a wgpu::Buffer,
16    output: &'a wgpu::Buffer,
17    num_items: u32,
18    mode: ScanMode,
19    profile_prefix: &'a str,
20}
21
22/// Performs inclusive and exclusive unsigned 32-bit prefix scans on a wgpu device.
23pub struct Scanner {
24    pipeline: ScanPipeline,
25    device: wgpu::Device,
26    queue: wgpu::Queue,
27    scratch_buffer: Option<wgpu::Buffer>,
28    scratch_size_bytes: u64,
29}
30
31impl Scanner {
32    /// Creates a scanner that submits work through an existing wgpu device and queue.
33    pub fn new(device: &wgpu::Device, queue: &wgpu::Queue) -> Self {
34        Self {
35            pipeline: ScanPipeline::new(device),
36            device: device.clone(),
37            queue: queue.clone(),
38            scratch_buffer: None,
39            scratch_size_bytes: 0,
40        }
41    }
42
43    /// Creates a scanner from the crate's optional convenience context.
44    pub fn from_context(ctx: &Context) -> Self {
45        Self::new(&ctx.device, &ctx.queue)
46    }
47
48    /// Uploads values, scans them on the GPU, and downloads the inclusive prefixes.
49    pub async fn scan(&mut self, input: &[u32]) -> Result<Vec<u32>, Error> {
50        self.scan_slice(input, ScanMode::Inclusive).await
51    }
52
53    /// Uploads values, scans them on the GPU, and downloads the exclusive prefixes.
54    pub async fn scan_exclusive(&mut self, input: &[u32]) -> Result<Vec<u32>, Error> {
55        self.scan_slice(input, ScanMode::Exclusive).await
56    }
57
58    async fn scan_slice(&mut self, input: &[u32], mode: ScanMode) -> Result<Vec<u32>, Error> {
59        if input.is_empty() {
60            return Ok(Vec::new());
61        }
62
63        let num_items = common::math::checked_u32(input.len() as u64)?;
64        let data_buffer = common::buffers::create_storage_buffer(&self.device, input);
65        let dst_buffer =
66            common::buffers::create_empty_storage_buffer(&self.device, data_buffer.size());
67
68        self.submit_scan(&data_buffer, &dst_buffer, num_items, mode)?;
69
70        common::buffers::download_buffer(&self.device, &self.queue, &dst_buffer, input.len()).await
71    }
72
73    /// Scans caller-owned GPU buffers and submits the work immediately.
74    pub fn scan_gpu_to_gpu(
75        &mut self,
76        input_buf: &wgpu::Buffer,
77        output_buf: &wgpu::Buffer,
78        num_items: u32,
79    ) -> Result<(), Error> {
80        self.submit_scan(input_buf, output_buf, num_items, ScanMode::Inclusive)
81    }
82
83    /// Exclusively scans caller-owned GPU buffers and submits the work immediately.
84    pub fn scan_exclusive_gpu_to_gpu(
85        &mut self,
86        input_buf: &wgpu::Buffer,
87        output_buf: &wgpu::Buffer,
88        num_items: u32,
89    ) -> Result<(), Error> {
90        self.submit_scan(input_buf, output_buf, num_items, ScanMode::Exclusive)
91    }
92
93    fn submit_scan(
94        &mut self,
95        input_buf: &wgpu::Buffer,
96        output_buf: &wgpu::Buffer,
97        num_items: u32,
98        mode: ScanMode,
99    ) -> Result<(), Error> {
100        let mut encoder = self
101            .device
102            .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
103        self.record_scan_with_mode(
104            &mut encoder,
105            ScanRecording {
106                input: input_buf,
107                output: output_buf,
108                num_items,
109                mode,
110                profile_prefix: "scan",
111            },
112            None,
113        )?;
114        self.queue.submit(Some(encoder.finish()));
115        Ok(())
116    }
117
118    /// Profiles an inclusive scan of caller-owned GPU buffers using GPU timestamps.
119    pub async fn profile_scan_gpu_to_gpu(
120        &mut self,
121        input_buf: &wgpu::Buffer,
122        output_buf: &wgpu::Buffer,
123        num_items: u32,
124    ) -> Result<GpuProfile, Error> {
125        self.profile_scan_with_mode(input_buf, output_buf, num_items, ScanMode::Inclusive)
126            .await
127    }
128
129    /// Profiles an exclusive scan of caller-owned GPU buffers using GPU timestamps.
130    pub async fn profile_exclusive_scan_gpu_to_gpu(
131        &mut self,
132        input_buf: &wgpu::Buffer,
133        output_buf: &wgpu::Buffer,
134        num_items: u32,
135    ) -> Result<GpuProfile, Error> {
136        self.profile_scan_with_mode(input_buf, output_buf, num_items, ScanMode::Exclusive)
137            .await
138    }
139
140    async fn profile_scan_with_mode(
141        &mut self,
142        input_buf: &wgpu::Buffer,
143        output_buf: &wgpu::Buffer,
144        num_items: u32,
145        mode: ScanMode,
146    ) -> Result<GpuProfile, Error> {
147        let span_count = self.pipeline.compute_pass_count(num_items);
148        if span_count == 0 {
149            let mut encoder = self
150                .device
151                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
152                    label: Some("Profiled Trivial Scan"),
153                });
154            self.record_scan_with_mode(
155                &mut encoder,
156                ScanRecording {
157                    input: input_buf,
158                    output: output_buf,
159                    num_items,
160                    mode,
161                    profile_prefix: "scan",
162                },
163                None,
164            )?;
165            let submission = self.queue.submit(Some(encoder.finish()));
166            self.device.poll(wgpu::PollType::Wait {
167                submission_index: Some(submission),
168                timeout: None,
169            })?;
170            return Ok(GpuProfile::empty());
171        }
172
173        let mut profiler = TimestampRecorder::new(&self.device, &self.queue, span_count)?;
174        let mut encoder = self
175            .device
176            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
177                label: Some("Profiled Prefix Scan"),
178            });
179        self.record_scan_with_mode(
180            &mut encoder,
181            ScanRecording {
182                input: input_buf,
183                output: output_buf,
184                num_items,
185                mode,
186                profile_prefix: "scan",
187            },
188            Some(&mut profiler),
189        )?;
190        profiler.resolve(&mut encoder);
191        let submission = self.queue.submit(Some(encoder.finish()));
192        profiler.read(&self.device, submission).await
193    }
194
195    /// Records a GPU prefix scan without submitting or waiting for the work.
196    pub fn record_scan(
197        &mut self,
198        encoder: &mut wgpu::CommandEncoder,
199        input_buf: &wgpu::Buffer,
200        output_buf: &wgpu::Buffer,
201        num_items: u32,
202    ) -> Result<(), Error> {
203        self.record_scan_with_mode(
204            encoder,
205            ScanRecording {
206                input: input_buf,
207                output: output_buf,
208                num_items,
209                mode: ScanMode::Inclusive,
210                profile_prefix: "scan",
211            },
212            None,
213        )
214    }
215
216    /// Records an exclusive GPU prefix scan without submitting or waiting for the work.
217    pub fn record_exclusive_scan(
218        &mut self,
219        encoder: &mut wgpu::CommandEncoder,
220        input_buf: &wgpu::Buffer,
221        output_buf: &wgpu::Buffer,
222        num_items: u32,
223    ) -> Result<(), Error> {
224        self.record_scan_with_mode(
225            encoder,
226            ScanRecording {
227                input: input_buf,
228                output: output_buf,
229                num_items,
230                mode: ScanMode::Exclusive,
231                profile_prefix: "scan",
232            },
233            None,
234        )
235    }
236
237    pub(crate) fn record_profiled_scan(
238        &mut self,
239        encoder: &mut wgpu::CommandEncoder,
240        input_buf: &wgpu::Buffer,
241        output_buf: &wgpu::Buffer,
242        num_items: u32,
243        profile_prefix: &str,
244        profiler: &mut TimestampRecorder,
245    ) -> Result<(), Error> {
246        self.record_scan_with_mode(
247            encoder,
248            ScanRecording {
249                input: input_buf,
250                output: output_buf,
251                num_items,
252                mode: ScanMode::Inclusive,
253                profile_prefix,
254            },
255            Some(profiler),
256        )
257    }
258
259    pub(crate) fn compute_pass_count(&self, num_items: u32) -> u32 {
260        self.pipeline.compute_pass_count(num_items)
261    }
262
263    fn record_scan_with_mode(
264        &mut self,
265        encoder: &mut wgpu::CommandEncoder,
266        recording: ScanRecording<'_>,
267        mut profiler: Option<&mut TimestampRecorder>,
268    ) -> Result<(), Error> {
269        let ScanRecording {
270            input,
271            output,
272            num_items,
273            mode,
274            profile_prefix,
275        } = recording;
276        if num_items == 0 {
277            return Ok(());
278        }
279
280        let size_bytes = common::math::checked_byte_size(u64::from(num_items), 4)?;
281        common::buffers::validate_buffer(
282            input,
283            "scan input",
284            size_bytes,
285            wgpu::BufferUsages::COPY_SRC,
286        )?;
287        common::buffers::validate_buffer(
288            output,
289            "scan output",
290            size_bytes,
291            wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::STORAGE,
292        )?;
293
294        if num_items == 1 {
295            match mode {
296                ScanMode::Inclusive => {
297                    encoder.copy_buffer_to_buffer(input, 0, output, 0, size_bytes);
298                }
299                ScanMode::Exclusive => encoder.clear_buffer(output, 0, Some(size_bytes)),
300            }
301            return Ok(());
302        }
303
304        encoder.copy_buffer_to_buffer(input, 0, output, 0, size_bytes);
305
306        self.prepare_scratch(num_items);
307
308        let scratch = self
309            .scratch_buffer
310            .as_ref()
311            .expect("scan scratch exists for multi-element inputs");
312
313        struct Level<'a> {
314            buf: &'a wgpu::Buffer,
315            offset: u64,
316            count: u32,
317        }
318
319        let mut levels = Vec::new();
320        levels.push(Level {
321            buf: output,
322            offset: 0,
323            count: num_items,
324        });
325
326        let mut current_scratch_offset = 0u64;
327
328        loop {
329            let current = levels.last().unwrap();
330            if current.count <= 1 {
331                break;
332            }
333
334            let items_per_block = self.pipeline.vt * self.pipeline.block_size;
335
336            let aux_count = current.count.div_ceil(items_per_block);
337            let aux_size = (aux_count * 4) as u64;
338            let aux_offset = crate::common::math::align_to(current_scratch_offset, 256);
339
340            let scan_pipeline = match (levels.len(), mode) {
341                (1, ScanMode::Exclusive) => &self.pipeline.exclusive_scan_pipeline,
342                _ => &self.pipeline.inclusive_scan_pipeline,
343            };
344            let profile_label = profiler
345                .is_some()
346                .then(|| format!("{profile_prefix}.level.{}", levels.len() - 1));
347
348            self.pipeline.dispatch(
349                &self.device,
350                encoder,
351                ScanDispatch {
352                    pipeline: scan_pipeline,
353                    data: (current.buf, current.offset),
354                    auxiliary: (scratch, aux_offset),
355                    num_items: current.count,
356                    pass_label: "Prefix Scan",
357                    profile_label,
358                },
359                profiler.as_deref_mut(),
360            );
361
362            levels.push(Level {
363                buf: scratch,
364                offset: aux_offset,
365                count: aux_count,
366            });
367            current_scratch_offset = aux_offset + aux_size;
368        }
369
370        for i in (0..levels.len() - 1).rev() {
371            let data_level = &levels[i];
372            let aux_level = &levels[i + 1];
373            let profile_label = profiler
374                .is_some()
375                .then(|| format!("{profile_prefix}.add.{i}"));
376
377            self.pipeline.dispatch(
378                &self.device,
379                encoder,
380                ScanDispatch {
381                    pipeline: &self.pipeline.add_pipeline,
382                    data: (data_level.buf, data_level.offset),
383                    auxiliary: (aux_level.buf, aux_level.offset),
384                    num_items: data_level.count,
385                    pass_label: "Prefix Add",
386                    profile_label,
387                },
388                profiler.as_deref_mut(),
389            );
390        }
391
392        Ok(())
393    }
394
395    fn prepare_scratch(&mut self, num_items: u32) {
396        let needed_bytes = self.pipeline.get_scratch_size(num_items);
397        if self.scratch_buffer.is_none() || needed_bytes > self.scratch_size_bytes {
398            self.scratch_buffer = Some(self.device.create_buffer(&wgpu::BufferDescriptor {
399                label: Some("Scanner Scratch"),
400                size: needed_bytes,
401                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
402                mapped_at_creation: false,
403            }));
404            self.scratch_size_bytes = needed_bytes;
405        }
406    }
407}
wgpu_primitives/scan/scanner.rs

wgpu_primitives/scan/
scanner.rs
