[Merged by Bors] - Separate out PBR lighting, shadows, clustered forward, and utils from pbr.wgsl

assets/shaders/animate_shader.wgsl

@@ -4,6 +4,9 @@
 [[group(1), binding(0)]]
 var<uniform> mesh: Mesh;
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
     [[location(1)]] normal: vec3<f32>;
@@ -17,10 +20,8 @@ struct VertexOutput {
 
 [[stage(vertex)]]
 fn vertex(vertex: Vertex) -> VertexOutput {
-    let world_position = mesh.model * vec4<f32>(vertex.position, 1.0);
-
     var out: VertexOutput;
-    out.clip_position = view.view_proj * world_position;
+    out.clip_position = mesh_position_local_to_clip(mesh.model, vec4<f32>(vertex.position, 1.0));
     out.uv = vertex.uv;
     return out;
 }

assets/shaders/custom_vertex_attribute.wgsl

@@ -1,28 +1,29 @@
 #import bevy_pbr::mesh_view_bindings
 #import bevy_pbr::mesh_bindings
 
-struct Vertex {
-    [[location(0)]] position: vec3<f32>;
-    [[location(1)]] blend_color: vec4<f32>;
-};
-
 struct CustomMaterial {
     color: vec4<f32>;
 };
 [[group(1), binding(0)]]
 var<uniform> material: CustomMaterial;
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
+struct Vertex {
+    [[location(0)]] position: vec3<f32>;
+    [[location(1)]] blend_color: vec4<f32>;
+};
+
 struct VertexOutput {
     [[builtin(position)]] clip_position: vec4<f32>;
     [[location(0)]] blend_color: vec4<f32>;
 };
 
 [[stage(vertex)]]
 fn vertex(vertex: Vertex) -> VertexOutput {
-    let world_position = mesh.model * vec4<f32>(vertex.position, 1.0);
-
     var out: VertexOutput;
-    out.clip_position = view.view_proj * world_position;
+    out.clip_position = mesh_position_local_to_clip(mesh.model, vec4<f32>(vertex.position, 1.0));
     out.blend_color = vertex.blend_color;
     return out;
 }

assets/shaders/instancing.wgsl

@@ -4,6 +4,9 @@
 [[group(1), binding(0)]]
 var<uniform> mesh: Mesh;
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
     [[location(1)]] normal: vec3<f32>;
@@ -21,10 +24,8 @@ struct VertexOutput {
 [[stage(vertex)]]
 fn vertex(vertex: Vertex) -> VertexOutput {
     let position = vertex.position * vertex.i_pos_scale.w + vertex.i_pos_scale.xyz;
-    let world_position = mesh.model * vec4<f32>(position, 1.0);
-
     var out: VertexOutput;
-    out.clip_position = view.view_proj * world_position;
+    out.clip_position = mesh_position_local_to_clip(mesh.model, vec4<f32>(position, 1.0));
     out.color = vertex.i_color;
     return out;
 }

assets/shaders/shader_defs.wgsl

@@ -4,6 +4,9 @@
 [[group(1), binding(0)]]
 var<uniform> mesh: Mesh;
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
     [[location(1)]] normal: vec3<f32>;
@@ -16,10 +19,8 @@ struct VertexOutput {
 
 [[stage(vertex)]]
 fn vertex(vertex: Vertex) -> VertexOutput {
-    let world_position = mesh.model * vec4<f32>(vertex.position, 1.0);
-
     var out: VertexOutput;
-    out.clip_position = view.view_proj * world_position;
+    out.clip_position = mesh_position_local_to_clip(mesh.model, vec4<f32>(vertex.position, 1.0));
     return out;
 }
 

crates/bevy_pbr/src/lib.rs

@@ -54,6 +54,14 @@ pub const PBR_TYPES_SHADER_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 1708015359337029744);
 pub const PBR_BINDINGS_SHADER_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 5635987986427308186);
+pub const UTILS_HANDLE: HandleUntyped =
+    HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 1900548483293416725);
+pub const CLUSTERED_FORWARD_HANDLE: HandleUntyped =
+    HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 166852093121196815);
+pub const PBR_LIGHTING_HANDLE: HandleUntyped =
+    HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 14170772752254856967);
+pub const SHADOWS_HANDLE: HandleUntyped =
+    HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 11350275143789590502);
 pub const PBR_SHADER_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 4805239651767701046);
 pub const SHADOW_SHADER_HANDLE: HandleUntyped =
