1use crate::buffer::GpuBuffer;
2use crate::{BindingKind, BindingSpec, ComputePipeline, ComputeRecorder, GpuContext};
3use wgpu::{BindGroup, BindGroupEntry, Device};
4
5const THREADS: u32 = 256;
6const BLOCK: u32 = 256;
7
8fn bucketed_shader(source: &str, buckets: u32) -> String {
9 source.replace("__BUCKETS__", &buckets.to_string())
10}
11
12#[derive(PartialEq, Eq)]
13struct Channels {
14 keys: u64,
15 values: u64,
16 keys_out: u64,
17 values_out: u64,
18 count: u64,
19}
20
21struct BucketBindGroups {
22 channels: Channels,
23 histogram: BindGroup,
24 scatter: BindGroup,
25}
26
27impl BucketBindGroups {
28 #[expect(
29 clippy::too_many_arguments,
30 reason = "bucket sort carries both key/value channels and their outputs explicitly"
31 )]
32 fn build(
33 device: &Device,
34 bucket: &GpuBucketSort,
35 channels: Channels,
36 keys: &GpuBuffer,
37 values: &GpuBuffer,
38 keys_out: &GpuBuffer,
39 values_out: &GpuBuffer,
40 count_holder: &GpuBuffer,
41 ) -> Self {
42 let histogram = bucket.histogram_pipeline.create_bind_group(
43 device,
44 0,
45 &[
46 BindGroupEntry {
47 binding: 0,
48 resource: keys.as_binding(),
49 },
50 BindGroupEntry {
51 binding: 1,
52 resource: bucket.histogram.as_binding(),
53 },
54 BindGroupEntry {
55 binding: 2,
56 resource: bucket.block_histogram.as_binding(),
57 },
58 BindGroupEntry {
59 binding: 3,
60 resource: count_holder.as_binding(),
61 },
62 ],
63 );
64 let scatter = bucket.scatter_pipeline.create_bind_group(
65 device,
66 0,
67 &[
68 BindGroupEntry {
69 binding: 0,
70 resource: keys.as_binding(),
71 },
72 BindGroupEntry {
73 binding: 1,
74 resource: values.as_binding(),
75 },
76 BindGroupEntry {
77 binding: 2,
78 resource: bucket.cursor.as_binding(),
79 },
80 BindGroupEntry {
81 binding: 3,
82 resource: bucket.block_prefix.as_binding(),
83 },
84 BindGroupEntry {
85 binding: 4,
86 resource: keys_out.as_binding(),
87 },
88 BindGroupEntry {
89 binding: 5,
90 resource: values_out.as_binding(),
91 },
92 BindGroupEntry {
93 binding: 6,
94 resource: count_holder.as_binding(),
95 },
96 ],
97 );
98 Self {
99 channels,
100 histogram,
101 scatter,
102 }
103 }
104}
105
106pub struct GpuBucketSort {
107 histogram_pipeline: ComputePipeline,
108 prefix_pipeline: ComputePipeline,
109 block_prefix_pipeline: ComputePipeline,
110 scatter_pipeline: ComputePipeline,
111 prefix_group: BindGroup,
112 block_prefix_group: BindGroup,
113 histogram: GpuBuffer,
114 cursor: GpuBuffer,
115 block_histogram: GpuBuffer,
116 block_prefix: GpuBuffer,
117 bindings: std::sync::Mutex<Option<BucketBindGroups>>,
118}
119
120impl GpuBucketSort {
121 pub fn new(context: &GpuContext, label: &str, buckets: u32, data_capacity: u32) -> Self {
122 let device = context.device();
123 let histogram_spec = [
124 BindingSpec {
125 binding: 0,
126 kind: BindingKind::ReadOnlyStorage,
127 },
128 BindingSpec {
129 binding: 1,
130 kind: BindingKind::ReadWriteStorage,
131 },
132 BindingSpec {
133 binding: 2,
134 kind: BindingKind::ReadWriteStorage,
135 },
136 BindingSpec {
137 binding: 3,
138 kind: BindingKind::ReadOnlyStorage,
139 },
140 ];
141 let scatter_spec = [
142 BindingSpec {
143 binding: 0,
144 kind: BindingKind::ReadOnlyStorage,
145 },
146 BindingSpec {
147 binding: 1,
148 kind: BindingKind::ReadOnlyStorage,
149 },
150 BindingSpec {
151 binding: 2,
152 kind: BindingKind::ReadWriteStorage,
153 },
154 BindingSpec {
155 binding: 3,
156 kind: BindingKind::ReadOnlyStorage,
157 },
158 BindingSpec {
159 binding: 4,
160 kind: BindingKind::ReadWriteStorage,
161 },
162 BindingSpec {
163 binding: 5,
164 kind: BindingKind::ReadWriteStorage,
165 },
166 BindingSpec {
167 binding: 6,
168 kind: BindingKind::ReadOnlyStorage,
169 },
170 ];
171 let prefix_spec = [
172 BindingSpec {
173 binding: 0,
174 kind: BindingKind::ReadWriteStorage,
175 },
176 BindingSpec {
177 binding: 1,
178 kind: BindingKind::ReadWriteStorage,
179 },
180 ];
181 let block_prefix_spec = [
182 BindingSpec {
183 binding: 0,
184 kind: BindingKind::ReadWriteStorage,
185 },
186 BindingSpec {
187 binding: 1,
188 kind: BindingKind::ReadWriteStorage,
189 },
190 ];
191 let histogram_pipeline = context.compute_pipeline(
192 &format!("{label} histogram"),
193 &bucketed_shader(include_str!("shaders/bucket_histogram.wgsl"), buckets),
194 "main",
195 &[&histogram_spec[..]],
196 THREADS,
197 );
198 let prefix_pipeline = context.compute_pipeline(
199 &format!("{label} prefix"),
200 &bucketed_shader(include_str!("shaders/bucket_prefix.wgsl"), buckets),
201 "main",
202 &[&prefix_spec[..]],
203 THREADS,
204 );
205 let block_prefix_pipeline = context.compute_pipeline(
206 &format!("{label} block prefix"),
207 &bucketed_shader(include_str!("shaders/bucket_block_prefix.wgsl"), buckets),
208 "main",
209 &[&block_prefix_spec[..]],
210 THREADS,
211 );
212 let scatter_pipeline = context.compute_pipeline(
213 &format!("{label} scatter"),
214 &bucketed_shader(include_str!("shaders/bucket_scatter.wgsl"), buckets),
215 "main",
216 &[&scatter_spec[..]],
217 THREADS,
218 );
219 let bucket_bytes = (buckets as u64) * 4;
220 let data_blocks = data_capacity.div_ceil(BLOCK);
221 let block_bytes = data_blocks as u64 * buckets as u64 * 4;
222 let histogram = GpuBuffer::zeroed(
223 device,
224 &format!("{label} histogram"),
225 bucket_bytes,
226 wgpu::BufferUsages::STORAGE,
227 );
228 let cursor = GpuBuffer::new(
229 device,
230 &format!("{label} cursor"),
231 bucket_bytes,
232 wgpu::BufferUsages::STORAGE,
233 );
234 let block_histogram = GpuBuffer::zeroed(
235 device,
236 &format!("{label} block histogram"),
237 block_bytes,
238 wgpu::BufferUsages::STORAGE,
239 );
240 let block_prefix = GpuBuffer::new(
241 device,
242 &format!("{label} block prefix"),
243 block_bytes,
244 wgpu::BufferUsages::STORAGE,
245 );
246 let prefix_group = prefix_pipeline.create_bind_group(
247 device,
248 0,
249 &[
250 BindGroupEntry {
251 binding: 0,
252 resource: histogram.as_binding(),
253 },
254 BindGroupEntry {
255 binding: 1,
256 resource: cursor.as_binding(),
257 },
258 ],
259 );
260 let block_prefix_group = block_prefix_pipeline.create_bind_group(
261 device,
262 0,
263 &[
264 BindGroupEntry {
265 binding: 0,
266 resource: block_histogram.as_binding(),
267 },
268 BindGroupEntry {
269 binding: 1,
270 resource: block_prefix.as_binding(),
271 },
272 ],
273 );
274 Self {
275 histogram_pipeline,
276 prefix_pipeline,
277 block_prefix_pipeline,
278 scatter_pipeline,
279 prefix_group,
280 block_prefix_group,
281 histogram,
282 cursor,
283 block_histogram,
284 block_prefix,
285 bindings: std::sync::Mutex::new(None),
286 }
287 }
288
289 #[expect(
290 clippy::too_many_arguments,
291 reason = "bucket sort carries both key/value channels and their outputs explicitly"
292 )]
293 fn bindings(
294 &self,
295 device: &Device,
296 channels: Channels,
297 keys: &GpuBuffer,
298 values: &GpuBuffer,
299 keys_out: &GpuBuffer,
300 values_out: &GpuBuffer,
301 count_holder: &GpuBuffer,
302 ) -> std::sync::MutexGuard<'_, Option<BucketBindGroups>> {
303 let mut guard = self.bindings.lock().unwrap();
304 if guard
305 .as_ref()
306 .is_none_or(|cached| cached.channels != channels)
307 {
308 *guard = Some(BucketBindGroups::build(
309 device,
310 self,
311 channels,
312 keys,
313 values,
314 keys_out,
315 values_out,
316 count_holder,
317 ));
318 }
319 guard
320 }
321
322 #[expect(
323 clippy::too_many_arguments,
324 reason = "bucket sort carries both key/value channels and their outputs explicitly"
325 )]
326 pub fn sort(
327 &self,
328 device: &Device,
329 recorder: &mut ComputeRecorder,
330 count_holder: &GpuBuffer,
331 args: &GpuBuffer,
332 keys: &GpuBuffer,
333 values: &GpuBuffer,
334 keys_out: &GpuBuffer,
335 values_out: &GpuBuffer,
336 ) {
337 let channels = Channels {
338 keys: keys.token(),
339 values: values.token(),
340 keys_out: keys_out.token(),
341 values_out: values_out.token(),
342 count: count_holder.token(),
343 };
344 let guard = self.bindings(
345 device,
346 channels,
347 keys,
348 values,
349 keys_out,
350 values_out,
351 count_holder,
352 );
353 let bindings = guard.as_ref().expect("bindings ensured just above");
354 recorder.record_indirect(&self.histogram_pipeline, &[&bindings.histogram], args, 16);
355 recorder.record(&self.prefix_pipeline, &[&self.prefix_group], 1);
356 recorder.record(&self.block_prefix_pipeline, &[&self.block_prefix_group], 1);
357 recorder.record_indirect(&self.scatter_pipeline, &[&bindings.scatter], args, 16);
358 }
359}