wgpu_primitives/sort/
sorter.rs

1use wgpu::util::DeviceExt;
2
3use crate::Error;
4use crate::common;
5use crate::context::Context;
6use crate::scan::Scanner;
7use crate::sort::pipeline::SortPipeline;
8
9const RADIX_PASSES: u32 = 16;
10const WORKSPACE_GROWTH_BYTES: u64 = 16 * 1024 * 1024;
11const UNIFORM_SIZE_BYTES: u64 = 16;
12
13struct SortWorkspace {
14    capacity_bytes: u64,
15    scratch: wgpu::Buffer,
16    histogram: wgpu::Buffer,
17    scanned_histogram: wgpu::Buffer,
18}
19
20#[derive(Clone, Copy)]
21struct PreparedSort {
22    num_items: u32,
23    num_blocks: u32,
24    size_bytes: u64,
25}
26
27/// Performs an unsigned 32-bit LSD radix sort on a wgpu device.
28pub struct Sorter {
29    device: wgpu::Device,
30    queue: wgpu::Queue,
31    scanner: Scanner,
32    pipeline: SortPipeline,
33    workspace: Option<SortWorkspace>,
34}
35
36impl Sorter {
37    /// Creates a sorter that submits work through an existing wgpu device and queue.
38    pub fn new(device: &wgpu::Device, queue: &wgpu::Queue) -> Self {
39        Self {
40            device: device.clone(),
41            queue: queue.clone(),
42            scanner: Scanner::new(device, queue),
43            pipeline: SortPipeline::new(device),
44            workspace: None,
45        }
46    }
47
48    /// Creates a sorter from the crate's optional convenience context.
49    pub fn from_context(ctx: &Context) -> Self {
50        Self::new(&ctx.device, &ctx.queue)
51    }
52
53    /// Uploads values, sorts them on the GPU, and downloads the sorted result.
54    pub async fn sort(&mut self, input: &[u32]) -> Result<Vec<u32>, Error> {
55        if input.is_empty() {
56            return Ok(Vec::new());
57        }
58
59        let num_items = common::math::checked_u32(input.len() as u64)?;
60        let size_bytes = common::math::checked_byte_size(input.len() as u64, 4)?;
61        let input_buffer = common::buffers::create_storage_buffer(&self.device, input);
62        let output_buffer = common::buffers::create_empty_storage_buffer(&self.device, size_bytes);
63
64        self.sort_gpu_to_gpu(&input_buffer, &output_buffer, num_items)?;
65        common::buffers::download_buffer(&self.device, &self.queue, &output_buffer, size_bytes)
66            .await
67    }
68
69    /// Sorts caller-owned GPU buffers and submits the work immediately.
70    pub fn sort_gpu_to_gpu(
71        &mut self,
72        input: &wgpu::Buffer,
73        output: &wgpu::Buffer,
74        num_items: u32,
75    ) -> Result<(), Error> {
76        let mut encoder = self
77            .device
78            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
79                label: Some("Radix Sort"),
80            });
81        self.record_sort(&mut encoder, input, output, num_items)?;
82        self.queue.submit(Some(encoder.finish()));
83        Ok(())
84    }
85
86    /// Records a GPU radix sort without submitting or waiting for the work.
87    pub fn record_sort(
88        &mut self,
89        encoder: &mut wgpu::CommandEncoder,
90        input: &wgpu::Buffer,
91        output: &wgpu::Buffer,
92        num_items: u32,
93    ) -> Result<(), Error> {
94        if num_items == 0 {
95            return Ok(());
96        }
97
98        let problem = self.describe_sort(num_items)?;
99        common::buffers::validate_buffer(
100            input,
101            "sort input",
102            problem.size_bytes,
103            wgpu::BufferUsages::STORAGE,
104        )?;
105        common::buffers::validate_buffer(
106            output,
107            "sort output",
108            problem.size_bytes,
109            wgpu::BufferUsages::STORAGE,
110        )?;
111
112        self.ensure_workspace(problem.size_bytes)?;
113        self.record_radix_passes(encoder, input, output, problem)
114    }
115
116    fn describe_sort(&self, num_items: u32) -> Result<PreparedSort, Error> {
117        let size_bytes = common::math::checked_byte_size(u64::from(num_items), 4)?;
118        let items_per_block = self.pipeline.vt * self.pipeline.block_size;
119        let num_blocks = num_items.div_ceil(items_per_block);
120
121        Ok(PreparedSort {
122            num_items,
123            num_blocks,
124            size_bytes,
125        })
126    }
127
128    fn ensure_workspace(&mut self, size_bytes: u64) -> Result<(), Error> {
129        let needs_allocation = self
130            .workspace
131            .as_ref()
132            .is_none_or(|workspace| workspace.capacity_bytes < size_bytes);
133        if needs_allocation {
134            self.allocate_workspace(size_bytes)?;
135        }
136        Ok(())
137    }
138
139    fn record_radix_passes(
140        &mut self,
141        encoder: &mut wgpu::CommandEncoder,
142        input: &wgpu::Buffer,
143        output: &wgpu::Buffer,
144        problem: PreparedSort,
145    ) -> Result<(), Error> {
146        let max_dispatch = 65_535;
147        let x_groups = problem.num_blocks.min(max_dispatch);
148        let y_groups = problem.num_blocks.div_ceil(max_dispatch);
149        let histogram_items = problem
150            .num_blocks
151            .checked_mul(4)
152            .ok_or(Error::SizeOverflow)?;
153
154        let workspace = self.workspace.as_ref().expect("sort workspace is prepared");
155        let scanner = &mut self.scanner;
156        let (uniform, uniform_stride) = create_uniform_buffer(&self.device, problem);
157
158        for radix_pass in 0..RADIX_PASSES {
159            let (source, destination) = pass_buffers(radix_pass, input, output, &workspace.scratch);
160            let uniform_offset = u64::from(radix_pass) * uniform_stride;
161            let reduce_bind_group = create_sort_bind_group(
162                &self.device,
163                &self.pipeline.bind_group_layout,
164                "Reduce Bind Group",
165                (source, &workspace.histogram, destination),
166                &uniform,
167                uniform_offset,
168            );
169            let scatter_bind_group = create_sort_bind_group(
170                &self.device,
171                &self.pipeline.bind_group_layout,
172                "Scatter Bind Group",
173                (source, &workspace.scanned_histogram, destination),
174                &uniform,
175                uniform_offset,
176            );
177
178            record_compute_pass(
179                encoder,
180                &self.pipeline.reduce_pipeline,
181                &reduce_bind_group,
182                x_groups,
183                y_groups,
184            );
185            scanner.record_scan(
186                encoder,
187                &workspace.histogram,
188                &workspace.scanned_histogram,
189                histogram_items,
190            )?;
191            record_compute_pass(
192                encoder,
193                &self.pipeline.scatter_pipeline,
194                &scatter_bind_group,
195                x_groups,
196                y_groups,
197            );
198        }
199
200        Ok(())
201    }
202
203    fn allocate_workspace(&mut self, requested_size: u64) -> Result<(), Error> {
204        let capacity = if requested_size < WORKSPACE_GROWTH_BYTES {
205            requested_size
206                .max(4)
207                .checked_next_power_of_two()
208                .ok_or(Error::SizeOverflow)?
209        } else {
210            common::math::checked_align_to(requested_size, WORKSPACE_GROWTH_BYTES)?
211        };
212        let limits = self.device.limits();
213        let buffer_limit = limits
214            .max_buffer_size
215            .min(u64::from(limits.max_storage_buffer_binding_size));
216        if capacity > buffer_limit {
217            return Err(Error::BufferLimitExceeded {
218                requested: capacity,
219                limit: buffer_limit,
220            });
221        }
222
223        let items_per_block = u64::from(self.pipeline.vt * self.pipeline.block_size);
224        let max_blocks = (capacity / 4).div_ceil(items_per_block);
225        let histogram_bytes = common::math::checked_byte_size(max_blocks, 16)?;
226        let histogram_capacity = common::math::checked_align_to(histogram_bytes, 256)?;
227
228        self.workspace = Some(SortWorkspace {
229            capacity_bytes: capacity,
230            scratch: common::buffers::create_empty_storage_buffer(&self.device, capacity),
231            histogram: common::buffers::create_empty_storage_buffer(
232                &self.device,
233                histogram_capacity,
234            ),
235            scanned_histogram: common::buffers::create_empty_storage_buffer(
236                &self.device,
237                histogram_capacity,
238            ),
239        });
240        Ok(())
241    }
242}
243
244fn pass_buffers<'a>(
245    radix_pass: u32,
246    input: &'a wgpu::Buffer,
247    output: &'a wgpu::Buffer,
248    scratch: &'a wgpu::Buffer,
249) -> (&'a wgpu::Buffer, &'a wgpu::Buffer) {
250    if radix_pass == 0 {
251        (input, scratch)
252    } else if radix_pass.is_multiple_of(2) {
253        (output, scratch)
254    } else {
255        (scratch, output)
256    }
257}
258
259fn create_uniform_buffer(device: &wgpu::Device, problem: PreparedSort) -> (wgpu::Buffer, u64) {
260    let uniform_stride =
261        u64::from(device.limits().min_uniform_buffer_offset_alignment).max(UNIFORM_SIZE_BYTES);
262    let words_per_uniform = (uniform_stride / size_of::<u32>() as u64) as usize;
263    let mut data = vec![0_u32; words_per_uniform * RADIX_PASSES as usize];
264
265    for radix_pass in 0..RADIX_PASSES as usize {
266        let offset = radix_pass * words_per_uniform;
267        data[offset..offset + 4].copy_from_slice(&[
268            radix_pass as u32 * 2,
269            problem.num_items,
270            problem.num_blocks,
271            0,
272        ]);
273    }
274
275    let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
276        label: Some("Sort Uniform"),
277        contents: bytemuck::cast_slice(&data),
278        usage: wgpu::BufferUsages::UNIFORM,
279    });
280    (buffer, uniform_stride)
281}
282
283fn create_sort_bind_group(
284    device: &wgpu::Device,
285    layout: &wgpu::BindGroupLayout,
286    label: &'static str,
287    buffers: (&wgpu::Buffer, &wgpu::Buffer, &wgpu::Buffer),
288    uniform: &wgpu::Buffer,
289    uniform_offset: u64,
290) -> wgpu::BindGroup {
291    let (source, histogram, destination) = buffers;
292    device.create_bind_group(&wgpu::BindGroupDescriptor {
293        label: Some(label),
294        layout,
295        entries: &[
296            wgpu::BindGroupEntry {
297                binding: 0,
298                resource: source.as_entire_binding(),
299            },
300            wgpu::BindGroupEntry {
301                binding: 1,
302                resource: histogram.as_entire_binding(),
303            },
304            wgpu::BindGroupEntry {
305                binding: 2,
306                resource: destination.as_entire_binding(),
307            },
308            wgpu::BindGroupEntry {
309                binding: 3,
310                resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
311                    buffer: uniform,
312                    offset: uniform_offset,
313                    size: wgpu::BufferSize::new(UNIFORM_SIZE_BYTES),
314                }),
315            },
316        ],
317    })
318}
319
320fn record_compute_pass(
321    encoder: &mut wgpu::CommandEncoder,
322    pipeline: &wgpu::ComputePipeline,
323    bind_group: &wgpu::BindGroup,
324    x_groups: u32,
325    y_groups: u32,
326) {
327    let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
328    pass.set_pipeline(pipeline);
329    pass.set_bind_group(0, bind_group, &[]);
330    pass.dispatch_workgroups(x_groups, y_groups, 1);
331}
wgpu_primitives/sort/sorter.rs

wgpu_primitives/sort/
sorter.rs