@@ -77,6 +85,25 @@ impl Plugin for PbrPlugin {
             "render/pbr_bindings.wgsl",
             Shader::from_wgsl
         );
+        load_internal_asset!(app, UTILS_HANDLE, "render/utils.wgsl", Shader::from_wgsl);
+        load_internal_asset!(
+            app,
+            CLUSTERED_FORWARD_HANDLE,
+            "render/clustered_forward.wgsl",
+            Shader::from_wgsl
+        );
+        load_internal_asset!(
+            app,
+            PBR_LIGHTING_HANDLE,
+            "render/pbr_lighting.wgsl",
+            Shader::from_wgsl
+        );
+        load_internal_asset!(
+            app,
+            SHADOWS_HANDLE,
+            "render/shadows.wgsl",
+            Shader::from_wgsl
+        );
         load_internal_asset!(app, PBR_SHADER_HANDLE, "render/pbr.wgsl", Shader::from_wgsl);
         load_internal_asset!(
             app,

crates/bevy_pbr/src/render/clustered_forward.wgsl

@@ -0,0 +1,100 @@
+#define_import_path bevy_pbr::clustered_forward
+
+// NOTE: Keep in sync with bevy_pbr/src/light.rs
+fn view_z_to_z_slice(view_z: f32, is_orthographic: bool) -> u32 {
+    var z_slice: u32 = 0u;
+    if (is_orthographic) {
+        // NOTE: view_z is correct in the orthographic case
+        z_slice = u32(floor((view_z - lights.cluster_factors.z) * lights.cluster_factors.w));
+    } else {
+        // NOTE: had to use -view_z to make it positive else log(negative) is nan
+        z_slice = u32(log(-view_z) * lights.cluster_factors.z - lights.cluster_factors.w + 1.0);
+    }
+    // NOTE: We use min as we may limit the far z plane used for clustering to be closeer than
+    // the furthest thing being drawn. This means that we need to limit to the maximum cluster.
+    return min(z_slice, lights.cluster_dimensions.z - 1u);
+}
+
+fn fragment_cluster_index(frag_coord: vec2<f32>, view_z: f32, is_orthographic: bool) -> u32 {
+    let xy = vec2<u32>(floor(frag_coord * lights.cluster_factors.xy));
+    let z_slice = view_z_to_z_slice(view_z, is_orthographic);
+    // NOTE: Restricting cluster index to avoid undefined behavior when accessing uniform buffer
+    // arrays based on the cluster index.
+    return min(
+        (xy.y * lights.cluster_dimensions.x + xy.x) * lights.cluster_dimensions.z + z_slice,
+        lights.cluster_dimensions.w - 1u
+    );
+}
+
+// this must match CLUSTER_COUNT_SIZE in light.rs
+let CLUSTER_COUNT_SIZE = 13u;
+fn unpack_offset_and_count(cluster_index: u32) -> vec2<u32> {
+#ifdef NO_STORAGE_BUFFERS_SUPPORT
+    let offset_and_count = cluster_offsets_and_counts.data[cluster_index >> 2u][cluster_index & ((1u << 2u) - 1u)];
+    return vec2<u32>(
+        // The offset is stored in the upper 32 - CLUSTER_COUNT_SIZE = 19 bits
+        (offset_and_count >> CLUSTER_COUNT_SIZE) & ((1u << 32u - CLUSTER_COUNT_SIZE) - 1u),
+        // The count is stored in the lower CLUSTER_COUNT_SIZE = 13 bits
+        offset_and_count & ((1u << CLUSTER_COUNT_SIZE) - 1u)
+    );
+#else
+    return cluster_offsets_and_counts.data[cluster_index];
+#endif
+}
+
+fn get_light_id(index: u32) -> u32 {
+#ifdef NO_STORAGE_BUFFERS_SUPPORT
+    // The index is correct but in cluster_light_index_lists we pack 4 u8s into a u32
+    // This means the index into cluster_light_index_lists is index / 4
+    let indices = cluster_light_index_lists.data[index >> 4u][(index >> 2u) & ((1u << 2u) - 1u)];
+    // And index % 4 gives the sub-index of the u8 within the u32 so we shift by 8 * sub-index
+    return (indices >> (8u * (index & ((1u << 2u) - 1u)))) & ((1u << 8u) - 1u);
+#else
+    return cluster_light_index_lists.data[index];
+#endif
+}
+
+fn cluster_debug_visualization(
+    output_color: vec4<f32>,
+    view_z: f32,
+    is_orthographic: bool,
+    offset_and_count: vec2<u32>,
+    cluster_index: u32,
+) -> vec4<f32> {
+    // Cluster allocation debug (using 'over' alpha blending)
+#ifdef CLUSTERED_FORWARD_DEBUG_Z_SLICES
+    // NOTE: This debug mode visualises the z-slices
+    let cluster_overlay_alpha = 0.1;
+    var z_slice: u32 = view_z_to_z_slice(view_z, is_orthographic);
+    // A hack to make the colors alternate a bit more
+    if ((z_slice & 1u) == 1u) {
+        z_slice = z_slice + lights.cluster_dimensions.z / 2u;
+    }
+    let slice_color = hsv2rgb(f32(z_slice) / f32(lights.cluster_dimensions.z + 1u), 1.0, 0.5);
+    output_color = vec4<f32>(
+        (1.0 - cluster_overlay_alpha) * output_color.rgb + cluster_overlay_alpha * slice_color,
+        output_color.a
+    );
+#endif // CLUSTERED_FORWARD_DEBUG_Z_SLICES
+#ifdef CLUSTERED_FORWARD_DEBUG_CLUSTER_LIGHT_COMPLEXITY
+    // NOTE: This debug mode visualises the number of lights within the cluster that contains
+    // the fragment. It shows a sort of lighting complexity measure.
+    let cluster_overlay_alpha = 0.1;
+    let max_light_complexity_per_cluster = 64.0;
+    output_color.r = (1.0 - cluster_overlay_alpha) * output_color.r
+        + cluster_overlay_alpha * smoothStep(0.0, max_light_complexity_per_cluster, f32(offset_and_count[1]));
+    output_color.g = (1.0 - cluster_overlay_alpha) * output_color.g
+        + cluster_overlay_alpha * (1.0 - smoothStep(0.0, max_light_complexity_per_cluster, f32(offset_and_count[1])));
+#endif // CLUSTERED_FORWARD_DEBUG_CLUSTER_LIGHT_COMPLEXITY
+#ifdef CLUSTERED_FORWARD_DEBUG_CLUSTER_COHERENCY
+    // NOTE: Visualizes the cluster to which the fragment belongs
+    let cluster_overlay_alpha = 0.1;
+    let cluster_color = hsv2rgb(random1D(f32(cluster_index)), 1.0, 0.5);
+    output_color = vec4<f32>(
+        (1.0 - cluster_overlay_alpha) * output_color.rgb + cluster_overlay_alpha * cluster_color,
+        output_color.a
+    );
+#endif // CLUSTERED_FORWARD_DEBUG_CLUSTER_COHERENCY
+
+    return output_color;
+}

crates/bevy_pbr/src/render/depth.wgsl

@@ -13,6 +13,9 @@ var<uniform> joint_matrices: SkinnedMesh;
 #import bevy_pbr::skinning
 #endif
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
 #ifdef SKINNED
@@ -34,6 +37,6 @@ fn vertex(vertex: Vertex) -> VertexOutput {
 #endif
 
     var out: VertexOutput;
-    out.clip_position = view.view_proj * model * vec4<f32>(vertex.position, 1.0);
+    out.clip_position = mesh_position_local_to_clip(model, vec4<f32>(vertex.position, 1.0));
     return out;
 }

crates/bevy_pbr/src/render/mesh.rs

@@ -43,6 +43,8 @@ pub const MESH_TYPES_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 2506024101911992377);
 pub const MESH_BINDINGS_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 16831548636314682308);
