//! The shader that allocates the clustered object ID buffer.
//!
//! The clustered object ID buffer consists of many tightly-packed
//! variable-length arrays of clustered object IDs. The offset to the start of
//! each list is stored in `ClusterOffsetsAndCounts`. To allocate the lists in
//! the buffer, the number of objects of each type in each cluster must be known.
//!
//! Since the lists are tightly packed, determining the offsets in the global
//! buffer is a [prefix sum] problem. To deal with the fact that workgroup sizes
//! are limited to 256 in `wgpu`, and we will usually have more clusters than
//! that, we use a two-pass approach:
//!
//! 1. First, the *local* allocation pass runs on workgroups of 256 clusters
//! each. A [Hillis-Steele scan] is performed to allocate the position in the
//! buffer for each 256 clusters, relative to the first cluster in the workgroup.
//!
//! 2. Next, the *global* allocation pass runs a sequential loop over each 256
//! clusters to calculate the final offsets relative to the previous 256 cluster
//! chunk.
//!
//! At the end of this process, the clusters will have been assigned to their
//! final positions in the list.
//!
//! [prefix sum]: https://en.wikipedia.org/wiki/Prefix_sum
//!
//! [Hillis-Steele scan]: https://en.wikipedia.org/wiki/Prefix_sum#Algorithm_1:_Shorter_span,_more_parallel
import package::cluster::ClusterMetadata;
import package::render::mesh_view_types::{ClusterOffsetsAndCounts, Lights};
@group(0) @binding(0) var<storage, read_write> offsets_and_counts: ClusterOffsetsAndCounts;
@group(0) @binding(1) var<uniform> lights: Lights;
@group(0) @binding(2) var<storage, read_write> clustering_metadata: ClusterMetadata;
@group(0) @binding(3) var<storage, read_write> scratchpad_offsets_and_counts:
ClusterOffsetsAndCounts;
/// The offset of the first clustered light within the buffer for each thread.
var<workgroup> block_offsets: array<u32, 256>;
/// The local allocation pass that, for each chunk of 256 clusters, calculates
/// the offset for each cluster in the chunk.
@compute @workgroup_size(256, 1, 1)
fn allocate_local_main(
@builtin(local_invocation_id) local_id: vec3<u32>,
@builtin(workgroup_id) group_id: vec3<u32>,
@builtin(global_invocation_id) global_id: vec3<u32>
) {
let cluster_count =
lights.cluster_dimensions.x *
lights.cluster_dimensions.y *
lights.cluster_dimensions.z;
let block_start = group_id.x * 256u;
let block_end = min(block_start + 256u, cluster_count);
// Initialize the block offsets.
block_offsets[local_id.x] = 0u;
workgroupBarrier();
// Compute our "block offset" relative to the cluster before us.
// After a prefix scan, the offsets will be correct.
if (global_id.x < block_end && local_id.x < 255u) {
block_offsets[local_id.x + 1] = cluster_object_count(global_id.x);
}
workgroupBarrier();
// Do the Hillis-Steele scan.
for (var offset = 1u; offset < 256u; offset *= 2u) {
var term = 0u;
if (local_id.x >= offset) {
term = block_offsets[local_id.x - offset];
}
workgroupBarrier();
block_offsets[local_id.x] += term;
workgroupBarrier();
}
if (global_id.x < block_end) {
// Now write in the local offset.
offsets_and_counts.data[global_id.x][0].x = block_offsets[local_id.x];
// Zero out the scratchpad counts in preparation for the populate phase
// of the rasterizer.
scratchpad_offsets_and_counts.data[global_id.x][0u] = vec4(0u);
scratchpad_offsets_and_counts.data[global_id.x][1u] = vec4(0u);
}
}
/// The global allocation pass that propagates the offsets from each chunk of 256
/// clusters to later chunks sequentially.
@compute @workgroup_size(256, 1, 1)
fn allocate_global_main(@builtin(local_invocation_id) local_id: vec3<u32>) {
let cluster_count =
lights.cluster_dimensions.x *
lights.cluster_dimensions.y *
lights.cluster_dimensions.z;
// March along the chunks sequentially, accumulating as we go.
var current_offset = 0u;
for (var i = 0u; i < cluster_count; i += 256u) {
if (i + local_id.x < cluster_count) {
offsets_and_counts.data[i + local_id.x][0].x += current_offset;
}
storageBarrier();
if (i + 255u < cluster_count) {
current_offset = offsets_and_counts.data[i + 255u][0].x +
cluster_object_count(i + 255u);
}
}
storageBarrier();
// Write in the final size. This will be read back to the CPU so that the
// buffer can be resized if necessary.
if (local_id.x == 0u) {
clustering_metadata.index_list_size = offsets_and_counts.data[cluster_count - 1][0].x +
cluster_object_count(cluster_count - 1);
}
}
/// Returns the total number of objects in the given cluster.
///
/// The per-type counts are laid out as:
/// [0] = (offset, point lights, spot lights, rect lights)
/// [1] = (reflection probes, irradiance volumes, decals, pad)
fn cluster_object_count(cluster_index: u32) -> u32 {
return
offsets_and_counts.data[cluster_index][0].y +
offsets_and_counts.data[cluster_index][0].z +
offsets_and_counts.data[cluster_index][0].w +
offsets_and_counts.data[cluster_index][1].x +
offsets_and_counts.data[cluster_index][1].y +
offsets_and_counts.data[cluster_index][1].z;
}