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