+pub const MESH_FUNCTIONS_HANDLE: HandleUntyped =
+    HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 6300874327833745635);
 pub const MESH_SHADER_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 3252377289100772450);
 pub const SKINNING_HANDLE: HandleUntyped =
@@ -69,6 +71,12 @@ impl Plugin for MeshRenderPlugin {
             "mesh_bindings.wgsl",
             Shader::from_wgsl
         );
+        load_internal_asset!(
+            app,
+            MESH_FUNCTIONS_HANDLE,
+            "mesh_functions.wgsl",
+            Shader::from_wgsl
+        );
         load_internal_asset!(app, MESH_SHADER_HANDLE, "mesh.wgsl", Shader::from_wgsl);
         load_internal_asset!(app, SKINNING_HANDLE, "skinning.wgsl", Shader::from_wgsl);
 

crates/bevy_pbr/src/render/mesh.wgsl

@@ -1,6 +1,9 @@
 #import bevy_pbr::mesh_view_bindings
 #import bevy_pbr::mesh_bindings
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
     [[location(1)]] normal: vec3<f32>;
@@ -35,35 +38,21 @@ fn vertex(vertex: Vertex) -> VertexOutput {
     var out: VertexOutput;
 #ifdef SKINNED
     var model = skin_model(vertex.joint_indices, vertex.joint_weights);
-    out.world_position = model * vec4<f32>(vertex.position, 1.0);
     out.world_normal = skin_normals(model, vertex.normal);
-#ifdef VERTEX_TANGENTS
-    out.world_tangent = skin_tangents(model, vertex.tangent);
-#endif
 #else
-    out.world_position = mesh.model * vec4<f32>(vertex.position, 1.0);
-    out.world_normal = mat3x3<f32>(
-        mesh.inverse_transpose_model[0].xyz,
-        mesh.inverse_transpose_model[1].xyz,
-        mesh.inverse_transpose_model[2].xyz
-    ) * vertex.normal;
-#ifdef VERTEX_TANGENTS
-    out.world_tangent = vec4<f32>(
-        mat3x3<f32>(
-            mesh.model[0].xyz,
-            mesh.model[1].xyz,
-            mesh.model[2].xyz
-        ) * vertex.tangent.xyz,
-        vertex.tangent.w
-    );
+    var model = mesh.model;
 #endif
+    out.world_position = mesh_position_local_to_world(model, vec4<f32>(vertex.position, 1.0));
+    out.world_normal = mesh_normal_local_to_world(vertex.normal);
+    out.uv = vertex.uv;
+#ifdef VERTEX_TANGENTS
+    out.world_tangent = mesh_tangent_local_to_world(model, vertex.tangent);
 #endif
 #ifdef VERTEX_COLORS
     out.color = vertex.color;
 #endif
 
-    out.uv = vertex.uv;
-    out.clip_position = view.view_proj * out.world_position;
+    out.clip_position = mesh_position_world_to_clip(out.world_position);
     return out;
 }
 

