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