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    // https://github.com/gfx-rs/wgpu/issues/4394
186    if (instance_index == 0u) {
187        indirect_parameters_gpu_metadata[indirect_parameters_index].mesh_index = input_index;
188    } else if (work_items[instance_index - 1].output_or_indirect_parameters_index != indirect_parameters_index) {
189        indirect_parameters_gpu_metadata[indirect_parameters_index].mesh_index = input_index;
190    }
191#endif  // LATE_PHASE
192
193#else   // INDIRECT
194    let mesh_output_index = work_items[instance_index].output_or_indirect_parameters_index;
195#endif  // INDIRECT
196
197    // Unpack the input matrix.
198    let world_from_local_affine_transpose = current_input[input_index].world_from_local;
199    let world_from_local = maths::affine3_to_square(world_from_local_affine_transpose);
200
201    // Frustum cull if necessary.
202#ifdef FRUSTUM_CULLING
203    if ((current_input[input_index].flags & MESH_FLAGS_NO_FRUSTUM_CULLING_BIT) == 0u) {
204        let aabb_center = mesh_culling_data[input_index].aabb_center.xyz;
205        let aabb_half_extents = mesh_culling_data[input_index].aabb_half_extents.xyz;
206
207        // Do an OBB-based frustum cull.
208        let model_center = world_from_local * vec4(aabb_center, 1.0);
209        if (!view_frustum_intersects_obb(world_from_local, model_center, aabb_half_extents)) {
210            return;
211        }
212    }
213#endif
214
215    // See whether the `MeshInputUniform` was updated on this frame. If it
216    // wasn't, then we know the transforms of this mesh must be identical to
217    // those on the previous frame, and therefore we don't need to access the
218    // `previous_input_index` (in fact, we can't; that index are only valid for
219    // one frame and will be invalid).
220    let timestamp = current_input[input_index].timestamp;
221    let mesh_changed_this_frame = timestamp == view.frame_count;
222
223    // Look up the previous model matrix, if it could have been.
224    let previous_input_index = current_input[input_index].previous_input_index;
225    var previous_world_from_local_affine_transpose: mat3x4<f32>;
226    if (mesh_changed_this_frame && previous_input_index != 0xffffffffu) {
227        previous_world_from_local_affine_transpose =
228            previous_input[previous_input_index].world_from_local;
229    } else {
230        previous_world_from_local_affine_transpose = world_from_local_affine_transpose;
231    }
232    let previous_world_from_local =
233        maths::affine3_to_square(previous_world_from_local_affine_transpose);
234
235    // Occlusion cull if necessary. This is done by calculating the screen-space
236    // axis-aligned bounding box (AABB) of the mesh and testing it against the
237    // appropriate level of the depth pyramid (a.k.a. hierarchical Z-buffer). If
238    // no part of the AABB is in front of the corresponding pixel quad in the
239    // hierarchical Z-buffer, then this mesh must be occluded, and we can skip
240    // rendering it.
241#ifdef OCCLUSION_CULLING
242    let aabb_center = mesh_culling_data[input_index].aabb_center.xyz;
243    let aabb_half_extents = mesh_culling_data[input_index].aabb_half_extents.xyz;
244
245    // Initialize the AABB and the maximum depth.
246    let infinity = bitcast<f32>(0x7f800000u);
247    let neg_infinity = bitcast<f32>(0xff800000u);
248    var aabb = vec4(infinity, infinity, neg_infinity, neg_infinity);
249    var max_depth_view = neg_infinity;
250
251    // Build up the AABB by taking each corner of this mesh's OBB, transforming
252    // it, and updating the AABB and depth accordingly.
253    for (var i = 0u; i < 8u; i += 1u) {
254        let local_pos = aabb_center + select(
255            vec3(-1.0),
256            vec3(1.0),
257            vec3((i & 1) != 0, (i & 2) != 0, (i & 4) != 0)
258        ) * aabb_half_extents;
259
260#ifdef EARLY_PHASE
261        // If we're in the early phase, we're testing against the last frame's
262        // depth buffer, so we need to use the previous frame's transform.
263        let prev_world_pos = (previous_world_from_local * vec4(local_pos, 1.0)).xyz;
264        let view_pos = position_world_to_prev_view(prev_world_pos);
265        let ndc_pos = position_world_to_prev_ndc(prev_world_pos);
266#else   // EARLY_PHASE
267        // Otherwise, if this is the late phase, we use the current frame's
268        // transform.
269        let world_pos = (world_from_local * vec4(local_pos, 1.0)).xyz;
270        let view_pos = position_world_to_view(world_pos);
271        let ndc_pos = position_world_to_ndc(world_pos);
272#endif  // EARLY_PHASE
273
274        let uv_pos = ndc_to_uv(ndc_pos.xy);
275
276        // Update the AABB and maximum view-space depth.
277        aabb = vec4(min(aabb.xy, uv_pos), max(aabb.zw, uv_pos));
278        max_depth_view = max(max_depth_view, view_pos.z);
279    }
280
281    // Clip to the near plane to avoid the NDC depth becoming negative.
282#ifdef EARLY_PHASE
283    max_depth_view = min(-previous_view_uniforms.clip_from_view[3][2], max_depth_view);
284#else   // EARLY_PHASE
285    max_depth_view = min(-view.clip_from_view[3][2], max_depth_view);
286#endif  // EARLY_PHASE
287
288    // Figure out the depth of the occluder, and compare it to our own depth.
289
290    let aabb_pixel_size = occlusion_culling::get_aabb_size_in_pixels(aabb, depth_pyramid);
291    let occluder_depth_ndc =
292        occlusion_culling::get_occluder_depth(aabb, aabb_pixel_size, depth_pyramid);
293
294#ifdef EARLY_PHASE
295    let max_depth_ndc = prev_view_z_to_depth_ndc(max_depth_view);
296#else   // EARLY_PHASE
297    let max_depth_ndc = view_z_to_depth_ndc(max_depth_view);
298#endif
299
300    // Are we culled out?
301    if (max_depth_ndc < occluder_depth_ndc) {
302#ifdef EARLY_PHASE
303        // If this is the early phase, we need to make a note of this mesh so
304        // that we examine it again in the late phase, so that we handle the
305        // case in which a mesh that was invisible last frame became visible in
306        // this frame.
307        let output_work_item_index = atomicAdd(&late_preprocess_work_item_indirect_parameters[
308            push_constants.late_preprocess_work_item_indirect_offset].work_item_count, 1u);
309        if (output_work_item_index % 64u == 0u) {
310            // Our workgroup size is 64, and the indirect parameters for the
311            // late mesh preprocessing phase are counted in workgroups, so if
312            // we're the first thread in this workgroup, bump the workgroup
313            // count.
314            atomicAdd(&late_preprocess_work_item_indirect_parameters[
315                push_constants.late_preprocess_work_item_indirect_offset].dispatch_x, 1u);
316        }
317
318        // Enqueue a work item for the late prepass phase.
319        late_preprocess_work_items[output_work_item_index].input_index = input_index;
320        late_preprocess_work_items[output_work_item_index].output_or_indirect_parameters_index =
321            indirect_parameters_index;
322#endif  // EARLY_PHASE
323        // This mesh is culled. Skip it.
324        return;
325    }
326#endif  // OCCLUSION_CULLING
327
328    // Calculate inverse transpose.
329    let local_from_world_transpose = transpose(maths::inverse_affine3(transpose(
330        world_from_local_affine_transpose)));
331
332    // Pack inverse transpose.
333    let local_from_world_transpose_a = mat2x4<f32>(
334        vec4<f32>(local_from_world_transpose[0].xyz, local_from_world_transpose[1].x),
335        vec4<f32>(local_from_world_transpose[1].yz, local_from_world_transpose[2].xy));
336    let local_from_world_transpose_b = local_from_world_transpose[2].z;
337
338    // Figure out the output index. In indirect mode, this involves bumping the
339    // instance index in the indirect parameters metadata, which
340    // `build_indirect_params.wgsl` will use to generate the actual indirect
341    // parameters. Otherwise, this index was directly supplied to us.
342#ifdef INDIRECT
343#ifdef LATE_PHASE
344    let batch_output_index = atomicLoad(
345        &indirect_parameters_gpu_metadata[indirect_parameters_index].early_instance_count
346    ) + atomicAdd(
347        &indirect_parameters_gpu_metadata[indirect_parameters_index].late_instance_count,
348        1u
349    );
350#else   // LATE_PHASE
351    let batch_output_index = atomicAdd(
352        &indirect_parameters_gpu_metadata[indirect_parameters_index].early_instance_count,
353        1u
354    );
355#endif  // LATE_PHASE
356
357    let mesh_output_index =
358        indirect_parameters_cpu_metadata[indirect_parameters_index].base_output_index +
359        batch_output_index;
360
361#endif  // INDIRECT
362
363    // Write the output.
364    output[mesh_output_index].world_from_local = world_from_local_affine_transpose;
365    output[mesh_output_index].previous_world_from_local =
366        previous_world_from_local_affine_transpose;
367    output[mesh_output_index].local_from_world_transpose_a = local_from_world_transpose_a;
368    output[mesh_output_index].local_from_world_transpose_b = local_from_world_transpose_b;
369    output[mesh_output_index].flags = current_input[input_index].flags;
370    output[mesh_output_index].lightmap_uv_rect = current_input[input_index].lightmap_uv_rect;
371    output[mesh_output_index].first_vertex_index = current_input[input_index].first_vertex_index;
372    output[mesh_output_index].current_skin_index = current_input[input_index].current_skin_index;
373    output[mesh_output_index].material_and_lightmap_bind_group_slot =
374        current_input[input_index].material_and_lightmap_bind_group_slot;
375    output[mesh_output_index].tag = current_input[input_index].tag;
376}