crates/bevy_pbr/src/render/mesh_functions.wgsl

@@ -0,0 +1,36 @@
+#define_import_path bevy_pbr::mesh_functions
+
+fn mesh_position_local_to_world(model: mat4x4<f32>, vertex_position: vec4<f32>) -> vec4<f32> {
+    return model * vertex_position;
+}
+
+fn mesh_position_world_to_clip(world_position: vec4<f32>) -> vec4<f32> {
+    return view.view_proj * world_position;
+}
+
+// NOTE: The intermediate world_position assignment is important
+// for precision purposes when using the 'equals' depth comparison
+// function.
+fn mesh_position_local_to_clip(model: mat4x4<f32>, vertex_position: vec4<f32>) -> vec4<f32> {
+    let world_position = mesh_position_local_to_world(model, vertex_position);
+    return mesh_position_world_to_clip(world_position);
+}
+
+fn mesh_normal_local_to_world(vertex_normal: vec3<f32>) -> vec3<f32> {
+    return mat3x3<f32>(
+        mesh.inverse_transpose_model[0].xyz,
+        mesh.inverse_transpose_model[1].xyz,
+        mesh.inverse_transpose_model[2].xyz
+    ) * vertex_normal;
+}
+
+fn mesh_tangent_local_to_world(model: mat4x4<f32>, vertex_tangent: vec4<f32>) -> vec4<f32> {
+    return vec4<f32>(
+        mat3x3<f32>(
+            model[0].xyz,
+            model[1].xyz,
+            model[2].xyz
+        ) * vertex_tangent.xyz,
+        vertex_tangent.w
+    );
+}