1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
use crate::context::Context;
use crate::{Error, GpuProfile};
use super::core::RadixSorter;
use super::pipeline::SortItemKind;
/// A `u32` key and its associated `u32` value.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, bytemuck::Pod, bytemuck::Zeroable)]
pub struct KeyValue {
pub key: u32,
pub value: u32,
}
impl KeyValue {
pub const fn new(key: u32, value: u32) -> Self {
Self { key, value }
}
}
/// Performs a stable LSD radix sort of `KeyValue` items by key on a wgpu device.
pub struct KeyValueSorter {
core: RadixSorter,
}
impl KeyValueSorter {
/// Creates a sorter that submits work through an existing wgpu device and queue.
pub fn new(device: &wgpu::Device, queue: &wgpu::Queue) -> Self {
Self {
core: RadixSorter::new(device, queue, SortItemKind::KeyValue),
}
}
/// Creates a sorter specialized for the supplied adapter when a measured
/// fast path is available.
///
/// Discrete NVIDIA Vulkan adapters with 32-wide subgroups use the 8-bit
/// radix kernel. Other NVIDIA Vulkan devices use the 4-bit kernel, and all
/// remaining adapters use the portable 2-bit kernel.
pub fn new_for_adapter(
device: &wgpu::Device,
queue: &wgpu::Queue,
adapter_info: &wgpu::AdapterInfo,
) -> Self {
Self {
core: RadixSorter::new_for_adapter(device, queue, SortItemKind::KeyValue, adapter_info),
}
}
/// Creates a sorter from the crate's optional convenience context.
pub fn from_context(ctx: &Context) -> Self {
Self::new_for_adapter(&ctx.device, &ctx.queue, &ctx.adapter_info)
}
/// Uploads items, stably sorts them by key, and downloads the result.
pub async fn sort(&mut self, input: &[KeyValue]) -> Result<Vec<KeyValue>, Error> {
self.core.sort_slice(input).await
}
/// Stably sorts caller-owned GPU buffers and submits the work immediately.
pub fn sort_gpu_to_gpu(
&mut self,
input: &wgpu::Buffer,
output: &wgpu::Buffer,
num_items: u32,
) -> Result<(), Error> {
self.core.sort_gpu_to_gpu(input, output, num_items)
}
/// Profiles a stable GPU-buffer key-value radix sort using GPU timestamps.
pub async fn profile_sort_gpu_to_gpu(
&mut self,
input: &wgpu::Buffer,
output: &wgpu::Buffer,
num_items: u32,
) -> Result<GpuProfile, Error> {
self.core
.profile_sort_gpu_to_gpu(input, output, num_items)
.await
}
/// Records a stable GPU key-value radix sort without submitting or waiting.
pub fn record_sort(
&mut self,
encoder: &mut wgpu::CommandEncoder,
input: &wgpu::Buffer,
output: &wgpu::Buffer,
num_items: u32,
) -> Result<(), Error> {
self.core.record_sort(encoder, input, output, num_items)
}
}