Fix meshlet shaders for bindless mode. (#16825)

We have to extract the material ID from the mesh and stuff it in the
vertex during visibility buffer resolution.

Author: pcwalton

crates/bevy_pbr/src/meshlet/visibility_buffer_resolve.wgsl

@@ -94,6 +94,7 @@ struct VertexOutput {
     world_tangent: vec4<f32>,
     mesh_flags: u32,
     cluster_id: u32,
+    material_bind_group_slot: u32,
 #ifdef PREPASS_FRAGMENT
 #ifdef MOTION_VECTOR_PREPASS
     motion_vector: vec2<f32>,
@@ -173,6 +174,7 @@ fn resolve_vertex_output(frag_coord: vec4<f32>) -> VertexOutput {
         world_tangent,
         instance_uniform.flags,
         instance_id ^ meshlet_id,
+        instance_uniform.material_and_lightmap_bind_group_slot & 0xffffu,
 #ifdef PREPASS_FRAGMENT
 #ifdef MOTION_VECTOR_PREPASS
         motion_vector,

crates/bevy_pbr/src/render/mesh_bindings.wgsl

@@ -2,8 +2,10 @@
 
 #import bevy_pbr::mesh_types::Mesh
 
+#ifndef MESHLET_MESH_MATERIAL_PASS
 #ifdef PER_OBJECT_BUFFER_BATCH_SIZE
 @group(1) @binding(0) var<uniform> mesh: array<Mesh, #{PER_OBJECT_BUFFER_BATCH_SIZE}u>;
 #else
 @group(1) @binding(0) var<storage> mesh: array<Mesh>;
 #endif // PER_OBJECT_BUFFER_BATCH_SIZE
+#endif  // MESHLET_MESH_MATERIAL_PASS

crates/bevy_pbr/src/render/mesh_functions.wgsl

@@ -12,6 +12,7 @@
 }
 #import bevy_render::maths::{affine3_to_square, mat2x4_f32_to_mat3x3_unpack}
 
+#ifndef MESHLET_MESH_MATERIAL_PASS
 
 fn get_world_from_local(instance_index: u32) -> mat4x4<f32> {
     return affine3_to_square(mesh[instance_index].world_from_local);
@@ -21,6 +22,8 @@ fn get_previous_world_from_local(instance_index: u32) -> mat4x4<f32> {
     return affine3_to_square(mesh[instance_index].previous_world_from_local);
 }
 
+#endif  // MESHLET_MESH_MATERIAL_PASS
+
 fn mesh_position_local_to_world(world_from_local: mat4x4<f32>, vertex_position: vec4<f32>) -> vec4<f32> {
     return world_from_local * vertex_position;
 }
@@ -33,6 +36,8 @@ fn mesh_position_local_to_clip(world_from_local: mat4x4<f32>, vertex_position: v
     return position_world_to_clip(world_position.xyz);
 }
 
+#ifndef MESHLET_MESH_MATERIAL_PASS
+
 fn mesh_normal_local_to_world(vertex_normal: vec3<f32>, instance_index: u32) -> vec3<f32> {
     // NOTE: The mikktspace method of normal mapping requires that the world normal is
     // re-normalized in the vertex shader to match the way mikktspace bakes vertex tangents
@@ -53,6 +58,8 @@ fn mesh_normal_local_to_world(vertex_normal: vec3<f32>, instance_index: u32) ->
     }
 }
 
+#endif  // MESHLET_MESH_MATERIAL_PASS
+
 // Calculates the sign of the determinant of the 3x3 model matrix based on a
 // mesh flag
 fn sign_determinant_model_3x3m(mesh_flags: u32) -> f32 {
@@ -62,6 +69,8 @@ fn sign_determinant_model_3x3m(mesh_flags: u32) -> f32 {
     return f32(bool(mesh_flags & MESH_FLAGS_SIGN_DETERMINANT_MODEL_3X3_BIT)) * 2.0 - 1.0;
 }
 
+#ifndef MESHLET_MESH_MATERIAL_PASS
+
 fn mesh_tangent_local_to_world(world_from_local: mat4x4<f32>, vertex_tangent: vec4<f32>, instance_index: u32) -> vec4<f32> {
     // NOTE: The mikktspace method of normal mapping requires that the world tangent is
     // re-normalized in the vertex shader to match the way mikktspace bakes vertex tangents
@@ -88,6 +97,8 @@ fn mesh_tangent_local_to_world(world_from_local: mat4x4<f32>, vertex_tangent: ve
     }
 }
 
+#endif  // MESHLET_MESH_MATERIAL_PASS
+
 // Returns an appropriate dither level for the current mesh instance.
 //
 // This looks up the LOD range in the `visibility_ranges` table and compares the

crates/bevy_pbr/src/render/parallax_mapping.wgsl

@@ -5,8 +5,7 @@
     mesh_bindings::mesh
 }
 
-fn sample_depth_map(uv: vec2<f32>, instance_index: u32) -> f32 {
-    let slot = mesh[instance_index].material_and_lightmap_bind_group_slot & 0xffffu;
+fn sample_depth_map(uv: vec2<f32>, material_bind_group_slot: u32) -> f32 {
     // We use `textureSampleLevel` over `textureSample` because the wgpu DX12
     // backend (Fxc) panics when using "gradient instructions" inside a loop.
     // It results in the whole loop being unrolled by the shader compiler,
@@ -19,8 +18,8 @@ fn sample_depth_map(uv: vec2<f32>, instance_index: u32) -> f32 {
     // See https://stackoverflow.com/questions/56581141/direct3d11-gradient-instruction-used-in-a-loop-with-varying-iteration-forcing
     return textureSampleLevel(
 #ifdef BINDLESS
-        depth_map_texture[slot],
-        depth_map_sampler[slot],
+        depth_map_texture[material_bind_group_slot],
+        depth_map_sampler[material_bind_group_slot],
 #else   // BINDLESS
         depth_map_texture,
         depth_map_sampler,
@@ -40,7 +39,7 @@ fn parallaxed_uv(
     original_uv: vec2<f32>,
     // The vector from the camera to the fragment at the surface in tangent space
     Vt: vec3<f32>,
-    instance_index: u32,
+    material_bind_group_slot: u32,
 ) -> vec2<f32> {
     if max_layer_count < 1.0 {
         return original_uv;
@@ -68,15 +67,15 @@ fn parallaxed_uv(
     var delta_uv = depth_scale * layer_depth * Vt.xy * vec2(1.0, -1.0) / view_steepness;
 
     var current_layer_depth = 0.0;
-    var texture_depth = sample_depth_map(uv, instance_index);
+    var texture_depth = sample_depth_map(uv, material_bind_group_slot);
 
     // texture_depth > current_layer_depth means the depth map depth is deeper
     // than the depth the ray would be at this UV offset so the ray has not
     // intersected the surface
     for (var i: i32 = 0; texture_depth > current_layer_depth && i <= i32(layer_count); i++) {
         current_layer_depth += layer_depth;
         uv += delta_uv;
-        texture_depth = sample_depth_map(uv, instance_index);
+        texture_depth = sample_depth_map(uv, material_bind_group_slot);
     }
 
 #ifdef RELIEF_MAPPING
@@ -94,7 +93,7 @@ fn parallaxed_uv(
     current_layer_depth -= delta_depth;
 
     for (var i: u32 = 0u; i < max_steps; i++) {
-        texture_depth = sample_depth_map(uv, instance_index);
+        texture_depth = sample_depth_map(uv, material_bind_group_slot);
 
         // Halve the deltas for the next step
         delta_uv *= 0.5;
@@ -118,7 +117,8 @@ fn parallaxed_uv(
     // may skip small details and result in writhing material artifacts.
     let previous_uv = uv - delta_uv;
     let next_depth = texture_depth - current_layer_depth;
-    let previous_depth = sample_depth_map(previous_uv, instance_index) - current_layer_depth + layer_depth;
+    let previous_depth = sample_depth_map(previous_uv, material_bind_group_slot) -
+        current_layer_depth + layer_depth;
 
     let weight = next_depth / (next_depth - previous_depth);
 

crates/bevy_pbr/src/render/pbr_fragment.wgsl

@@ -71,7 +71,11 @@ fn pbr_input_from_standard_material(
     is_front: bool,
 ) -> pbr_types::PbrInput {
 #ifdef BINDLESS
+#ifdef MESHLET_MESH_MATERIAL_PASS
+    let slot = in.material_bind_group_slot;
+#else   // MESHLET_MESH_MATERIAL_PASS
     let slot = mesh[in.instance_index].material_and_lightmap_bind_group_slot & 0xffffu;
+#endif  // MESHLET_MESH_MATERIAL_PASS
     let flags = pbr_bindings::material[slot].flags;
     let base_color = pbr_bindings::material[slot].base_color;
     let deferred_lighting_pass_id = pbr_bindings::material[slot].deferred_lighting_pass_id;
@@ -146,7 +150,7 @@ fn pbr_input_from_standard_material(
             // parallax mapping algorithm easier to understand and reason
             // about.
             -Vt,
-            in.instance_index,
+            slot,
         );
 #endif