Skip to main content

wgpu_primitives/sort/
key_value_sorter.rs

1use crate::Error;
2use crate::context::Context;
3
4use super::core::RadixSorter;
5use super::pipeline::SortItemKind;
6
7/// A `u32` key and its associated `u32` value.
8#[repr(C)]
9#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, bytemuck::Pod, bytemuck::Zeroable)]
10pub struct KeyValue {
11    pub key: u32,
12    pub value: u32,
13}
14
15impl KeyValue {
16    pub const fn new(key: u32, value: u32) -> Self {
17        Self { key, value }
18    }
19}
20
21/// Performs a stable LSD radix sort of `KeyValue` items by key on a wgpu device.
22pub struct KeyValueSorter {
23    core: RadixSorter,
24}
25
26impl KeyValueSorter {
27    /// Creates a sorter that submits work through an existing wgpu device and queue.
28    pub fn new(device: &wgpu::Device, queue: &wgpu::Queue) -> Self {
29        Self {
30            core: RadixSorter::new(device, queue, SortItemKind::KeyValue),
31        }
32    }
33
34    /// Creates a sorter from the crate's optional convenience context.
35    pub fn from_context(ctx: &Context) -> Self {
36        Self::new(&ctx.device, &ctx.queue)
37    }
38
39    /// Uploads items, stably sorts them by key, and downloads the result.
40    pub async fn sort(&mut self, input: &[KeyValue]) -> Result<Vec<KeyValue>, Error> {
41        self.core.sort_slice(input).await
42    }
43
44    /// Stably sorts caller-owned GPU buffers and submits the work immediately.
45    pub fn sort_gpu_to_gpu(
46        &mut self,
47        input: &wgpu::Buffer,
48        output: &wgpu::Buffer,
49        num_items: u32,
50    ) -> Result<(), Error> {
51        self.core.sort_gpu_to_gpu(input, output, num_items)
52    }
53
54    /// Records a stable GPU key-value radix sort without submitting or waiting.
55    pub fn record_sort(
56        &mut self,
57        encoder: &mut wgpu::CommandEncoder,
58        input: &wgpu::Buffer,
59        output: &wgpu::Buffer,
60        num_items: u32,
61    ) -> Result<(), Error> {
62        self.core.record_sort(encoder, input, output, num_items)
63    }
64}