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 pages to be updated.
16    updated_page_count: u32,
17    // The base-2 logarithm of the page size.
18    page_size_log2: u32,
19};
20
21// The buffer we're copying to.
22@group(0) @binding(0) var<storage, read_write> dest_buffer: array<u32>;
23// The buffer we're copying from.
24@group(0) @binding(1) var<storage> src_buffer: array<u32>;
25// For each page in `src_buffer`, the page in `dest_buffer` that we should copy
26// it to.
27@group(0) @binding(2) var<storage> indices: array<u32>;
28// Metadata that describes the operation.
29@group(0) @binding(3) var<uniform> metadata: SparseBufferUpdateMetadata;
30
31@workgroup_size(256, 1, 1)
32@compute
33fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
34    // Calculate which word we are. Remember that this shader executes with one
35    // thread per word.
36    let invocation_index = global_id.x;
37    let total_word_count = (metadata.updated_page_count << metadata.page_size_log2) *
38        metadata.element_size;
39    if (invocation_index >= total_word_count) {
40        return;
41    }
42
43    // Calculate which element we are.
44    let element_index = invocation_index / metadata.element_size;
45    // Calculate which word *within* that element we're looking at.
46    let word_index = invocation_index % metadata.element_size;
47    // Calculate which page we're copying.
48    let update_index = element_index >> metadata.page_size_log2;
49    // Determine which element we're copying within that page.
50    let element_index_in_page = element_index & ((1u << metadata.page_size_log2) - 1u);
51
52    // Look up our destination page.
53    let page_index = indices[update_index];
54    // Calculate where we should write our word.
55    let dest_index = ((page_index << metadata.page_size_log2) + element_index_in_page) *
56        metadata.element_size + word_index;
57    if (dest_index >= arrayLength(&dest_buffer)) {
58        return;
59    }
60
61    // Copy the word over.
62    let src_index = element_index * metadata.element_size + word_index;
63    dest_buffer[dest_index] = src_buffer[src_index];
64}