1use crate::buffer::GpuBuffer;
2use crate::{BindingKind, BindingSpec, ComputePipeline, ComputeRecorder, GpuContext};
3use wgpu::{BindGroup, BindGroupEntry, Device};
4
5const THREADS: u32 = 256;
6const RADIX_PASSES: usize = 8;
7const BIN_COUNT: usize = 256;
8
9fn shifted_shader(source: &str, shift: u32) -> String {
10 source.replace("__SHIFT__", &shift.to_string())
11}
12
13#[derive(PartialEq, Eq)]
14struct SortChannels {
15 count: u64,
16 keys_lo: u64,
17 keys_hi: u64,
18 values: u64,
19 keys_lo_out: u64,
20 keys_hi_out: u64,
21 values_out: u64,
22}
23
24impl SortChannels {
25 fn of(
26 count: &GpuBuffer,
27 keys_lo: &GpuBuffer,
28 keys_hi: &GpuBuffer,
29 values: &GpuBuffer,
30 keys_lo_out: &GpuBuffer,
31 keys_hi_out: &GpuBuffer,
32 values_out: &GpuBuffer,
33 ) -> Self {
34 Self {
35 count: count.token(),
36 keys_lo: keys_lo.token(),
37 keys_hi: keys_hi.token(),
38 values: values.token(),
39 keys_lo_out: keys_lo_out.token(),
40 keys_hi_out: keys_hi_out.token(),
41 values_out: values_out.token(),
42 }
43 }
44}
45
46#[derive(PartialEq, Eq)]
47struct SortArgs {
48 args: u64,
49}
50
51struct SortBindGroups {
52 channels: SortChannels,
53 args: SortArgs,
54 histogram: [BindGroup; 2],
55 scatter: [BindGroup; 2],
56}
57
58impl SortBindGroups {
59 #[expect(
60 clippy::too_many_arguments,
61 reason = "radix sort carries three key channels and their outputs explicitly"
62 )]
63 fn build(
64 device: &Device,
65 sort: &GpuSort,
66 channels: SortChannels,
67 args: SortArgs,
68 keys_lo: &GpuBuffer,
69 keys_hi: &GpuBuffer,
70 values: &GpuBuffer,
71 keys_lo_out: &GpuBuffer,
72 keys_hi_out: &GpuBuffer,
73 values_out: &GpuBuffer,
74 count_holder: &GpuBuffer,
75 ) -> Self {
76 let in_lo = [keys_lo, keys_lo_out];
77 let in_hi = [keys_hi, keys_hi_out];
78 let in_values = [values, values_out];
79 let out_lo = [keys_lo_out, keys_lo];
80 let out_hi = [keys_hi_out, keys_hi];
81 let out_values = [values_out, values];
82 let histogram = std::array::from_fn(|parity| {
83 sort.histogram_pipelines[0].create_bind_group(
84 device,
85 0,
86 &[
87 BindGroupEntry {
88 binding: 0,
89 resource: in_lo[parity].as_binding(),
90 },
91 BindGroupEntry {
92 binding: 1,
93 resource: in_hi[parity].as_binding(),
94 },
95 BindGroupEntry {
96 binding: 2,
97 resource: sort.histogram.as_binding(),
98 },
99 BindGroupEntry {
100 binding: 3,
101 resource: sort.block_histogram.as_binding(),
102 },
103 BindGroupEntry {
104 binding: 4,
105 resource: count_holder.as_binding(),
106 },
107 ],
108 )
109 });
110 let scatter = std::array::from_fn(|parity| {
111 sort.scatter_pipelines[0].create_bind_group(
112 device,
113 0,
114 &[
115 BindGroupEntry {
116 binding: 0,
117 resource: in_lo[parity].as_binding(),
118 },
119 BindGroupEntry {
120 binding: 1,
121 resource: in_hi[parity].as_binding(),
122 },
123 BindGroupEntry {
124 binding: 2,
125 resource: in_values[parity].as_binding(),
126 },
127 BindGroupEntry {
128 binding: 3,
129 resource: sort.cursor.as_binding(),
130 },
131 BindGroupEntry {
132 binding: 4,
133 resource: sort.block_prefix.as_binding(),
134 },
135 BindGroupEntry {
136 binding: 5,
137 resource: out_lo[parity].as_binding(),
138 },
139 BindGroupEntry {
140 binding: 6,
141 resource: out_hi[parity].as_binding(),
142 },
143 BindGroupEntry {
144 binding: 7,
145 resource: out_values[parity].as_binding(),
146 },
147 BindGroupEntry {
148 binding: 8,
149 resource: count_holder.as_binding(),
150 },
151 ],
152 )
153 });
154 Self {
155 channels,
156 args,
157 histogram,
158 scatter,
159 }
160 }
161}
162
163pub struct GpuSort {
164 histogram_pipelines: [ComputePipeline; RADIX_PASSES],
165 scatter_pipelines: [ComputePipeline; RADIX_PASSES],
166 prefix_pipeline: ComputePipeline,
167 prefix_group: BindGroup,
168 histogram: GpuBuffer,
169 cursor: GpuBuffer,
170 block_histogram: GpuBuffer,
171 block_prefix: GpuBuffer,
172 bindings: std::sync::Mutex<Option<SortBindGroups>>,
173}
174
175impl GpuSort {
176 pub fn new(context: &GpuContext, label: &str, capacity: u32) -> Self {
177 let device = context.device();
178 let blocks = capacity.div_ceil(THREADS).max(1);
179 let histogram_spec = [
180 BindingSpec {
181 binding: 0,
182 kind: BindingKind::ReadOnlyStorage,
183 },
184 BindingSpec {
185 binding: 1,
186 kind: BindingKind::ReadOnlyStorage,
187 },
188 BindingSpec {
189 binding: 2,
190 kind: BindingKind::ReadWriteStorage,
191 },
192 BindingSpec {
193 binding: 3,
194 kind: BindingKind::ReadWriteStorage,
195 },
196 BindingSpec {
197 binding: 4,
198 kind: BindingKind::ReadOnlyStorage,
199 },
200 ];
201 let scatter_spec = [
202 BindingSpec {
203 binding: 0,
204 kind: BindingKind::ReadOnlyStorage,
205 },
206 BindingSpec {
207 binding: 1,
208 kind: BindingKind::ReadOnlyStorage,
209 },
210 BindingSpec {
211 binding: 2,
212 kind: BindingKind::ReadOnlyStorage,
213 },
214 BindingSpec {
215 binding: 3,
216 kind: BindingKind::ReadOnlyStorage,
217 },
218 BindingSpec {
219 binding: 4,
220 kind: BindingKind::ReadOnlyStorage,
221 },
222 BindingSpec {
223 binding: 5,
224 kind: BindingKind::ReadWriteStorage,
225 },
226 BindingSpec {
227 binding: 6,
228 kind: BindingKind::ReadWriteStorage,
229 },
230 BindingSpec {
231 binding: 7,
232 kind: BindingKind::ReadWriteStorage,
233 },
234 BindingSpec {
235 binding: 8,
236 kind: BindingKind::ReadOnlyStorage,
237 },
238 ];
239 let prefix_spec = [
240 BindingSpec {
241 binding: 0,
242 kind: BindingKind::ReadWriteStorage,
243 },
244 BindingSpec {
245 binding: 1,
246 kind: BindingKind::ReadWriteStorage,
247 },
248 BindingSpec {
249 binding: 2,
250 kind: BindingKind::ReadWriteStorage,
251 },
252 BindingSpec {
253 binding: 3,
254 kind: BindingKind::ReadWriteStorage,
255 },
256 ];
257 let prefix_pipeline = context.compute_pipeline(
258 &format!("{label} prefix"),
259 include_str!("shaders/sort_prefix.wgsl"),
260 "main",
261 &[&prefix_spec[..]],
262 THREADS,
263 );
264 let histogram_shader = include_str!("shaders/sort_histogram.wgsl");
265 let histogram_pipelines = std::array::from_fn(|index| {
266 context.compute_pipeline(
267 &format!("{label} histogram {index}"),
268 &shifted_shader(histogram_shader, (index * 8) as u32),
269 "main",
270 &[&histogram_spec[..]],
271 THREADS,
272 )
273 });
274 let scatter_shader = include_str!("shaders/sort_scatter.wgsl");
275 let scatter_pipelines = std::array::from_fn(|index| {
276 context.compute_pipeline(
277 &format!("{label} scatter {index}"),
278 &shifted_shader(scatter_shader, (index * 8) as u32),
279 "main",
280 &[&scatter_spec[..]],
281 THREADS,
282 )
283 });
284 let histogram = GpuBuffer::zeroed(
285 device,
286 &format!("{label} histogram"),
287 (BIN_COUNT * 4) as u64,
288 wgpu::BufferUsages::STORAGE,
289 );
290 let cursor = GpuBuffer::new(
291 device,
292 &format!("{label} cursor"),
293 (BIN_COUNT * 4) as u64,
294 wgpu::BufferUsages::STORAGE,
295 );
296 let block_bytes = blocks as u64 * BIN_COUNT as u64 * 4;
297 let block_histogram = GpuBuffer::zeroed(
298 device,
299 &format!("{label} block histogram"),
300 block_bytes,
301 wgpu::BufferUsages::STORAGE,
302 );
303 let block_prefix = GpuBuffer::new(
304 device,
305 &format!("{label} block prefix"),
306 block_bytes,
307 wgpu::BufferUsages::STORAGE,
308 );
309 let prefix_group = prefix_pipeline.create_bind_group(
310 device,
311 0,
312 &[
313 BindGroupEntry {
314 binding: 0,
315 resource: histogram.as_binding(),
316 },
317 BindGroupEntry {
318 binding: 1,
319 resource: cursor.as_binding(),
320 },
321 BindGroupEntry {
322 binding: 2,
323 resource: block_histogram.as_binding(),
324 },
325 BindGroupEntry {
326 binding: 3,
327 resource: block_prefix.as_binding(),
328 },
329 ],
330 );
331 Self {
332 histogram_pipelines,
333 scatter_pipelines,
334 prefix_pipeline,
335 prefix_group,
336 histogram,
337 cursor,
338 block_histogram,
339 block_prefix,
340 bindings: std::sync::Mutex::new(None),
341 }
342 }
343
344 #[expect(
345 clippy::too_many_arguments,
346 reason = "radix sort carries three key channels and their outputs explicitly"
347 )]
348 fn bindings(
349 &self,
350 device: &Device,
351 channels: SortChannels,
352 args: SortArgs,
353 keys_lo: &GpuBuffer,
354 keys_hi: &GpuBuffer,
355 values: &GpuBuffer,
356 keys_lo_out: &GpuBuffer,
357 keys_hi_out: &GpuBuffer,
358 values_out: &GpuBuffer,
359 count_holder: &GpuBuffer,
360 ) -> std::sync::MutexGuard<'_, Option<SortBindGroups>> {
361 let mut guard = self.bindings.lock().unwrap();
362 if guard
363 .as_ref()
364 .is_none_or(|cached| cached.channels != channels || cached.args != args)
365 {
366 *guard = Some(SortBindGroups::build(
367 device,
368 self,
369 channels,
370 args,
371 keys_lo,
372 keys_hi,
373 values,
374 keys_lo_out,
375 keys_hi_out,
376 values_out,
377 count_holder,
378 ));
379 }
380 guard
381 }
382
383 fn encode_all_passes(
384 &self,
385 recorder: &mut ComputeRecorder,
386 bindings: &SortBindGroups,
387 args: &GpuBuffer,
388 lo_words: u32,
389 hi_words: u32,
390 ) {
391 let pass_ranges = [0..lo_words, 4..(4 + hi_words)];
392 for (executed, pass_index) in pass_ranges.into_iter().flatten().enumerate() {
393 let parity = executed % 2;
394 let pass_index = pass_index as usize;
395 recorder.record_indirect(
396 &self.histogram_pipelines[pass_index],
397 &[&bindings.histogram[parity]],
398 args,
399 16,
400 );
401 recorder.record(&self.prefix_pipeline, &[&self.prefix_group], 1);
402 recorder.record_indirect(
403 &self.scatter_pipelines[pass_index],
404 &[&bindings.scatter[parity]],
405 args,
406 16,
407 );
408 }
409 }
410
411 #[expect(
412 clippy::too_many_arguments,
413 reason = "radix sort carries three key channels and their outputs explicitly"
414 )]
415 pub fn sort_64(
416 &self,
417 device: &Device,
418 recorder: &mut ComputeRecorder,
419 count_holder: &GpuBuffer,
420 args: &GpuBuffer,
421 lo_words: u32,
422 hi_words: u32,
423 keys_lo: &GpuBuffer,
424 keys_hi: &GpuBuffer,
425 values: &GpuBuffer,
426 keys_lo_out: &GpuBuffer,
427 keys_hi_out: &GpuBuffer,
428 values_out: &GpuBuffer,
429 ) {
430 let channels = SortChannels::of(
431 count_holder,
432 keys_lo,
433 keys_hi,
434 values,
435 keys_lo_out,
436 keys_hi_out,
437 values_out,
438 );
439 let guard = self.bindings(
440 device,
441 channels,
442 SortArgs { args: args.token() },
443 keys_lo,
444 keys_hi,
445 values,
446 keys_lo_out,
447 keys_hi_out,
448 values_out,
449 count_holder,
450 );
451 let bindings = guard.as_ref().expect("bindings ensured just above");
452 self.encode_all_passes(recorder, bindings, args, lo_words, hi_words);
453 }
454
455 pub fn debug_histogram(&self) -> &GpuBuffer {
456 &self.histogram
457 }
458}