pbr shader cleanup (bevyengine#10105)

# Objective

cleanup some pbr shader code. improve shader stage io consistency and
make pbr.wgsl (probably many people's first foray into bevy shader code)
a little more human-readable. also fix a couple of small issues with
deferred rendering.

## Solution

mesh_vertex_output: 
- rename to forward_io (to align with prepass_io)
- rename `MeshVertexOutput` to `VertexOutput` (to align with prepass_io)
- move `Vertex` from mesh.wgsl into here (to align with prepass_io)

prepass_io: 
- remove `FragmentInput`, use `VertexOutput` directly (to align with
forward_io)
- rename `VertexOutput::clip_position` to `position` (to align with
forward_io)

pbr.wgsl:
- restructure so we don't need `#ifdefs` on the actual entrypoint, use
VertexOutput and FragmentOutput in all cases and use #ifdefs to import
the right struct definitions.
- rearrange to make the flow clearer
- move alpha_discard up from `pbr_functions::pbr` to avoid needing to
call it on some branches and not others
- add a bunch of comments

deferred_lighting:
- move ssao into the `!unlit` block to reflect forward behaviour
correctly
- fix compile error with deferred + premultiply_alpha

## Migration Guide

in custom material shaders:
- `pbr_functions::pbr` no longer calls to
`pbr_functions::alpha_discard`. if you were using the `pbr` function in
a custom shader with alpha mask mode you now also need to call
alpha_discard manually
- rename imports of `bevy_pbr::mesh_vertex_output` to
`bevy_pbr::forward_io`
- rename instances of `MeshVertexOutput` to `VertexOutput`

in custom material prepass shaders:
- rename instances of `VertexOutput::clip_position` to
`VertexOutput::position`

assets/shaders/animate_shader.wgsl

@@ -1,6 +1,6 @@
 // The time since startup data is in the globals binding which is part of the mesh_view_bindings import
 #import bevy_pbr::mesh_view_bindings globals
-#import bevy_pbr::mesh_vertex_output MeshVertexOutput
+#import bevy_pbr::forward_io VertexOutput
 
 fn oklab_to_linear_srgb(c: vec3<f32>) -> vec3<f32> {
     let L = c.x;
@@ -23,7 +23,7 @@ fn oklab_to_linear_srgb(c: vec3<f32>) -> vec3<f32> {
 }
 
 @fragment
-fn fragment(in: MeshVertexOutput) -> @location(0) vec4<f32> {
+fn fragment(in: VertexOutput) -> @location(0) vec4<f32> {
     let speed = 2.0;
     // The globals binding contains various global values like time
     // which is the time since startup in seconds

assets/shaders/array_texture.wgsl

@@ -1,4 +1,4 @@
-#import bevy_pbr::mesh_vertex_output    MeshVertexOutput
+#import bevy_pbr::forward_io    VertexOutput
 #import bevy_pbr::mesh_view_bindings    view
 #import bevy_pbr::pbr_types             STANDARD_MATERIAL_FLAGS_DOUBLE_SIDED_BIT, PbrInput, pbr_input_new
 #import bevy_core_pipeline::tonemapping tone_mapping
@@ -10,7 +10,7 @@
 @fragment
 fn fragment(
     @builtin(front_facing) is_front: bool,
-    mesh: MeshVertexOutput,
+    mesh: VertexOutput,
 ) -> @location(0) vec4<f32> {
     let layer = i32(mesh.world_position.x) & 0x3;
 

assets/shaders/cubemap_unlit.wgsl

@@ -1,4 +1,4 @@
-#import bevy_pbr::mesh_vertex_output MeshVertexOutput
+#import bevy_pbr::forward_io VertexOutput
 
 #ifdef CUBEMAP_ARRAY
 @group(1) @binding(0) var base_color_texture: texture_cube_array<f32>;
@@ -10,7 +10,7 @@
 
 @fragment
 fn fragment(
-    mesh: MeshVertexOutput,
+    mesh: VertexOutput,
 ) -> @location(0) vec4<f32> {
     let fragment_position_view_lh = mesh.world_position.xyz * vec3<f32>(1.0, 1.0, -1.0);
     return textureSample(

assets/shaders/custom_material.wgsl

@@ -1,4 +1,4 @@
-#import bevy_pbr::mesh_vertex_output MeshVertexOutput
+#import bevy_pbr::forward_io VertexOutput
 // we can import items from shader modules in the assets folder with a quoted path
 #import "shaders/custom_material_import.wgsl" COLOR_MULTIPLIER
 
@@ -12,7 +12,7 @@ struct CustomMaterial {
 
 @fragment
 fn fragment(
-    mesh: MeshVertexOutput,
+    mesh: VertexOutput,
 ) -> @location(0) vec4<f32> {
     return material.color * textureSample(base_color_texture, base_color_sampler, mesh.uv) * COLOR_MULTIPLIER;
 }

assets/shaders/custom_material_screenspace_texture.wgsl

@@ -1,13 +1,13 @@
 #import bevy_pbr::mesh_view_bindings  view
-#import bevy_pbr::mesh_vertex_output  MeshVertexOutput
+#import bevy_pbr::forward_io  VertexOutput
 #import bevy_pbr::utils               coords_to_viewport_uv
 
 @group(1) @binding(0) var texture: texture_2d<f32>;
 @group(1) @binding(1) var texture_sampler: sampler;
 
 @fragment
 fn fragment(
-    mesh: MeshVertexOutput,
+    mesh: VertexOutput,
 ) -> @location(0) vec4<f32> {
     let viewport_uv = coords_to_viewport_uv(mesh.position.xy, view.viewport);
     let color = textureSample(texture, texture_sampler, viewport_uv);

assets/shaders/fallback_image_test.wgsl

@@ -1,6 +1,6 @@
 #import bevy_pbr::mesh_view_bindings
 #import bevy_pbr::mesh_bindings
-#import bevy_pbr::mesh_vertex_output MeshVertexOutput
+#import bevy_pbr::forward_io VertexOutput
 
 @group(1) @binding(0) var test_texture_1d: texture_1d<f32>;
 @group(1) @binding(1) var test_texture_1d_sampler: sampler;
@@ -21,4 +21,4 @@
 @group(1) @binding(11) var test_texture_3d_sampler: sampler;
 
 @fragment
-fn fragment(in: MeshVertexOutput) {}
+fn fragment(in: VertexOutput) {}

assets/shaders/line_material.wgsl

@@ -1,4 +1,4 @@
-#import bevy_pbr::mesh_vertex_output MeshVertexOutput
+#import bevy_pbr::forward_io VertexOutput
 
 struct LineMaterial {
     color: vec4<f32>,
@@ -8,7 +8,7 @@ struct LineMaterial {
 
 @fragment
 fn fragment(
-    mesh: MeshVertexOutput,
+    mesh: VertexOutput,
 ) -> @location(0) vec4<f32> {
     return material.color;
 }

assets/shaders/shader_defs.wgsl

@@ -1,4 +1,4 @@
-#import bevy_pbr::mesh_vertex_output MeshVertexOutput
+#import bevy_pbr::forward_io VertexOutput
 
 struct CustomMaterial {
     color: vec4<f32>,
@@ -8,7 +8,7 @@ struct CustomMaterial {
 
 @fragment
 fn fragment(
-    mesh: MeshVertexOutput,
+    mesh: VertexOutput,
 ) -> @location(0) vec4<f32> {
 #ifdef IS_RED
     return vec4<f32>(1.0, 0.0, 0.0, 1.0);

assets/shaders/show_prepass.wgsl

@@ -1,7 +1,7 @@
 #import bevy_pbr::mesh_types
 #import bevy_pbr::mesh_view_bindings  globals
 #import bevy_pbr::prepass_utils
-#import bevy_pbr::mesh_vertex_output  MeshVertexOutput
+#import bevy_pbr::forward_io  VertexOutput
 
 struct ShowPrepassSettings {
     show_depth: u32,
@@ -17,7 +17,7 @@ fn fragment(
 #ifdef MULTISAMPLED
     @builtin(sample_index) sample_index: u32,
 #endif
-    mesh: MeshVertexOutput,
+    mesh: VertexOutput,
 ) -> @location(0) vec4<f32> {
 #ifndef MULTISAMPLED
     let sample_index = 0u;

assets/shaders/texture_binding_array.wgsl

@@ -1,4 +1,4 @@
-#import bevy_pbr::mesh_vertex_output MeshVertexOutput
+#import bevy_pbr::forward_io VertexOutput
 
 @group(1) @binding(0) var textures: binding_array<texture_2d<f32>>;
 @group(1) @binding(1) var nearest_sampler: sampler;
@@ -7,7 +7,7 @@
 
 @fragment
 fn fragment(
-    mesh: MeshVertexOutput,
+    mesh: VertexOutput,
 ) -> @location(0) vec4<f32> {
     // Select the texture to sample from using non-uniform uv coordinates
     let coords = clamp(vec2<u32>(mesh.uv * 4.0), vec2<u32>(0u), vec2<u32>(3u));

assets/shaders/tonemapping_test_patterns.wgsl

@@ -1,6 +1,6 @@
 #import bevy_pbr::mesh_view_bindings
 #import bevy_pbr::mesh_bindings
-#import bevy_pbr::mesh_vertex_output  MeshVertexOutput
+#import bevy_pbr::forward_io  VertexOutput
 #import bevy_pbr::utils               PI
 
 #ifdef TONEMAP_IN_SHADER
@@ -43,7 +43,7 @@ fn continuous_hue(uv: vec2<f32>) -> vec3<f32> {
 
 @fragment
 fn fragment(
-    in: MeshVertexOutput,
+    in: VertexOutput,
 ) -> @location(0) vec4<f32> {
     var uv = in.uv;
     var out = vec3(0.0);

crates/bevy_pbr/src/deferred/deferred_lighting.wgsl

@@ -56,14 +56,15 @@ fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
     var pbr_input = pbr_input_from_deferred_gbuffer(frag_coord, deferred_data);
     var output_color = vec4(0.0);
 
-#ifdef SCREEN_SPACE_AMBIENT_OCCLUSION
-    let ssao = textureLoad(screen_space_ambient_occlusion_texture, vec2<i32>(in.position.xy), 0i).r;
-    let ssao_multibounce = gtao_multibounce(ssao, pbr_input.material.base_color.rgb);
-    pbr_input.occlusion = min(pbr_input.occlusion, ssao_multibounce);
-#endif // SCREEN_SPACE_AMBIENT_OCCLUSION
-
     // NOTE: Unlit bit not set means == 0 is true, so the true case is if lit
     if ((pbr_input.material.flags & STANDARD_MATERIAL_FLAGS_UNLIT_BIT) == 0u) {
+
+#ifdef SCREEN_SPACE_AMBIENT_OCCLUSION
+        let ssao = textureLoad(screen_space_ambient_occlusion_texture, vec2<i32>(in.position.xy), 0i).r;
+        let ssao_multibounce = gtao_multibounce(ssao, pbr_input.material.base_color.rgb);
+        pbr_input.occlusion = min(pbr_input.occlusion, ssao_multibounce);
+#endif // SCREEN_SPACE_AMBIENT_OCCLUSION
+
         output_color = pbr_functions::pbr(pbr_input);
     } else {
         output_color = pbr_input.material.base_color;
@@ -85,9 +86,6 @@ fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
     output_rgb = powsafe(output_rgb, 2.2);
     output_color = vec4(output_rgb, output_color.a);
 #endif
-#endif
-#ifdef PREMULTIPLY_ALPHA
-    output_color = pbr_functions::premultiply_alpha(material.flags, output_color);
 #endif
 
     return output_color;

crates/bevy_pbr/src/prepass/prepass.wgsl

@@ -1,6 +1,6 @@
 #import bevy_pbr::prepass_bindings
 #import bevy_pbr::mesh_functions
-#import bevy_pbr::prepass_io Vertex, VertexOutput, FragmentInput, FragmentOutput
+#import bevy_pbr::prepass_io Vertex, VertexOutput, FragmentOutput
 #import bevy_pbr::skinning
 #import bevy_pbr::morph
 #import bevy_pbr::mesh_bindings mesh
@@ -50,10 +50,10 @@ fn vertex(vertex_no_morph: Vertex) -> VertexOutput {
     var model = bevy_pbr::mesh_functions::get_model_matrix(vertex_no_morph.instance_index);
 #endif // SKINNED
 
-    out.clip_position = bevy_pbr::mesh_functions::mesh_position_local_to_clip(model, vec4(vertex.position, 1.0));
+    out.position = bevy_pbr::mesh_functions::mesh_position_local_to_clip(model, vec4(vertex.position, 1.0));
 #ifdef DEPTH_CLAMP_ORTHO
-    out.clip_position_unclamped = out.clip_position;
-    out.clip_position.z = min(out.clip_position.z, 1.0);
+    out.clip_position_unclamped = out.position;
+    out.position.z = min(out.position.z, 1.0);
 #endif // DEPTH_CLAMP_ORTHO
 
 #ifdef VERTEX_UVS
@@ -111,7 +111,7 @@ fn vertex(vertex_no_morph: Vertex) -> VertexOutput {
 
 #ifdef PREPASS_FRAGMENT
 @fragment
-fn fragment(in: FragmentInput) -> FragmentOutput {
+fn fragment(in: VertexOutput) -> FragmentOutput {
     var out: FragmentOutput;
 
 #ifdef NORMAL_PREPASS

crates/bevy_pbr/src/prepass/prepass_io.wgsl

@@ -32,7 +32,9 @@ struct Vertex {
 }
 
 struct VertexOutput {
-    @builtin(position) clip_position: vec4<f32>,
+    // This is `clip position` when the struct is used as a vertex stage output
+    // and `frag coord` when used as a fragment stage input
+    @builtin(position) position: vec4<f32>,
 
 #ifdef VERTEX_UVS
     @location(0) uv: vec2<f32>,
@@ -62,36 +64,6 @@ struct VertexOutput {
 #endif
 }
 
-struct FragmentInput {
-    @builtin(position) position: vec4<f32>,
-#ifdef VERTEX_UVS
-    @location(0) uv: vec2<f32>,
-#endif // VERTEX_UVS
-
-#ifdef NORMAL_PREPASS_OR_DEFERRED_PREPASS
-    @location(1) world_normal: vec3<f32>,
-#ifdef VERTEX_TANGENTS
-    @location(2) world_tangent: vec4<f32>,
-#endif
-#endif // NORMAL_PREPASS_OR_DEFERRED_PREPASS
-
-    @location(3) world_position: vec4<f32>,
-#ifdef MOTION_VECTOR_PREPASS
-    @location(4) previous_world_position: vec4<f32>,
-#endif // MOTION_VECTOR_PREPASS
-
-#ifdef DEPTH_CLAMP_ORTHO
-    @location(5) clip_position_unclamped: vec4<f32>,
-#endif // DEPTH_CLAMP_ORTHO
-#ifdef VERTEX_OUTPUT_INSTANCE_INDEX
-    @location(6) instance_index: u32,
-#endif
-
-#ifdef VERTEX_COLORS
-    @location(7) color: vec4<f32>,
-#endif
-};
-
 #ifdef PREPASS_FRAGMENT
 struct FragmentOutput {
 #ifdef NORMAL_PREPASS

crates/bevy_pbr/src/render/forward_io.wgsl

@@ -0,0 +1,52 @@
+#define_import_path bevy_pbr::forward_io
+
+struct Vertex {
+    @builtin(instance_index) instance_index: u32,
+#ifdef VERTEX_POSITIONS
+    @location(0) position: vec3<f32>,
+#endif
+#ifdef VERTEX_NORMALS
+    @location(1) normal: vec3<f32>,
+#endif
+#ifdef VERTEX_UVS
+    @location(2) uv: vec2<f32>,
+#endif
+// (Alternate UVs are at location 3, but they're currently unused here.)
+#ifdef VERTEX_TANGENTS
+    @location(4) tangent: vec4<f32>,
+#endif
+#ifdef VERTEX_COLORS
+    @location(5) color: vec4<f32>,
+#endif
+#ifdef SKINNED
+    @location(6) joint_indices: vec4<u32>,
+    @location(7) joint_weights: vec4<f32>,
+#endif
+#ifdef MORPH_TARGETS
+    @builtin(vertex_index) index: u32,
+#endif
+};
+
+struct VertexOutput {
+    // This is `clip position` when the struct is used as a vertex stage output
+    // and `frag coord` when used as a fragment stage input
+    @builtin(position) position: vec4<f32>,
+    @location(0) world_position: vec4<f32>,
+    @location(1) world_normal: vec3<f32>,
+#ifdef VERTEX_UVS
+    @location(2) uv: vec2<f32>,
+#endif
+#ifdef VERTEX_TANGENTS
+    @location(3) world_tangent: vec4<f32>,
+#endif
+#ifdef VERTEX_COLORS
+    @location(4) color: vec4<f32>,
+#endif
+#ifdef VERTEX_OUTPUT_INSTANCE_INDEX
+    @location(5) @interpolate(flat) instance_index: u32,
+#endif
+}
+
+struct FragmentOutput {
+    @location(0) color: vec4<f32>,
+}

crates/bevy_pbr/src/render/mesh.rs

@@ -62,7 +62,7 @@ use super::skin::SkinIndices;
 #[derive(Default)]
 pub struct MeshRenderPlugin;
 
-pub const MESH_VERTEX_OUTPUT: Handle<Shader> = Handle::weak_from_u128(2645551199423808407);
+pub const FORWARD_IO_HANDLE: Handle<Shader> = Handle::weak_from_u128(2645551199423808407);
 pub const MESH_VIEW_TYPES_HANDLE: Handle<Shader> = Handle::weak_from_u128(8140454348013264787);
 pub const MESH_VIEW_BINDINGS_HANDLE: Handle<Shader> = Handle::weak_from_u128(9076678235888822571);
 pub const MESH_TYPES_HANDLE: Handle<Shader> = Handle::weak_from_u128(2506024101911992377);
@@ -74,12 +74,7 @@ pub const MORPH_HANDLE: Handle<Shader> = Handle::weak_from_u128(9709828135876073
 
 impl Plugin for MeshRenderPlugin {
     fn build(&self, app: &mut bevy_app::App) {
-        load_internal_asset!(
-            app,
-            MESH_VERTEX_OUTPUT,
-            "mesh_vertex_output.wgsl",
-            Shader::from_wgsl
-        );
+        load_internal_asset!(app, FORWARD_IO_HANDLE, "forward_io.wgsl", Shader::from_wgsl);
         load_internal_asset!(
             app,
             MESH_VIEW_TYPES_HANDLE,