wgpu_primitives/sort/
key_value_sorter.rs

1use crate::context::Context;
2use crate::{Error, GpuProfile};
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 specialized for the supplied adapter when a measured
35    /// fast path is available.
36    ///
37    /// Discrete NVIDIA Vulkan adapters with 32-wide subgroups use the 8-bit
38    /// radix kernel. Other NVIDIA Vulkan devices use the 4-bit kernel, and all
39    /// remaining adapters use the portable 2-bit kernel.
40    pub fn new_for_adapter(
41        device: &wgpu::Device,
42        queue: &wgpu::Queue,
43        adapter_info: &wgpu::AdapterInfo,
44    ) -> Self {
45        Self {
46            core: RadixSorter::new_for_adapter(device, queue, SortItemKind::KeyValue, adapter_info),
47        }
48    }
49
50    /// Creates a sorter from the crate's optional convenience context.
51    pub fn from_context(ctx: &Context) -> Self {
52        Self::new_for_adapter(&ctx.device, &ctx.queue, &ctx.adapter_info)
53    }
54
55    /// Uploads items, stably sorts them by key, and downloads the result.
56    pub async fn sort(&mut self, input: &[KeyValue]) -> Result<Vec<KeyValue>, Error> {
57        self.core.sort_slice(input).await
58    }
59
60    /// Stably sorts caller-owned GPU buffers and submits the work immediately.
61    pub fn sort_gpu_to_gpu(
62        &mut self,
63        input: &wgpu::Buffer,
64        output: &wgpu::Buffer,
65        num_items: u32,
66    ) -> Result<(), Error> {
67        self.core.sort_gpu_to_gpu(input, output, num_items)
68    }
69
70    /// Profiles a stable GPU-buffer key-value radix sort using GPU timestamps.
71    pub async fn profile_sort_gpu_to_gpu(
72        &mut self,
73        input: &wgpu::Buffer,
74        output: &wgpu::Buffer,
75        num_items: u32,
76    ) -> Result<GpuProfile, Error> {
77        self.core
78            .profile_sort_gpu_to_gpu(input, output, num_items)
79            .await
80    }
81
82    /// Records a stable GPU key-value radix sort without submitting or waiting.
83    pub fn record_sort(
84        &mut self,
85        encoder: &mut wgpu::CommandEncoder,
86        input: &wgpu::Buffer,
87        output: &wgpu::Buffer,
88        num_items: u32,
89    ) -> Result<(), Error> {
90        self.core.record_sort(encoder, input, output, num_items)
91    }
92}
wgpu_primitives/sort/key_value_sorter.rs

wgpu_primitives/sort/
key_value_sorter.rs
