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