1//! GPU mesh transforming and culling.
2//!
3//! This is a compute shader that expands each `MeshInputUniform` out to a full
4//! `MeshUniform` for each view before rendering. (Thus `MeshInputUniform` and
5//! `MeshUniform` are in a 1:N relationship.) It runs in parallel for all meshes
6//! for all views. As part of this process, the shader gathers each mesh's
7//! transform on the previous frame and writes it into the `MeshUniform` so that
8//! TAA works. It also performs frustum culling and occlusion culling, if
9//! requested.
10//!
11//! If occlusion culling is on, this shader runs twice: once to prepare the
12//! meshes that were visible last frame, and once to prepare the meshes that
13//! weren't visible last frame but became visible this frame. The two invocations
14//! are known as *early mesh preprocessing* and *late mesh preprocessing*
15//! respectively.
16
17import bevy_render::occlusion_culling::mesh_preprocess_types::{
18    IndirectParametersMetadata, MeshInput, PreviousMeshInput, PreprocessWorkItem
19};
20import package::render::mesh_types::{
21    Mesh, MESH_FLAGS_AABB_BASED_VISIBILITY_RANGE_BIT, MESH_FLAGS_NO_FRUSTUM_CULLING_BIT,
22    MESH_FLAGS_VISIBILITY_RANGE_INDEX_BITS
23};
24import package::render::mesh_view_bindings::view;
25import package::render::occlusion_culling;
26import package::prepass::bindings::previous_view_uniforms;
27import package::render::view_transformations::{
28    position_world_to_ndc, position_world_to_view, ndc_to_uv, view_z_to_depth_ndc,
29    position_world_to_prev_ndc, position_world_to_prev_view, prev_view_z_to_depth_ndc
30};
31import bevy_render::maths;
32import bevy_render::view::View;
33
34/// Information about each mesh instance needed to cull it on GPU.
35///
36/// At the moment, this just consists of its axis-aligned bounding box (AABB).
37struct MeshCullingData {
38    /// The 3D center of the AABB in model space, padded with an extra unused
39    /// float value.
40    aabb_center: vec4<f32>,
41    /// The 3D extents of the AABB in model space, divided by two, padded with
42    /// an extra unused float value.
43    aabb_half_extents: vec4<f32>,
44}
45
46/// The parameters for the indirect compute dispatch for the late mesh
47/// preprocessing phase.
48struct LatePreprocessWorkItemIndirectParameters {
49    /// The number of workgroups we're going to dispatch.
50    ///
51    /// This value should always be equal to `ceil(work_item_count / 64)`.
52    ///
53    /// In the late phase this buffer is bound read-only, and the atomic built-ins
54    /// are only defined for `read_write` storage, so the fields are plain `u32`
55    /// there and `atomic<u32>` in the early phase. `u32` and `atomic<u32>` are
56    /// layout-identical, so the buffer's memory layout is unchanged either way.
57    @if(LATE_PHASE)
58    dispatch_x: u32,
59    @else   // LATE_PHASE
60    dispatch_x: atomic<u32>,
61    /// The number of workgroups in the Y direction; always 1.
62    dispatch_y: u32,
63    /// The number of workgroups in the Z direction; always 1.
64    dispatch_z: u32,
65    /// The precise number of work items.
66    @if(LATE_PHASE)
67    work_item_count: u32,
68    @else   // LATE_PHASE
69    work_item_count: atomic<u32>,
70    /// Padding.
71    ///
72    /// This isn't the usual structure padding; it's needed because some hardware
73    /// requires indirect compute dispatch parameters to be aligned on 64-byte
74    /// boundaries.
75    pad: vec4<u32>,
76}
77
78/// These have to be in a structure because of Naga limitations on DX12.
79struct Immediates {
80    /// The offset into the `late_preprocess_work_item_indirect_parameters`
81    /// buffer.
82    late_preprocess_work_item_indirect_offset: u32,
83}
84
85/// The current frame's `MeshInput`.
86@group(0) @binding(3) var<storage> current_input: array<MeshInput>;
87/// The `MeshInput` values from the previous frame.
88@group(0) @binding(4) var<storage> previous_input: array<PreviousMeshInput>;
89/// Indices into the `MeshInput` buffer.
90///
91/// There may be many indices that map to the same `MeshInput`.
92@group(0) @binding(5) var<storage> work_items: array<PreprocessWorkItem>;
93/// The output array of `Mesh`es.
94@group(0) @binding(6) var<storage, read_write> output: array<Mesh>;
95
96@if(INDIRECT)
97/// The array of indirect parameters for drawcalls.
98@group(0) @binding(7) var<storage, read_write> indirect_parameters_metadata:
99    array<IndirectParametersMetadata>;
100
101@if(FRUSTUM_CULLING) {
102/// Data needed to cull the meshes.
103///
104/// At the moment, this consists only of AABBs.
105@group(0) @binding(9) var<storage> mesh_culling_data: array<MeshCullingData>;
106
107@group(0) @binding(10) var<storage> visibility_ranges: array<vec4<f32>>;
108}
109
110@if(OCCLUSION_CULLING)
111@group(0) @binding(11) var depth_pyramid: texture_2d<f32>;
112
113@if(OCCLUSION_CULLING && EARLY_PHASE) {
114@group(0) @binding(12) var<storage, read_write> late_preprocess_work_items:
115    array<PreprocessWorkItem>;
116
117@group(0) @binding(13) var<storage, read_write> late_preprocess_work_item_indirect_parameters:
118    array<LatePreprocessWorkItemIndirectParameters>;
119}
120
121@if(OCCLUSION_CULLING && LATE_PHASE)
122@group(0) @binding(13) var<storage, read> late_preprocess_work_item_indirect_parameters:
123    array<LatePreprocessWorkItemIndirectParameters>;
124
125@if(OCCLUSION_CULLING)
126var<immediate> immediates: Immediates;
127
128@if(FRUSTUM_CULLING)
129/// Returns true if the view frustum intersects an oriented bounding box (OBB).
130///
131/// `aabb_center.w` should be 1.0.
132fn view_frustum_intersects_obb(
133    world_from_local: mat4x4<f32>,
134    aabb_center: vec4<f32>,
135    aabb_half_extents: vec3<f32>,
136) -> bool {
137
138    for (var i = 0; i < 5; i += 1) {
139        // Calculate relative radius of the sphere associated with this plane.
140        let plane_normal = view.frustum[i];
141        let relative_radius = dot(
142            abs(
143                vec3(
144                    dot(plane_normal.xyz, world_from_local[0].xyz),
145                    dot(plane_normal.xyz, world_from_local[1].xyz),
146                    dot(plane_normal.xyz, world_from_local[2].xyz),
147                )
148            ),
149            aabb_half_extents
150        );
151
152        // Check the frustum plane.
153        if (!maths::sphere_intersects_plane_half_space(
154                plane_normal, aabb_center, relative_radius)) {
155            return false;
156        }
157    }
158
159    return true;
160}
161
162@compute
163@workgroup_size(64)
164fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
165    // Figure out our instance index. If this thread doesn't correspond to any
166    // index, bail.
167    let instance_index = global_invocation_id.x;
168
169@if(LATE_PHASE)
170    if (instance_index >= late_preprocess_work_item_indirect_parameters[
171            immediates.late_preprocess_work_item_indirect_offset].work_item_count) {
172        return;
173    }
174@else   // LATE_PHASE
175    if (instance_index >= arrayLength(&work_items)) {
176        return;
177    }
178
179    // Unpack the work item.
180    let input_index = work_items[instance_index].input_index;
181@if(INDIRECT)
182    let indirect_parameters_index = work_items[instance_index].output_or_indirect_parameters_index;
183
184    // If we're the first mesh instance in this batch, write the index of our
185    // `MeshInput` into the appropriate slot so that the indirect parameters
186    // building shader can access it.
187@if(INDIRECT && !LATE_PHASE)
188    if (instance_index == 0u) || (work_items[instance_index - 1].output_or_indirect_parameters_index != indirect_parameters_index) {
189        indirect_parameters_metadata[indirect_parameters_index].mesh_index = input_index;
190    }
191
192@if(!INDIRECT)
193    let mesh_output_index = work_items[instance_index].output_or_indirect_parameters_index;
194
195    // Unpack the input matrix.
196    let world_from_local_affine_transpose = current_input[input_index].world_from_local;
197    let world_from_local = maths::affine3_to_square(world_from_local_affine_transpose);
198
199@if(FRUSTUM_CULLING) {
200    // Frustum cull if necessary.
201    if ((current_input[input_index].flags & MESH_FLAGS_NO_FRUSTUM_CULLING_BIT) == 0u) {
202        let aabb_center = mesh_culling_data[input_index].aabb_center.xyz;
203        let aabb_half_extents = mesh_culling_data[input_index].aabb_half_extents.xyz;
204
205        // Do an OBB-based frustum cull.
206        let model_center = world_from_local * vec4(aabb_center, 1.0);
207        if (!view_frustum_intersects_obb(world_from_local, model_center, aabb_half_extents)) {
208            return;
209        }
210    }
211
212    // Visibility range cull if necessary.
213    let visibility_buffer_array_len = arrayLength(&visibility_ranges);
214    let visibility_buffer_index =
215        current_input[input_index].flags & MESH_FLAGS_VISIBILITY_RANGE_INDEX_BITS;
216    if (visibility_buffer_index < visibility_buffer_array_len) {
217        let lod_range = visibility_ranges[visibility_buffer_index];
218
219        // If we're using the AABB as the mesh center, determine its world space position.
220        // Otherwise, just use the center of the transform.
221        var world_pos: vec3<f32>;
222        if ((current_input[input_index].flags & MESH_FLAGS_AABB_BASED_VISIBILITY_RANGE_BIT) != 0u) {
223            let aabb_center = mesh_culling_data[input_index].aabb_center.xyz;
224            world_pos = (world_from_local * vec4(aabb_center, 1.0)).xyz;
225        } else {
226            world_pos = world_from_local[3].xyz;
227        }
228
229        let camera_distance = length(world_pos - view.lod_view_world_position);
230        // `x` is the minimum range; `w` is the largest range.
231        if (camera_distance < lod_range.x || camera_distance >= lod_range.w) {
232            return;
233        }
234    }
235}
236
237    // See whether the `MeshInputUniform` was updated on this frame. If it
238    // wasn't, then we know the transforms of this mesh must be identical to
239    // those on the previous frame, and therefore we don't need to access the
240    // `previous_input_index` (in fact, we can't; that index are only valid for
241    // one frame and will be invalid).
242    let timestamp = current_input[input_index].timestamp;
243    let mesh_changed_this_frame = timestamp == view.frame_count;
244
245    // Look up the previous model matrix, if it could have been.
246    let previous_input_index = current_input[input_index].previous_input_index;
247    var previous_world_from_local_affine_transpose: mat3x4<f32>;
248    if (mesh_changed_this_frame && previous_input_index != 0xffffffffu) {
249        previous_world_from_local_affine_transpose =
250            previous_input[previous_input_index].world_from_local;
251    } else {
252        previous_world_from_local_affine_transpose = world_from_local_affine_transpose;
253    }
254    let previous_world_from_local =
255        maths::affine3_to_square(previous_world_from_local_affine_transpose);
256
257    // Occlusion cull if necessary. This is done by calculating the screen-space
258    // axis-aligned bounding box (AABB) of the mesh and testing it against the
259    // appropriate level of the depth pyramid (a.k.a. hierarchical Z-buffer). If
260    // no part of the AABB is in front of the corresponding pixel quad in the
261    // hierarchical Z-buffer, then this mesh must be occluded, and we can skip
262    // rendering it.
263@if(OCCLUSION_CULLING) {
264    let aabb_center = mesh_culling_data[input_index].aabb_center.xyz;
265    let aabb_half_extents = mesh_culling_data[input_index].aabb_half_extents.xyz;
266
267    // Initialize the AABB and the maximum depth.
268    let infinity = bitcast<f32>(0x7f800000u);
269    let neg_infinity = bitcast<f32>(0xff800000u);
270    var aabb = vec4(infinity, infinity, neg_infinity, neg_infinity);
271    var max_depth_view = neg_infinity;
272
273    // Build up the AABB by taking each corner of this mesh's OBB, transforming
274    // it, and updating the AABB and depth accordingly.
275    for (var i = 0u; i < 8u; i += 1u) {
276        let local_pos = aabb_center + select(
277            vec3(-1.0),
278            vec3(1.0),
279            vec3((i & 1) != 0, (i & 2) != 0, (i & 4) != 0)
280        ) * aabb_half_extents;
281
282        // If we're in the early phase, we're testing against the last frame's
283        // depth buffer, so we need to use the previous frame's transform.
284        @if(EARLY_PHASE)
285        let prev_world_pos = (previous_world_from_local * vec4(local_pos, 1.0)).xyz;
286        @if(EARLY_PHASE)
287        let view_pos = position_world_to_prev_view(prev_world_pos);
288        @if(EARLY_PHASE)
289        let ndc_pos = position_world_to_prev_ndc(prev_world_pos);
290        // Otherwise, if this is the late phase, we use the current frame's
291        // transform.
292        @if(!EARLY_PHASE)
293        let world_pos = (world_from_local * vec4(local_pos, 1.0)).xyz;
294        @if(!EARLY_PHASE)
295        let view_pos = position_world_to_view(world_pos);
296        @if(!EARLY_PHASE)
297        let ndc_pos = position_world_to_ndc(world_pos);
298
299        let uv_pos = ndc_to_uv(ndc_pos.xy);
300
301        // Update the AABB and maximum view-space depth.
302        aabb = vec4(min(aabb.xy, uv_pos), max(aabb.zw, uv_pos));
303        max_depth_view = max(max_depth_view, view_pos.z);
304    }
305
306    // Clip to the near plane to avoid the NDC depth becoming negative.
307    @if(EARLY_PHASE)
308    max_depth_view = min(-previous_view_uniforms.clip_from_view[3][2], max_depth_view);
309    @else   // EARLY_PHASE
310    max_depth_view = min(-view.clip_from_view[3][2], max_depth_view);
311
312    // Figure out the depth of the occluder, and compare it to our own depth.
313
314    let aabb_pixel_size = occlusion_culling::get_aabb_size_in_pixels(aabb, depth_pyramid);
315    let occluder_depth_ndc =
316        occlusion_culling::get_occluder_depth(aabb, aabb_pixel_size, depth_pyramid);
317
318    @if(EARLY_PHASE)
319    let max_depth_ndc = prev_view_z_to_depth_ndc(max_depth_view);
320    @else   // EARLY_PHASE
321    let max_depth_ndc = view_z_to_depth_ndc(max_depth_view);
322
323    // Are we culled out?
324    if (max_depth_ndc < occluder_depth_ndc) {
325@if(EARLY_PHASE) {
326        // If this is the early phase, we need to make a note of this mesh so
327        // that we examine it again in the late phase, so that we handle the
328        // case in which a mesh that was invisible last frame became visible in
329        // this frame.
330        let output_work_item_index = atomicAdd(&late_preprocess_work_item_indirect_parameters[
331            immediates.late_preprocess_work_item_indirect_offset].work_item_count, 1u);
332        if (output_work_item_index % 64u == 0u) {
333            // Our workgroup size is 64, and the indirect parameters for the
334            // late mesh preprocessing phase are counted in workgroups, so if
335            // we're the first thread in this workgroup, bump the workgroup
336            // count.
337            atomicAdd(&late_preprocess_work_item_indirect_parameters[
338                immediates.late_preprocess_work_item_indirect_offset].dispatch_x, 1u);
339        }
340
341        // Enqueue a work item for the late prepass phase.
342        late_preprocess_work_items[output_work_item_index].input_index = input_index;
343        late_preprocess_work_items[output_work_item_index].output_or_indirect_parameters_index =
344            indirect_parameters_index;
345}
346        // This mesh is culled. Skip it.
347        return;
348    }
349}
350
351    // Calculate inverse transpose.
352    let local_from_world_transpose = transpose(maths::inverse_affine3(transpose(
353        world_from_local_affine_transpose)));
354
355    // Pack inverse transpose.
356    let local_from_world_transpose_a = mat2x4<f32>(
357        vec4<f32>(local_from_world_transpose[0].xyz, local_from_world_transpose[1].x),
358        vec4<f32>(local_from_world_transpose[1].yz, local_from_world_transpose[2].xy));
359    let local_from_world_transpose_b = local_from_world_transpose[2].z;
360
361    // Figure out the output index. In indirect mode, this involves bumping the
362    // instance index in the indirect parameters metadata, which
363    // `build_indirect_params.wesl` will use to generate the actual indirect
364    // parameters. Otherwise, this index was directly supplied to us.
365@if(INDIRECT && LATE_PHASE)
366    let batch_output_index = atomicLoad(
367        &indirect_parameters_metadata[indirect_parameters_index].early_instance_count
368    ) + atomicAdd(
369        &indirect_parameters_metadata[indirect_parameters_index].late_instance_count,
370        1u
371    );
372@if(INDIRECT && !LATE_PHASE)
373    let batch_output_index = atomicAdd(
374        &indirect_parameters_metadata[indirect_parameters_index].early_instance_count,
375        1u
376    );
377
378@if(INDIRECT)
379    let mesh_output_index =
380        indirect_parameters_metadata[indirect_parameters_index].base_output_index +
381        batch_output_index;
382
383    // Write the output.
384    output[mesh_output_index].world_from_local = world_from_local_affine_transpose;
385    output[mesh_output_index].previous_world_from_local =
386        previous_world_from_local_affine_transpose;
387    output[mesh_output_index].local_from_world_transpose_a = local_from_world_transpose_a;
388    output[mesh_output_index].local_from_world_transpose_b = local_from_world_transpose_b;
389    output[mesh_output_index].flags = current_input[input_index].flags;
390    output[mesh_output_index].lightmap_uv_rect = current_input[input_index].lightmap_uv_rect;
391    output[mesh_output_index].first_vertex_index = current_input[input_index].first_vertex_index;
392    output[mesh_output_index].current_skin_index = current_input[input_index].current_skin_index;
393    output[mesh_output_index].material_and_lightmap_bind_group_slot =
394        current_input[input_index].material_and_lightmap_bind_group_slot;
395    output[mesh_output_index].tag = current_input[input_index].tag;
396    output[mesh_output_index].morph_descriptor_index = current_input[input_index].morph_descriptor_index;
397    output[mesh_output_index].metadata_index = current_input[input_index].metadata_index;
398}