1import package::render::mesh_types::Mesh;
2import bevy_render::view::View;
3import package::prepass::bindings::PreviousViewUniforms;
4import bevy_render::utils::octahedral_decode_signed;
5
6struct BvhNode {
7    aabbs: array<MeshletAabbErrorOffset, 8>,
8    lod_bounds: array<vec4<f32>, 8>,
9    child_counts: array<u32, 2>,
10    _padding: vec2<u32>,
11}
12
13struct Meshlet {
14    start_vertex_position_bit: u32,
15    start_vertex_attribute_id: u32,
16    start_index_id: u32,
17    packed_a: u32,
18    packed_b: u32,
19    min_vertex_position_channel_x: f32,
20    min_vertex_position_channel_y: f32,
21    min_vertex_position_channel_z: f32,
22}
23
24fn get_meshlet_vertex_count(meshlet: ptr<function, Meshlet>) -> u32 {
25    return extractBits((*meshlet).packed_a, 0u, 8u) + 1u;
26}
27
28fn get_meshlet_triangle_count(meshlet: ptr<function, Meshlet>) -> u32 {
29    return extractBits((*meshlet).packed_a, 8u, 8u);
30}
31
32struct MeshletCullData {
33    aabb: MeshletAabbErrorOffset,
34    lod_group_sphere: vec4<f32>,
35}
36
37struct MeshletAabb {
38    center: vec3<f32>,
39    half_extent: vec3<f32>,
40}
41
42struct MeshletAabbErrorOffset {
43    center_and_error: vec4<f32>,
44    half_extent_and_child_offset: vec4<f32>,
45}
46
47fn get_aabb(aabb: ptr<function, MeshletAabbErrorOffset>) -> MeshletAabb {
48    return MeshletAabb(
49        (*aabb).center_and_error.xyz,
50        (*aabb).half_extent_and_child_offset.xyz,
51    );
52}
53
54fn get_aabb_error(aabb: ptr<function, MeshletAabbErrorOffset>) -> f32 {
55    return (*aabb).center_and_error.w;
56}
57
58fn get_aabb_child_offset(aabb: ptr<function, MeshletAabbErrorOffset>) -> u32 {
59    return bitcast<u32>((*aabb).half_extent_and_child_offset.w);
60}
61
62struct DispatchIndirectArgs {
63    x: atomic<u32>,
64    y: u32,
65    z: u32,
66}
67
68struct DrawIndirectArgs {
69    vertex_count: u32,
70    instance_count: atomic<u32>,
71    first_vertex: u32,
72    first_instance: u32,
73}
74
75/// Either a BVH node or a meshlet, along with the instance it is associated with.
76/// Refers to BVH nodes in `meshlet_bvh_cull_queue` and `meshlet_second_pass_bvh_queue`, where `offset` is the index into `meshlet_bvh_nodes`.
77/// Refers to meshlets in `meshlet_meshlet_cull_queue` and `meshlet_raster_clusters`.
78/// In `meshlet_meshlet_cull_queue`, `offset` is the index into `meshlet_cull_data`.
79/// In `meshlet_raster_clusters`, `offset` is the index into `meshlets`.
80struct InstancedOffset {
81    instance_id: u32,
82    offset: u32,
83}
84
85const CENTIMETERS_PER_METER = 100.0;
86
87@if(MESHLET_INSTANCE_CULLING_PASS)
88struct Constants { scene_instance_count: u32 }
89@if(MESHLET_INSTANCE_CULLING_PASS) {
90var<immediate> constants: Constants;
91
92// Cull data
93@group(0) @binding(0) var depth_pyramid: texture_2d<f32>;
94@group(0) @binding(1) var<uniform> view: View;
95@group(0) @binding(2) var<uniform> previous_view: PreviousViewUniforms;
96
97// Per entity instance data
98@group(0) @binding(3) var<storage, read> meshlet_instance_uniforms: array<Mesh>;
99@group(0) @binding(4) var<storage, read> meshlet_view_instance_visibility: array<u32>; // 1 bit per entity instance, packed as a bitmask
100@group(0) @binding(5) var<storage, read> meshlet_instance_aabbs: array<MeshletAabb>;
101@group(0) @binding(6) var<storage, read> meshlet_instance_bvh_root_nodes: array<u32>;
102
103// BVH cull queue data
104@group(0) @binding(7) var<storage, read_write> meshlet_bvh_cull_count_write: atomic<u32>;
105@group(0) @binding(8) var<storage, read_write> meshlet_bvh_cull_dispatch: DispatchIndirectArgs;
106@group(0) @binding(9) var<storage, read_write> meshlet_bvh_cull_queue: array<InstancedOffset>;
107}
108
109// Second pass queue data
110@if(MESHLET_INSTANCE_CULLING_PASS && MESHLET_FIRST_CULLING_PASS) {
111@group(0) @binding(10) var<storage, read_write> meshlet_second_pass_instance_count: atomic<u32>;
112@group(0) @binding(11) var<storage, read_write> meshlet_second_pass_instance_dispatch: DispatchIndirectArgs;
113@group(0) @binding(12) var<storage, read_write> meshlet_second_pass_instance_candidates: array<u32>;
114}
115@if(MESHLET_INSTANCE_CULLING_PASS && !MESHLET_FIRST_CULLING_PASS) {
116@group(0) @binding(10) var<storage, read> meshlet_second_pass_instance_count: u32;
117@group(0) @binding(11) var<storage, read> meshlet_second_pass_instance_candidates: array<u32>;
118}
119
120@if(MESHLET_BVH_CULLING_PASS)
121struct Constants { read_from_front: u32, rightmost_slot: u32 }
122@if(MESHLET_BVH_CULLING_PASS) {
123var<immediate> constants: Constants;
124
125// Cull data
126@group(0) @binding(0) var depth_pyramid: texture_2d<f32>; // From the end of the last frame for the first culling pass, and from the first raster pass for the second culling pass
127@group(0) @binding(1) var<uniform> view: View;
128@group(0) @binding(2) var<uniform> previous_view: PreviousViewUniforms;
129
130// Global mesh data
131@group(0) @binding(3) var<storage, read> meshlet_bvh_nodes: array<BvhNode>;
132
133// Per entity instance data
134@group(0) @binding(4) var<storage, read> meshlet_instance_uniforms: array<Mesh>;
135
136// BVH cull queue data
137@group(0) @binding(5) var<storage, read> meshlet_bvh_cull_count_read: u32;
138@group(0) @binding(6) var<storage, read_write> meshlet_bvh_cull_count_write: atomic<u32>;
139@group(0) @binding(7) var<storage, read_write> meshlet_bvh_cull_dispatch: DispatchIndirectArgs;
140@group(0) @binding(8) var<storage, read_write> meshlet_bvh_cull_queue: array<InstancedOffset>;
141
142// Meshlet cull queue data
143@group(0) @binding(9) var<storage, read_write> meshlet_meshlet_cull_count_early: atomic<u32>;
144@group(0) @binding(10) var<storage, read_write> meshlet_meshlet_cull_count_late: atomic<u32>;
145@group(0) @binding(11) var<storage, read_write> meshlet_meshlet_cull_dispatch_early: DispatchIndirectArgs;
146@group(0) @binding(12) var<storage, read_write> meshlet_meshlet_cull_dispatch_late: DispatchIndirectArgs;
147@group(0) @binding(13) var<storage, read_write> meshlet_meshlet_cull_queue: array<InstancedOffset>;
148}
149
150// Second pass queue data
151@if(MESHLET_BVH_CULLING_PASS && MESHLET_FIRST_CULLING_PASS) {
152@group(0) @binding(14) var<storage, read_write> meshlet_second_pass_bvh_count: atomic<u32>;
153@group(0) @binding(15) var<storage, read_write> meshlet_second_pass_bvh_dispatch: DispatchIndirectArgs;
154@group(0) @binding(16) var<storage, read_write> meshlet_second_pass_bvh_queue: array<InstancedOffset>;
155}
156
157@if(MESHLET_CLUSTER_CULLING_PASS)
158struct Constants { rightmost_slot: u32 }
159@if(MESHLET_CLUSTER_CULLING_PASS) {
160var<immediate> constants: Constants;
161
162// Cull data
163@group(0) @binding(0) var depth_pyramid: texture_2d<f32>; // From the end of the last frame for the first culling pass, and from the first raster pass for the second culling pass
164@group(0) @binding(1) var<uniform> view: View;
165@group(0) @binding(2) var<uniform> previous_view: PreviousViewUniforms;
166
167// Global mesh data
168@group(0) @binding(3) var<storage, read> meshlet_cull_data: array<MeshletCullData>;
169
170// Per entity instance data
171@group(0) @binding(4) var<storage, read> meshlet_instance_uniforms: array<Mesh>;
172
173// Raster queue data
174@group(0) @binding(5) var<storage, read_write> meshlet_software_raster_indirect_args: DispatchIndirectArgs;
175@group(0) @binding(6) var<storage, read_write> meshlet_hardware_raster_indirect_args: DrawIndirectArgs;
176@group(0) @binding(7) var<storage, read> meshlet_previous_raster_counts: array<u32>;
177@group(0) @binding(8) var<storage, read_write> meshlet_raster_clusters: array<InstancedOffset>;
178
179// Meshlet cull queue data
180@group(0) @binding(9) var<storage, read> meshlet_meshlet_cull_count_read: u32;
181}
182
183// Second pass queue data
184@if(MESHLET_CLUSTER_CULLING_PASS && MESHLET_FIRST_CULLING_PASS) {
185@group(0) @binding(10) var<storage, read_write> meshlet_meshlet_cull_count_write: atomic<u32>;
186@group(0) @binding(11) var<storage, read_write> meshlet_meshlet_cull_dispatch: DispatchIndirectArgs;
187@group(0) @binding(12) var<storage, read_write> meshlet_meshlet_cull_queue: array<InstancedOffset>;
188}
189@if(MESHLET_CLUSTER_CULLING_PASS && !MESHLET_FIRST_CULLING_PASS)
190@group(0) @binding(10) var<storage, read> meshlet_meshlet_cull_queue: array<InstancedOffset>;
191
192@if(MESHLET_VISIBILITY_BUFFER_RASTER_PASS) {
193@group(0) @binding(0) var<storage, read> meshlet_raster_clusters: array<InstancedOffset>; // Per cluster
194@group(0) @binding(1) var<storage, read> meshlets: array<Meshlet>; // Per meshlet
195@group(0) @binding(2) var<storage, read> meshlet_indices: array<u32>; // Many per meshlet
196@group(0) @binding(3) var<storage, read> meshlet_vertex_positions: array<u32>; // Many per meshlet
197@group(0) @binding(4) var<storage, read> meshlet_instance_uniforms: array<Mesh>; // Per entity instance
198@group(0) @binding(5) var<storage, read> meshlet_previous_raster_counts: array<u32>;
199@group(0) @binding(6) var<storage, read> meshlet_software_raster_cluster_count: u32;
200}
201@if(MESHLET_VISIBILITY_BUFFER_RASTER_PASS && MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT)
202@group(0) @binding(7) var meshlet_visibility_buffer: texture_storage_2d<r64uint, atomic>;
203@if(MESHLET_VISIBILITY_BUFFER_RASTER_PASS && !MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT)
204@group(0) @binding(7) var meshlet_visibility_buffer: texture_storage_2d<r32uint, atomic>;
205@if(MESHLET_VISIBILITY_BUFFER_RASTER_PASS)
206@group(0) @binding(8) var<uniform> view: View;
207
208// TODO: Load only twice, instead of 3x in cases where you load 3 indices per thread?
209@if(MESHLET_VISIBILITY_BUFFER_RASTER_PASS)
210fn get_meshlet_vertex_id(index_id: u32) -> u32 {
211    let packed_index = meshlet_indices[index_id / 4u];
212    let bit_offset = (index_id % 4u) * 8u;
213    return extractBits(packed_index, bit_offset, 8u);
214}
215
216@if(MESHLET_VISIBILITY_BUFFER_RASTER_PASS)
217fn get_meshlet_vertex_position(meshlet: ptr<function, Meshlet>, vertex_id: u32) -> vec3<f32> {
218    // Get bitstream start for the vertex
219    let unpacked = unpack4xU8((*meshlet).packed_b);
220    let bits_per_channel = unpacked.xyz;
221    let bits_per_vertex = bits_per_channel.x + bits_per_channel.y + bits_per_channel.z;
222    var start_bit = (*meshlet).start_vertex_position_bit + (vertex_id * bits_per_vertex);
223
224    // Read each vertex channel from the bitstream
225    var vertex_position_packed = vec3(0u);
226    for (var i = 0u; i < 3u; i++) {
227        let lower_word_index = start_bit / 32u;
228        let lower_word_bit_offset = start_bit & 31u;
229        var next_32_bits = meshlet_vertex_positions[lower_word_index] >> lower_word_bit_offset;
230        if lower_word_bit_offset + bits_per_channel[i] > 32u {
231            next_32_bits |= meshlet_vertex_positions[lower_word_index + 1u] << (32u - lower_word_bit_offset);
232        }
233        vertex_position_packed[i] = extractBits(next_32_bits, 0u, bits_per_channel[i]);
234        start_bit += bits_per_channel[i];
235    }
236
237    // Remap [0, range_max - range_min] vec3<u32> to [range_min, range_max] vec3<f32>
238    var vertex_position = vec3<f32>(vertex_position_packed) + vec3(
239        (*meshlet).min_vertex_position_channel_x,
240        (*meshlet).min_vertex_position_channel_y,
241        (*meshlet).min_vertex_position_channel_z,
242    );
243
244    // Reverse vertex quantization
245    let vertex_position_quantization_factor = unpacked.w;
246    vertex_position /= f32(1u << vertex_position_quantization_factor) * CENTIMETERS_PER_METER;
247
248    return vertex_position;
249}
250
251@if(MESHLET_MESH_MATERIAL_PASS) {
252@group(2) @binding(0) var meshlet_visibility_buffer: texture_storage_2d<r64uint, read>;
253@group(2) @binding(1) var<storage, read> meshlet_raster_clusters: array<InstancedOffset>; // Per cluster
254@group(2) @binding(2) var<storage, read> meshlets: array<Meshlet>; // Per meshlet
255@group(2) @binding(3) var<storage, read> meshlet_indices: array<u32>; // Many per meshlet
256@group(2) @binding(4) var<storage, read> meshlet_vertex_positions: array<u32>; // Many per meshlet
257@group(2) @binding(5) var<storage, read> meshlet_vertex_normals: array<u32>; // Many per meshlet
258@group(2) @binding(6) var<storage, read> meshlet_vertex_uvs: array<vec2<f32>>; // Many per meshlet
259@group(2) @binding(7) var<storage, read> meshlet_instance_uniforms: array<Mesh>; // Per entity instance
260}
261
262// TODO: Load only twice, instead of 3x in cases where you load 3 indices per thread?
263@if(MESHLET_MESH_MATERIAL_PASS)
264fn get_meshlet_vertex_id(index_id: u32) -> u32 {
265    let packed_index = meshlet_indices[index_id / 4u];
266    let bit_offset = (index_id % 4u) * 8u;
267    return extractBits(packed_index, bit_offset, 8u);
268}
269
270@if(MESHLET_MESH_MATERIAL_PASS)
271fn get_meshlet_vertex_position(meshlet: ptr<function, Meshlet>, vertex_id: u32) -> vec3<f32> {
272    // Get bitstream start for the vertex
273    let unpacked = unpack4xU8((*meshlet).packed_b);
274    let bits_per_channel = unpacked.xyz;
275    let bits_per_vertex = bits_per_channel.x + bits_per_channel.y + bits_per_channel.z;
276    var start_bit = (*meshlet).start_vertex_position_bit + (vertex_id * bits_per_vertex);
277
278    // Read each vertex channel from the bitstream
279    var vertex_position_packed = vec3(0u);
280    for (var i = 0u; i < 3u; i++) {
281        let lower_word_index = start_bit / 32u;
282        let lower_word_bit_offset = start_bit & 31u;
283        var next_32_bits = meshlet_vertex_positions[lower_word_index] >> lower_word_bit_offset;
284        if lower_word_bit_offset + bits_per_channel[i] > 32u {
285            next_32_bits |= meshlet_vertex_positions[lower_word_index + 1u] << (32u - lower_word_bit_offset);
286        }
287        vertex_position_packed[i] = extractBits(next_32_bits, 0u, bits_per_channel[i]);
288        start_bit += bits_per_channel[i];
289    }
290
291    // Remap [0, range_max - range_min] vec3<u32> to [range_min, range_max] vec3<f32>
292    var vertex_position = vec3<f32>(vertex_position_packed) + vec3(
293        (*meshlet).min_vertex_position_channel_x,
294        (*meshlet).min_vertex_position_channel_y,
295        (*meshlet).min_vertex_position_channel_z,
296    );
297
298    // Reverse vertex quantization
299    let vertex_position_quantization_factor = unpacked.w;
300    vertex_position /= f32(1u << vertex_position_quantization_factor) * CENTIMETERS_PER_METER;
301
302    return vertex_position;
303}
304
305@if(MESHLET_MESH_MATERIAL_PASS)
306fn get_meshlet_vertex_normal(meshlet: ptr<function, Meshlet>, vertex_id: u32) -> vec3<f32> {
307    let packed_normal = meshlet_vertex_normals[(*meshlet).start_vertex_attribute_id + vertex_id];
308    return octahedral_decode_signed(unpack2x16snorm(packed_normal));
309}
310
311@if(MESHLET_MESH_MATERIAL_PASS)
312fn get_meshlet_vertex_uv(meshlet: ptr<function, Meshlet>, vertex_id: u32) -> vec2<f32> {
313    return meshlet_vertex_uvs[(*meshlet).start_vertex_attribute_id + vertex_id];
314}