1//! The shader that allocates the clustered object ID buffer.
2//!
3//! The clustered object ID buffer consists of many tightly-packed
4//! variable-length arrays of clustered object IDs. The offset to the start of
5//! each list is stored in `ClusterOffsetsAndCounts`. To allocate the lists in
6//! the buffer, the number of objects of each type in each cluster must be known.
7//!
8//! Since the lists are tightly packed, determining the offsets in the global
9//! buffer is a [prefix sum] problem. To deal with the fact that workgroup sizes
10//! are limited to 256 in `wgpu`, and we will usually have more clusters than
11//! that, we use a two-pass approach:
12//!
13//! 1. First, the *local* allocation pass runs on workgroups of 256 clusters
14//! each. A [Hillis-Steele scan] is performed to allocate the position in the
15//! buffer for each 256 clusters, relative to the first cluster in the workgroup.
16//!
17//! 2. Next, the *global* allocation pass runs a sequential loop over each 256
18//! clusters to calculate the final offsets relative to the previous 256 cluster
19//! chunk.
20//!
21//! At the end of this process, the clusters will have been assigned to their
22//! final positions in the list.
23//!
24//! [prefix sum]: https://en.wikipedia.org/wiki/Prefix_sum
25//!
26//! [Hillis-Steele scan]: https://en.wikipedia.org/wiki/Prefix_sum#Algorithm_1:_Shorter_span,_more_parallel
27
28import package::cluster::ClusterMetadata;
29import package::render::mesh_view_types::{ClusterOffsetsAndCounts, Lights};
30
31@group(0) @binding(0) var<storage, read_write> offsets_and_counts: ClusterOffsetsAndCounts;
32@group(0) @binding(1) var<uniform> lights: Lights;
33@group(0) @binding(2) var<storage, read_write> clustering_metadata: ClusterMetadata;
34@group(0) @binding(3) var<storage, read_write> scratchpad_offsets_and_counts:
35    ClusterOffsetsAndCounts;
36
37/// The offset of the first clustered light within the buffer for each thread.
38var<workgroup> block_offsets: array<u32, 256>;
39
40/// The local allocation pass that, for each chunk of 256 clusters, calculates
41/// the offset for each cluster in the chunk.
42@compute @workgroup_size(256, 1, 1)
43fn allocate_local_main(
44    @builtin(local_invocation_id) local_id: vec3<u32>,
45    @builtin(workgroup_id) group_id: vec3<u32>,
46    @builtin(global_invocation_id) global_id: vec3<u32>
47) {
48    let cluster_count =
49        lights.cluster_dimensions.x *
50        lights.cluster_dimensions.y *
51        lights.cluster_dimensions.z;
52
53    let block_start = group_id.x * 256u;
54    let block_end = min(block_start + 256u, cluster_count);
55
56    // Initialize the block offsets.
57    block_offsets[local_id.x] = 0u;
58    workgroupBarrier();
59
60    // Compute our "block offset" relative to the cluster before us.
61    // After a prefix scan, the offsets will be correct.
62    if (global_id.x < block_end && local_id.x < 255u) {
63        block_offsets[local_id.x + 1] = cluster_object_count(global_id.x);
64    }
65    workgroupBarrier();
66
67    // Do the Hillis-Steele scan.
68    for (var offset = 1u; offset < 256u; offset *= 2u) {
69        var term = 0u;
70        if (local_id.x >= offset) {
71            term = block_offsets[local_id.x - offset];
72        }
73        workgroupBarrier();
74        block_offsets[local_id.x] += term;
75        workgroupBarrier();
76    }
77
78    if (global_id.x < block_end) {
79        // Now write in the local offset.
80        offsets_and_counts.data[global_id.x][0].x = block_offsets[local_id.x];
81
82        // Zero out the scratchpad counts in preparation for the populate phase
83        // of the rasterizer.
84        scratchpad_offsets_and_counts.data[global_id.x][0u] = vec4(0u);
85        scratchpad_offsets_and_counts.data[global_id.x][1u] = vec4(0u);
86    }
87}
88
89/// The global allocation pass that propagates the offsets from each chunk of 256
90/// clusters to later chunks sequentially.
91@compute @workgroup_size(256, 1, 1)
92fn allocate_global_main(@builtin(local_invocation_id) local_id: vec3<u32>) {
93    let cluster_count =
94        lights.cluster_dimensions.x *
95        lights.cluster_dimensions.y *
96        lights.cluster_dimensions.z;
97
98    // March along the chunks sequentially, accumulating as we go.
99    var current_offset = 0u;
100    for (var i = 0u; i < cluster_count; i += 256u) {
101        if (i + local_id.x < cluster_count) {
102            offsets_and_counts.data[i + local_id.x][0].x += current_offset;
103        }
104        storageBarrier();
105
106        if (i + 255u < cluster_count) {
107            current_offset = offsets_and_counts.data[i + 255u][0].x +
108                cluster_object_count(i + 255u);
109        }
110    }
111    storageBarrier();
112
113    // Write in the final size. This will be read back to the CPU so that the
114    // buffer can be resized if necessary.
115    if (local_id.x == 0u) {
116        clustering_metadata.index_list_size = offsets_and_counts.data[cluster_count - 1][0].x +
117            cluster_object_count(cluster_count - 1);
118    }
119}
120
121/// Returns the total number of objects in the given cluster.
122///
123/// The per-type counts are laid out as:
124///   [0] = (offset, point lights, spot lights, rect lights)
125///   [1] = (reflection probes, irradiance volumes, decals, pad)
126fn cluster_object_count(cluster_index: u32) -> u32 {
127    return
128        offsets_and_counts.data[cluster_index][0].y +
129        offsets_and_counts.data[cluster_index][0].z +
130        offsets_and_counts.data[cluster_index][0].w +
131        offsets_and_counts.data[cluster_index][1].x +
132        offsets_and_counts.data[cluster_index][1].y +
133        offsets_and_counts.data[cluster_index][1].z;
134}