1#import bevy_pbr::cluster::ClusterMetadata
2#import bevy_pbr::mesh_view_types::{ClusterOffsetsAndCounts, Lights}
3
4// The shader that allocates the clustered object ID buffer.
5//
6// The clustered object ID buffer consists of many tightly-packed
7// variable-length arrays of clustered object IDs. The offset to the start of
8// each list is stored in `ClusterOffsetsAndCounts`. To allocate the lists in
9// the buffer, the number of objects of each type in each cluster must be known.
10//
11// Since the lists are tightly packed, determining the offsets in the global
12// buffer is a [prefix sum] problem. To deal with the fact that workgroup sizes
13// are limited to 256 in `wgpu`, and we will usually have more clusters than
14// that, we use a two-pass approach:
15//
16// 1. First, the *local* allocation pass runs on workgroups of 256 clusters
17// each. A [Hillis-Steele scan] is performed to allocate the position in the
18// buffer for each 256 clusters, relative to the first cluster in the workgroup.
19//
20// 2. Next, the *global* allocation pass runs a sequential loop over each 256
21// clusters to calculate the final offsets relative to the previous 256 cluster
22// chunk.
23//
24// At the end of this process, the clusters will have been assigned to their
25// final positions in the list.
26//
27// [prefix sum]: https://en.wikipedia.org/wiki/Prefix_sum
28//
29// [Hillis-Steele scan]: https://en.wikipedia.org/wiki/Prefix_sum#Algorithm_1:_Shorter_span,_more_parallel
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.
122fn cluster_object_count(cluster_index: u32) -> u32 {
123    return
124        offsets_and_counts.data[cluster_index][0].y +
125        offsets_and_counts.data[cluster_index][0].z +
126        offsets_and_counts.data[cluster_index][0].w +
127        offsets_and_counts.data[cluster_index][1].x +
128        offsets_and_counts.data[cluster_index][1].y;
129}