1//! A compute shader that performs scattering sparse updates to
2//! `AtomicSparseBufferVec` types.
3//!
4//! This shader isn't used for every update. Only if the number of updates is
5//! small is this shader used. Otherwise, the standard `write_buffer` `wgpu`
6//! command is used to update the buffer in bulk.
7//!
8//! We issue one thread per *word*, not per element. That allows us to achieve
9//! maximum parallelism, without any loops.
10
11/// Metadata that describes the update.
12struct SparseBufferUpdateMetadata {
13    /// The size of a single element in words.
14    element_size: u32,
15    /// The total number of elements to be updated.
16    updated_element_count: u32,
17};
18
19/// The buffer we're copying to.
20@group(0) @binding(0) var<storage, read_write> dest_buffer: array<u32>;
21/// The buffer we're copying from.
22@group(0) @binding(1) var<storage> src_buffer: array<u32>;
23/// For each element in `src_buffer`, the element in `dest_buffer` that we should
24/// copy it to.
25@group(0) @binding(2) var<storage> indices: array<u32>;
26/// Metadata that describes the operation.
27@group(0) @binding(3) var<uniform> metadata: SparseBufferUpdateMetadata;
28
29@workgroup_size(256, 1, 1)
30@compute
31fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
32    // Calculate which word we are. Remember that this shader executes with one
33    // thread per word.
34    let invocation_index = global_id.x;
35    let total_word_count = metadata.updated_element_count * metadata.element_size;
36    if (invocation_index >= total_word_count) {
37        return;
38    }
39
40    // Calculate which element we are.
41    let element_index = invocation_index / metadata.element_size;
42    // Calculate which word *within* that element we're looking at.
43    let word_index = invocation_index % metadata.element_size;
44
45    // Look up our destination element.
46    let dest_element_index = indices[element_index];
47    // Calculate where we should write our word.
48    let dest_index = dest_element_index * metadata.element_size + word_index;
49    if (dest_index >= arrayLength(&dest_buffer)) {
50        return;
51    }
52
53    // Copy the word over.
54    let src_index = element_index * metadata.element_size + word_index;
55    dest_buffer[dest_index] = src_buffer[src_index];
56}