1#import bevy_pbr::cluster::{
2    CLUSTERABLE_OBJECT_TYPE_DECAL, CLUSTERABLE_OBJECT_TYPE_IRRADIANCE_VOLUME,
3    CLUSTERABLE_OBJECT_TYPE_POINT_LIGHT, CLUSTERABLE_OBJECT_TYPE_REFLECTION_PROBE,
4    CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT, ClusterMetadata, ClusterableObjectZSlice,
5    calculate_sphere_cluster_bounds, compute_view_from_world_scale
6}
7#import bevy_pbr::mesh_view_types::{
8    ClusteredDecals, ClusteredLights, LightProbes, Lights, POINT_LIGHT_FLAGS_SPOT_LIGHT_BIT
9}
10#import bevy_render::view::View
11
12// The shader that divides clusterable objects into Z slices.
13//
14// Treating the cluster froxel space as a grid of size WxHxD, for each
15// clusterable object, we seek to rasterize D overlapping quads into a viewport
16// of size WxH. Each quad represents a *Z slice* of the object. This shader
17// calculates the number of Z slices needed for each object in parallel and
18// prepares the `ClusterableObjectZSlice` data that the rasterization passes
19// will consume. It also updates the indirect draw parameters for the
20// rasterization pass and calculates statistics about the farthest Z value
21// encountered that the CPU can later read back for dynamic froxel range tuning.
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 (either point light or spot light).
79        object_index = id;
80        let flags = clustered_lights.data[object_index].flags;
81        object_type = select(
82            CLUSTERABLE_OBJECT_TYPE_POINT_LIGHT,
83            CLUSTERABLE_OBJECT_TYPE_SPOT_LIGHT,
84            (flags & POINT_LIGHT_FLAGS_SPOT_LIGHT_BIT) != 0u
85        );
86        position = clustered_lights.data[object_index].position_radius.xyz;
87        radius = clustered_lights.data[object_index].range;
88    } else if (id < last_reflection_probe_id) {
89        // We're a reflection probe.
90        object_index = id - last_clustered_light_id;
91        object_type = CLUSTERABLE_OBJECT_TYPE_REFLECTION_PROBE;
92        position = light_probes.reflection_probes[object_index].world_position;
93        radius = light_probes.reflection_probes[object_index].bounding_sphere_radius;
94    } else if (id < last_irradiance_volume_id) {
95        // We're an irradiance volume.
96        object_index = id - last_reflection_probe_id;
97        object_type = CLUSTERABLE_OBJECT_TYPE_IRRADIANCE_VOLUME;
98        position = light_probes.irradiance_volumes[object_index].world_position;
99        radius = light_probes.irradiance_volumes[object_index].bounding_sphere_radius;
100    } else if (id < last_decal_id) {
101        // We're a clustered decal.
102        object_index = id - last_irradiance_volume_id;
103        object_type = CLUSTERABLE_OBJECT_TYPE_DECAL;
104        position = clustered_decals.decals[object_index].world_position;
105        radius = clustered_decals.decals[object_index].bounding_sphere_radius;
106    }
107
108    let view_from_world_scale = compute_view_from_world_scale(view.world_from_view);
109    let is_orthographic = view.clip_from_view[3].w == 1.0;
110
111    // Gather the farthest Z value among all clusters in this workgroup.
112    // We want to do this *before* bailing out below so that all threads hit the
113    // same workgroup barriers, which this function uses.
114    accumulate_farthest_z_value(local_id, position, radius, view_from_world_scale, is_orthographic);
115
116    // Bail out if we have no clusterable object to work on.
117    if (object_type == 0xffffffffu) {
118        return;
119    }
120
121    let cluster_bounds = calculate_sphere_cluster_bounds(
122        position,
123        radius,
124        view.view_from_world,
125        view.clip_from_view,
126        view_from_world_scale,
127        lights.cluster_dimensions.xyz,
128        lights.cluster_factors.zw,
129        is_orthographic
130    );
131
132    // Write out our Z slices.
133    for (var z_slice = cluster_bounds.min.z; z_slice <= cluster_bounds.max.z; z_slice += 1u) {
134        try_write_z_slice(object_index, object_type, z_slice);
135    }
136}
137
138// Writes a Z slice to the list.
139//
140// This silently fails if the list is too small, but it still updates the
141// instance count, which the CPU reads back. So, if the list is too small, the
142// CPU will end up being notified and can resize the buffer.
143fn try_write_z_slice(object_index: u32, object_type: u32, z_slice: u32) {
144    let z_slice_offset = atomicAdd(&cluster_metadata.indirect_draw_params.instance_count, 1u);
145    if (z_slice_offset >= cluster_metadata.z_slice_list_capacity) {
146        return;
147    }
148
149    z_slices[z_slice_offset].object_index = object_index;
150    z_slices[z_slice_offset].object_type = object_type;
151    z_slices[z_slice_offset].z_slice = z_slice;
152}
153
154// Records the farthest Z value for clusterable objects in this workgroup for
155// the CPU to read back.
156fn accumulate_farthest_z_value(
157    local_id: u32,
158    position: vec3<f32>,
159    radius: f32,
160    view_from_world_scale: vec3<f32>,
161    is_orthographic: bool
162) {
163    // Compute the maximum Z extent for our clusterable object.
164    let view_from_world_row_2 = transpose(view.view_from_world)[2];
165    let far_z = dot(-view_from_world_row_2, vec4(position, 1.0)) + radius * view_from_world_scale.z;
166    shared_farthest_z[local_id] = far_z;
167    workgroupBarrier();
168
169    // Reduce in local memory to quickly find the maximum Z extent of the
170    // objects in our workgroup.
171    for (var stride = 32u; stride > 0u; stride /= 2u) {
172        if (local_id < stride) {
173            shared_farthest_z[local_id] = max(
174                shared_farthest_z[local_id],
175                shared_farthest_z[local_id + stride]
176            );
177        }
178        workgroupBarrier();
179    }
180
181    // Only the first thread will continue.
182    if (local_id != 0u) {
183        return;
184    }
185
186    // Have the first thread update the global farthest-Z value.
187    // WGSL has no `atomicMax` for floats, so we encode into a u32 that
188    // preserves float ordering and use integer `atomicMax` instead.
189    atomicMax(&cluster_metadata.farthest_z, f32_bits_to_sortable_u32(bitcast<u32>(shared_farthest_z[0u])));
190}
191
192// Encodes f32 bits into a u32 whose unsigned ordering matches
193// the float's total order. Positive floats get the sign bit set so they
194// sort above negative floats; negative floats get all bits flipped so
195// their ordering is reversed (more negative -> smaller u32).
196//
197// The CPU decodes with `sortable_u32_to_f32_bits` in `gpu.rs`.
198fn f32_bits_to_sortable_u32(bits: u32) -> u32 {
199    let mask = bitcast<u32>(bitcast<i32>(bits) >> 31) | 0x80000000u;
200    return bits ^ mask;
201}