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 Meshlet {
9    start_vertex_position_bit: u32,
10    start_vertex_attribute_id: u32,
11    start_index_id: u32,
12    packed_a: u32,
13    packed_b: u32,
14    min_vertex_position_channel_x: f32,
15    min_vertex_position_channel_y: f32,
16    min_vertex_position_channel_z: f32,
17}
18
19fn get_meshlet_vertex_count(meshlet: ptr<function, Meshlet>) -> u32 {
20    return extractBits((*meshlet).packed_a, 0u, 8u);
21}
22
23fn get_meshlet_triangle_count(meshlet: ptr<function, Meshlet>) -> u32 {
24    return extractBits((*meshlet).packed_a, 8u, 8u);
25}
26
27struct MeshletBoundingSpheres {
28    culling_sphere: MeshletBoundingSphere,
29    lod_group_sphere: MeshletBoundingSphere,
30    lod_parent_group_sphere: MeshletBoundingSphere,
31}
32
33struct MeshletBoundingSphere {
34    center: vec3<f32>,
35    radius: f32,
36}
37
38struct DispatchIndirectArgs {
39    x: atomic<u32>,
40    y: u32,
41    z: u32,
42}
43
44struct DrawIndirectArgs {
45    vertex_count: u32,
46    instance_count: atomic<u32>,
47    first_vertex: u32,
48    first_instance: u32,
49}
50
51const CENTIMETERS_PER_METER = 100.0;
52
53#ifdef MESHLET_FILL_CLUSTER_BUFFERS_PASS
54var<push_constant> scene_instance_count: u32;
55@group(0) @binding(0) var<storage, read> meshlet_instance_meshlet_counts: array<u32>; // Per entity instance
56@group(0) @binding(1) var<storage, read> meshlet_instance_meshlet_slice_starts: array<u32>; // Per entity instance
57@group(0) @binding(2) var<storage, read_write> meshlet_cluster_instance_ids: array<u32>; // Per cluster
58@group(0) @binding(3) var<storage, read_write> meshlet_cluster_meshlet_ids: array<u32>; // Per cluster
59@group(0) @binding(4) var<storage, read_write> meshlet_global_cluster_count: atomic<u32>; // Single object shared between all workgroups
60#endif
61
62#ifdef MESHLET_CULLING_PASS
63struct Constants { scene_cluster_count: u32, meshlet_raster_cluster_rightmost_slot: u32 }
64var<push_constant> constants: Constants;
65@group(0) @binding(0) var<storage, read> meshlet_cluster_meshlet_ids: array<u32>; // Per cluster
66@group(0) @binding(1) var<storage, read> meshlet_bounding_spheres: array<MeshletBoundingSpheres>; // Per meshlet
67@group(0) @binding(2) var<storage, read> meshlet_simplification_errors: array<u32>; // Per meshlet
68@group(0) @binding(3) var<storage, read> meshlet_cluster_instance_ids: array<u32>; // Per cluster
69@group(0) @binding(4) var<storage, read> meshlet_instance_uniforms: array<Mesh>; // Per entity instance
70@group(0) @binding(5) var<storage, read> meshlet_view_instance_visibility: array<u32>; // 1 bit per entity instance, packed as a bitmask
71@group(0) @binding(6) var<storage, read_write> meshlet_second_pass_candidates: array<atomic<u32>>; // 1 bit per cluster , packed as a bitmask
72@group(0) @binding(7) var<storage, read_write> meshlet_software_raster_indirect_args: DispatchIndirectArgs; // Single object shared between all workgroups
73@group(0) @binding(8) var<storage, read_write> meshlet_hardware_raster_indirect_args: DrawIndirectArgs; // Single object shared between all workgroups
74@group(0) @binding(9) var<storage, read_write> meshlet_raster_clusters: array<u32>; // Single object shared between all workgroups
75@group(0) @binding(10) 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
76@group(0) @binding(11) var<uniform> view: View;
77@group(0) @binding(12) var<uniform> previous_view: PreviousViewUniforms;
78
79fn should_cull_instance(instance_id: u32) -> bool {
80    let bit_offset = instance_id % 32u;
81    let packed_visibility = meshlet_view_instance_visibility[instance_id / 32u];
82    return bool(extractBits(packed_visibility, bit_offset, 1u));
83}
84
85// TODO: Load 4x per workgroup instead of once per thread?
86fn cluster_is_second_pass_candidate(cluster_id: u32) -> bool {
87    let packed_candidates = meshlet_second_pass_candidates[cluster_id / 32u];
88    let bit_offset = cluster_id % 32u;
89    return bool(extractBits(packed_candidates, bit_offset, 1u));
90}
91#endif
92
93#ifdef MESHLET_VISIBILITY_BUFFER_RASTER_PASS
94@group(0) @binding(0) var<storage, read> meshlet_cluster_meshlet_ids: array<u32>; // Per cluster
95@group(0) @binding(1) var<storage, read> meshlets: array<Meshlet>; // Per meshlet
96@group(0) @binding(2) var<storage, read> meshlet_indices: array<u32>; // Many per meshlet
97@group(0) @binding(3) var<storage, read> meshlet_vertex_positions: array<u32>; // Many per meshlet
98@group(0) @binding(4) var<storage, read> meshlet_cluster_instance_ids: array<u32>; // Per cluster
99@group(0) @binding(5) var<storage, read> meshlet_instance_uniforms: array<Mesh>; // Per entity instance
100@group(0) @binding(6) var<storage, read> meshlet_raster_clusters: array<u32>; // Single object shared between all workgroups
101@group(0) @binding(7) var<storage, read> meshlet_software_raster_cluster_count: u32;
102#ifdef MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT
103@group(0) @binding(8) var<storage, read_write> meshlet_visibility_buffer: array<atomic<u64>>; // Per pixel
104#else
105@group(0) @binding(8) var<storage, read_write> meshlet_visibility_buffer: array<atomic<u32>>; // Per pixel
106#endif
107@group(0) @binding(9) var<uniform> view: View;
108
109// TODO: Load only twice, instead of 3x in cases where you load 3 indices per thread?
110fn get_meshlet_vertex_id(index_id: u32) -> u32 {
111    let packed_index = meshlet_indices[index_id / 4u];
112    let bit_offset = (index_id % 4u) * 8u;
113    return extractBits(packed_index, bit_offset, 8u);
114}
115
116fn get_meshlet_vertex_position(meshlet: ptr<function, Meshlet>, vertex_id: u32) -> vec3<f32> {
117    // Get bitstream start for the vertex
118    let unpacked = unpack4xU8((*meshlet).packed_b);
119    let bits_per_channel = unpacked.xyz;
120    let bits_per_vertex = bits_per_channel.x + bits_per_channel.y + bits_per_channel.z;
121    var start_bit = (*meshlet).start_vertex_position_bit + (vertex_id * bits_per_vertex);
122
123    // Read each vertex channel from the bitstream
124    var vertex_position_packed = vec3(0u);
125    for (var i = 0u; i < 3u; i++) {
126        let lower_word_index = start_bit / 32u;
127        let lower_word_bit_offset = start_bit & 31u;
128        var next_32_bits = meshlet_vertex_positions[lower_word_index] >> lower_word_bit_offset;
129        if lower_word_bit_offset + bits_per_channel[i] > 32u {
130            next_32_bits |= meshlet_vertex_positions[lower_word_index + 1u] << (32u - lower_word_bit_offset);
131        }
132        vertex_position_packed[i] = extractBits(next_32_bits, 0u, bits_per_channel[i]);
133        start_bit += bits_per_channel[i];
134    }
135
136    // Remap [0, range_max - range_min] vec3<u32> to [range_min, range_max] vec3<f32>
137    var vertex_position = vec3<f32>(vertex_position_packed) + vec3(
138        (*meshlet).min_vertex_position_channel_x,
139        (*meshlet).min_vertex_position_channel_y,
140        (*meshlet).min_vertex_position_channel_z,
141    );
142
143    // Reverse vertex quantization
144    let vertex_position_quantization_factor = unpacked.w;
145    vertex_position /= f32(1u << vertex_position_quantization_factor) * CENTIMETERS_PER_METER;
146
147    return vertex_position;
148}
149#endif
150
151#ifdef MESHLET_MESH_MATERIAL_PASS
152@group(1) @binding(0) var<storage, read> meshlet_visibility_buffer: array<u64>; // Per pixel
153@group(1) @binding(1) var<storage, read> meshlet_cluster_meshlet_ids: array<u32>; // Per cluster
154@group(1) @binding(2) var<storage, read> meshlets: array<Meshlet>; // Per meshlet
155@group(1) @binding(3) var<storage, read> meshlet_indices: array<u32>; // Many per meshlet
156@group(1) @binding(4) var<storage, read> meshlet_vertex_positions: array<u32>; // Many per meshlet
157@group(1) @binding(5) var<storage, read> meshlet_vertex_normals: array<u32>; // Many per meshlet
158@group(1) @binding(6) var<storage, read> meshlet_vertex_uvs: array<vec2<f32>>; // Many per meshlet
159@group(1) @binding(7) var<storage, read> meshlet_cluster_instance_ids: array<u32>; // Per cluster
160@group(1) @binding(8) var<storage, read> meshlet_instance_uniforms: array<Mesh>; // Per entity instance
161
162// TODO: Load only twice, instead of 3x in cases where you load 3 indices per thread?
163fn get_meshlet_vertex_id(index_id: u32) -> u32 {
164    let packed_index = meshlet_indices[index_id / 4u];
165    let bit_offset = (index_id % 4u) * 8u;
166    return extractBits(packed_index, bit_offset, 8u);
167}
168
169fn get_meshlet_vertex_position(meshlet: ptr<function, Meshlet>, vertex_id: u32) -> vec3<f32> {
170    // Get bitstream start for the vertex
171    let unpacked = unpack4xU8((*meshlet).packed_b);
172    let bits_per_channel = unpacked.xyz;
173    let bits_per_vertex = bits_per_channel.x + bits_per_channel.y + bits_per_channel.z;
174    var start_bit = (*meshlet).start_vertex_position_bit + (vertex_id * bits_per_vertex);
175
176    // Read each vertex channel from the bitstream
177    var vertex_position_packed = vec3(0u);
178    for (var i = 0u; i < 3u; i++) {
179        let lower_word_index = start_bit / 32u;
180        let lower_word_bit_offset = start_bit & 31u;
181        var next_32_bits = meshlet_vertex_positions[lower_word_index] >> lower_word_bit_offset;
182        if lower_word_bit_offset + bits_per_channel[i] > 32u {
183            next_32_bits |= meshlet_vertex_positions[lower_word_index + 1u] << (32u - lower_word_bit_offset);
184        }
185        vertex_position_packed[i] = extractBits(next_32_bits, 0u, bits_per_channel[i]);
186        start_bit += bits_per_channel[i];
187    }
188
189    // Remap [0, range_max - range_min] vec3<u32> to [range_min, range_max] vec3<f32>
190    var vertex_position = vec3<f32>(vertex_position_packed) + vec3(
191        (*meshlet).min_vertex_position_channel_x,
192        (*meshlet).min_vertex_position_channel_y,
193        (*meshlet).min_vertex_position_channel_z,
194    );
195
196    // Reverse vertex quantization
197    let vertex_position_quantization_factor = unpacked.w;
198    vertex_position /= f32(1u << vertex_position_quantization_factor) * CENTIMETERS_PER_METER;
199
200    return vertex_position;
201}
202
203fn get_meshlet_vertex_normal(meshlet: ptr<function, Meshlet>, vertex_id: u32) -> vec3<f32> {
204    let packed_normal = meshlet_vertex_normals[(*meshlet).start_vertex_attribute_id + vertex_id];
205    return octahedral_decode_signed(unpack2x16snorm(packed_normal));
206}
207
208fn get_meshlet_vertex_uv(meshlet: ptr<function, Meshlet>, vertex_id: u32) -> vec2<f32> {
209    return meshlet_vertex_uvs[(*meshlet).start_vertex_attribute_id + vertex_id];
210}
211#endif