1// GPU mesh uniform building.
2//
3// This is a compute shader that expands each `MeshInputUniform` out to a full
4// `MeshUniform` for each view before rendering. (Thus `MeshInputUniform`
5// and `MeshUniform` are in a 1:N relationship.) It runs in parallel for all
6// meshes for all views. As part of this process, the shader gathers each
7// mesh's transform on the previous frame and writes it into the `MeshUniform`
8// so that TAA works.
9
10#import bevy_pbr::mesh_types::Mesh
11#import bevy_render::maths
12#import bevy_render::view::View
13
14// Per-frame data that the CPU supplies to the GPU.
15struct MeshInput {
16    // The model transform.
17    world_from_local: mat3x4<f32>,
18    // The lightmap UV rect, packed into 64 bits.
19    lightmap_uv_rect: vec2<u32>,
20    // Various flags.
21    flags: u32,
22    // The index of this mesh's `MeshInput` in the `previous_input` array, if
23    // applicable. If not present, this is `u32::MAX`.
24    previous_input_index: u32,
25    first_vertex_index: u32,
26    pad_a: u32,
27    pad_b: u32,
28    pad_c: u32,
29}
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_index: u32,
52}
53
54// The `wgpu` indirect parameters structure. This is a union of two structures.
55// For more information, see the corresponding comment in
56// `gpu_preprocessing.rs`.
57struct IndirectParameters {
58    // `vertex_count` or `index_count`.
59    data0: u32,
60    // `instance_count` in both structures.
61    instance_count: atomic<u32>,
62    // `first_vertex` in both structures.
63    first_vertex: u32,
64    // `first_instance` or `base_vertex`.
65    data1: u32,
66    // A read-only copy of `instance_index`.
67    instance_index: u32,
68}
69
70// The current frame's `MeshInput`.
71@group(0) @binding(0) var<storage> current_input: array<MeshInput>;
72// The `MeshInput` values from the previous frame.
73@group(0) @binding(1) var<storage> previous_input: array<MeshInput>;
74// Indices into the `MeshInput` buffer.
75//
76// There may be many indices that map to the same `MeshInput`.
77@group(0) @binding(2) var<storage> work_items: array<PreprocessWorkItem>;
78// The output array of `Mesh`es.
79@group(0) @binding(3) var<storage, read_write> output: array<Mesh>;
80
81#ifdef INDIRECT
82// The array of indirect parameters for drawcalls.
83@group(0) @binding(4) var<storage, read_write> indirect_parameters: array<IndirectParameters>;
84#endif
85
86#ifdef FRUSTUM_CULLING
87// Data needed to cull the meshes.
88//
89// At the moment, this consists only of AABBs.
90@group(0) @binding(5) var<storage> mesh_culling_data: array<MeshCullingData>;
91
92// The view data, including the view matrix.
93@group(0) @binding(6) var<uniform> view: View;
94
95// Returns true if the view frustum intersects an oriented bounding box (OBB).
96//
97// `aabb_center.w` should be 1.0.
98fn view_frustum_intersects_obb(
99    world_from_local: mat4x4<f32>,
100    aabb_center: vec4<f32>,
101    aabb_half_extents: vec3<f32>,
102) -> bool {
103
104    for (var i = 0; i < 5; i += 1) {
105        // Calculate relative radius of the sphere associated with this plane.
106        let plane_normal = view.frustum[i];
107        let relative_radius = dot(
108            abs(
109                vec3(
110                    dot(plane_normal, world_from_local[0]),
111                    dot(plane_normal, world_from_local[1]),
112                    dot(plane_normal, world_from_local[2]),
113                )
114            ),
115            aabb_half_extents
116        );
117
118        // Check the frustum plane.
119        if (!maths::sphere_intersects_plane_half_space(
120                plane_normal, aabb_center, relative_radius)) {
121            return false;
122        }
123    }
124
125    return true;
126}
127#endif
128
129@compute
130@workgroup_size(64)
131fn main(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
132    // Figure out our instance index. If this thread doesn't correspond to any
133    // index, bail.
134    let instance_index = global_invocation_id.x;
135    if (instance_index >= arrayLength(&work_items)) {
136        return;
137    }
138
139    // Unpack.
140    let input_index = work_items[instance_index].input_index;
141    let output_index = work_items[instance_index].output_index;
142    let world_from_local_affine_transpose = current_input[input_index].world_from_local;
143    let world_from_local = maths::affine3_to_square(world_from_local_affine_transpose);
144
145    // Cull if necessary.
146#ifdef FRUSTUM_CULLING
147    let aabb_center = mesh_culling_data[input_index].aabb_center.xyz;
148    let aabb_half_extents = mesh_culling_data[input_index].aabb_half_extents.xyz;
149
150    // Do an OBB-based frustum cull.
151    let model_center = world_from_local * vec4(aabb_center, 1.0);
152    if (!view_frustum_intersects_obb(world_from_local, model_center, aabb_half_extents)) {
153        return;
154    }
155#endif
156
157    // Calculate inverse transpose.
158    let local_from_world_transpose = transpose(maths::inverse_affine3(transpose(
159        world_from_local_affine_transpose)));
160
161    // Pack inverse transpose.
162    let local_from_world_transpose_a = mat2x4<f32>(
163        vec4<f32>(local_from_world_transpose[0].xyz, local_from_world_transpose[1].x),
164        vec4<f32>(local_from_world_transpose[1].yz, local_from_world_transpose[2].xy));
165    let local_from_world_transpose_b = local_from_world_transpose[2].z;
166
167    // Look up the previous model matrix.
168    let previous_input_index = current_input[input_index].previous_input_index;
169    var previous_world_from_local: mat3x4<f32>;
170    if (previous_input_index == 0xffffffff) {
171        previous_world_from_local = world_from_local_affine_transpose;
172    } else {
173        previous_world_from_local = previous_input[previous_input_index].world_from_local;
174    }
175
176    // Figure out the output index. In indirect mode, this involves bumping the
177    // instance index in the indirect parameters structure. Otherwise, this
178    // index was directly supplied to us.
179#ifdef INDIRECT
180    let mesh_output_index = indirect_parameters[output_index].instance_index +
181        atomicAdd(&indirect_parameters[output_index].instance_count, 1u);
182#else
183    let mesh_output_index = output_index;
184#endif
185
186    // Write the output.
187    output[mesh_output_index].world_from_local = world_from_local_affine_transpose;
188    output[mesh_output_index].previous_world_from_local = previous_world_from_local;
189    output[mesh_output_index].local_from_world_transpose_a = local_from_world_transpose_a;
190    output[mesh_output_index].local_from_world_transpose_b = local_from_world_transpose_b;
191    output[mesh_output_index].flags = current_input[input_index].flags;
192    output[mesh_output_index].lightmap_uv_rect = current_input[input_index].lightmap_uv_rect;
193    output[mesh_output_index].first_vertex_index = current_input[input_index].first_vertex_index;
194}