1//! The shader that divides clusterable objects into Z slices.
2//!
3//! Treating the cluster froxel space as a grid of size WxHxD, for each
4//! clusterable object, we seek to rasterize D overlapping quads into a viewport
5//! of size WxH. Each quad represents a *Z slice* of the object. This shader
6//! calculates the number of Z slices needed for each object in parallel and
7//! prepares the `ClusterableObjectZSlice` data that the rasterization passes
8//! will consume. It also updates the indirect draw parameters for the
9//! rasterization pass and calculates statistics about the farthest Z value
10//! encountered that the CPU can later read back for dynamic froxel range tuning.
11
12import package::cluster::{
13    CLUSTERABLE_OBJECT_TYPE_DECAL, CLUSTERABLE_OBJECT_TYPE_IRRADIANCE_VOLUME,
14    CLUSTERABLE_OBJECT_TYPE_POINT_LIGHT, CLUSTERABLE_OBJECT_TYPE_REFLECTION_PROBE,
15    CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT, CLUSTERABLE_OBJECT_TYPE_RECT_LIGHT, ClusterMetadata,
16    ClusterableObjectZSlice, calculate_sphere_cluster_bounds, compute_view_from_world_scale
17};
18import package::render::mesh_view_types::{
19    ClusteredDecals, ClusteredLights, LightProbes, Lights, POINT_LIGHT_FLAGS_SPOT_LIGHT_BIT, POINT_LIGHT_FLAGS_RECT_LIGHT_BIT
20};
21import bevy_render::view::View;
22
23/// Metadata, including the indirect draw parameters that we write to.
24@group(0) @binding(0) var<storage, read_write> cluster_metadata: ClusterMetadata;
25/// The list of clusterable object Z slices that we write to.
26@group(0) @binding(1) var<storage, read_write> z_slices: array<ClusterableObjectZSlice>;
27/// Information about each light.
28@group(0) @binding(2) var<storage> clustered_lights: ClusteredLights;
29/// Information about each light probe (reflection probe or irradiance volume).
30@group(0) @binding(3) var<uniform> light_probes: LightProbes;
31/// Information about each clustered decal.
32@group(0) @binding(4) var<storage> clustered_decals: ClusteredDecals;
33/// Information about the clusters as a whole, including the dimensions of the
34/// cluster grid.
35@group(0) @binding(5) var<uniform> lights: Lights;
36/// Information about the view.
37@group(0) @binding(6) var<uniform> view: View;
38
39/// A temporary workgroup-local buffer used to accelerate the "farthest depth of
40/// any object" calculation.
41var<workgroup> shared_farthest_z: array<f32, 64>;
42
43/// The shader entry point.
44///
45/// We have one invocation per clusterable object.
46@compute @workgroup_size(64, 1, 1)
47fn z_slice_main(
48    @builtin(global_invocation_id) global_invocation_id: vec3<u32>,
49    @builtin(local_invocation_id) local_invocation_id: vec3<u32>
50) {
51    let id = global_invocation_id.x;
52    let local_id = local_invocation_id.x;
53
54    var object_index: u32 = 0u;
55    var object_type: u32 = 0xffffffffu;
56    var position: vec3<f32> = vec3<f32>(0.0);
57    var radius: f32 = 0.0;
58
59    // Figure out what the bounds are for each type of clusterable object, in
60    // preparation to determining which kind of object we're clustering.
61    // In thread order, the threads are assigned to cluster clusterable lights,
62    // reflection probes, irradiance volumes, and decals, in that order.
63    // It might look like we should have done prefix sum on
64    // `clustered_light_count`, `reflection_probe_count`, etc. to avoid all
65    // this ID math. But that would make life harder for plugins that want to
66    // add clustered objects in compute shaders, because they might have to
67    // update multiple fields (atomically!) when adding, for example, a light.
68    let last_clustered_light_id = cluster_metadata.clustered_light_count;
69    let last_reflection_probe_id = last_clustered_light_id +
70        cluster_metadata.reflection_probe_count;
71    let last_irradiance_volume_id = last_reflection_probe_id +
72        cluster_metadata.irradiance_volume_count;
73    let last_decal_id = last_irradiance_volume_id + cluster_metadata.decal_count;
74
75    // Figure out which type of object we are, and calculate our position and range.
76    // We use a sphere to conservatively construct our AABB.
77    if (id < last_clustered_light_id) {
78        // We're a light (point, spot, or rect light).
79        object_index = id;
80        let flags = clustered_lights.data[object_index].flags;
81        if ((flags & POINT_LIGHT_FLAGS_RECT_LIGHT_BIT) != 0u) {
82            object_type = CLUSTERABLE_OBJECT_TYPE_RECT_LIGHT;
83        } else {
84            object_type = select(
85                CLUSTERABLE_OBJECT_TYPE_POINT_LIGHT,
86                CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT,
87                (flags & POINT_LIGHT_FLAGS_SPOT_LIGHT_BIT) != 0u
88            );
89        }
90        position = clustered_lights.data[object_index].position_radius.xyz;
91        radius = clustered_lights.data[object_index].range;
92    } else if (id < last_reflection_probe_id) {
93        // We're a reflection probe.
94        object_index = id - last_clustered_light_id;
95        object_type = CLUSTERABLE_OBJECT_TYPE_REFLECTION_PROBE;
96        position = light_probes.reflection_probes[object_index].world_position;
97        radius = light_probes.reflection_probes[object_index].bounding_sphere_radius;
98    } else if (id < last_irradiance_volume_id) {
99        // We're an irradiance volume.
100        object_index = id - last_reflection_probe_id;
101        object_type = CLUSTERABLE_OBJECT_TYPE_IRRADIANCE_VOLUME;
102        position = light_probes.irradiance_volumes[object_index].world_position;
103        radius = light_probes.irradiance_volumes[object_index].bounding_sphere_radius;
104    } else if (id < last_decal_id) {
105        // We're a clustered decal.
106        object_index = id - last_irradiance_volume_id;
107        object_type = CLUSTERABLE_OBJECT_TYPE_DECAL;
108        position = clustered_decals.decals[object_index].world_position;
109        radius = clustered_decals.decals[object_index].bounding_sphere_radius;
110    }
111
112    let view_from_world_scale = compute_view_from_world_scale(view.world_from_view);
113    let is_orthographic = view.clip_from_view[3].w == 1.0;
114
115    // Gather the farthest Z value among all clusters in this workgroup.
116    // We want to do this *before* bailing out below so that all threads hit the
117    // same workgroup barriers, which this function uses.
118    accumulate_farthest_z_value(local_id, position, radius, view_from_world_scale, is_orthographic);
119
120    // Bail out if we have no clusterable object to work on.
121    if (object_type == 0xffffffffu) {
122        return;
123    }
124
125    let cluster_bounds = calculate_sphere_cluster_bounds(
126        position,
127        radius,
128        view.view_from_world,
129        view.clip_from_view,
130        view_from_world_scale,
131        lights.cluster_dimensions.xyz,
132        lights.cluster_factors.zw,
133        is_orthographic
134    );
135
136    // Write out our Z slices.
137    for (var z_slice = cluster_bounds.min.z; z_slice <= cluster_bounds.max.z; z_slice += 1u) {
138        try_write_z_slice(object_index, object_type, z_slice);
139    }
140}
141
142/// Writes a Z slice to the list.
143///
144/// This silently fails if the list is too small, but it still updates the
145/// instance count, which the CPU reads back. So, if the list is too small, the
146/// CPU will end up being notified and can resize the buffer.
147fn try_write_z_slice(object_index: u32, object_type: u32, z_slice: u32) {
148    let z_slice_offset = atomicAdd(&cluster_metadata.indirect_draw_params.instance_count, 1u);
149    if (z_slice_offset >= cluster_metadata.z_slice_list_capacity) {
150        return;
151    }
152
153    z_slices[z_slice_offset].object_index = object_index;
154    z_slices[z_slice_offset].object_type = object_type;
155    z_slices[z_slice_offset].z_slice = z_slice;
156}
157
158/// Records the farthest Z value for clusterable objects in this workgroup for
159/// the CPU to read back.
160fn accumulate_farthest_z_value(
161    local_id: u32,
162    position: vec3<f32>,
163    radius: f32,
164    view_from_world_scale: vec3<f32>,
165    is_orthographic: bool
166) {
167    // Compute the maximum Z extent for our clusterable object.
168    let view_from_world_row_2 = transpose(view.view_from_world)[2];
169    let far_z = dot(-view_from_world_row_2, vec4(position, 1.0)) + radius * view_from_world_scale.z;
170    shared_farthest_z[local_id] = far_z;
171    workgroupBarrier();
172
173    // Reduce in local memory to quickly find the maximum Z extent of the
174    // objects in our workgroup.
175    for (var stride = 32u; stride > 0u; stride /= 2u) {
176        if (local_id < stride) {
177            shared_farthest_z[local_id] = max(
178                shared_farthest_z[local_id],
179                shared_farthest_z[local_id + stride]
180            );
181        }
182        workgroupBarrier();
183    }
184
185    // Only the first thread will continue.
186    if (local_id != 0u) {
187        return;
188    }
189
190    // Have the first thread update the global farthest-Z value.
191    // WGSL has no `atomicMax` for floats, so we encode into a u32 that
192    // preserves float ordering and use integer `atomicMax` instead.
193    atomicMax(&cluster_metadata.farthest_z, f32_bits_to_sortable_u32(bitcast<u32>(shared_farthest_z[0u])));
194}
195
196/// Encodes f32 bits into a u32 whose unsigned ordering matches
197/// the float's total order. Positive floats get the sign bit set so they
198/// sort above negative floats; negative floats get all bits flipped so
199/// their ordering is reversed (more negative -> smaller u32).
200///
201/// The CPU decodes with `sortable_u32_to_f32_bits` in `gpu.rs`.
202fn f32_bits_to_sortable_u32(bits: u32) -> u32 {
203    let mask = bitcast<u32>(bitcast<i32>(bits) >> 31) | 0x80000000u;
204    return bits ^ mask;
205}