PBR/material support

crates/bevy_pbr/src/material.rs

@@ -447,7 +447,7 @@ pub const fn tonemapping_pipeline_key(tonemapping: Tonemapping) -> MeshPipelineK
     }
 }
 
-const fn screen_space_specular_transmission_pipeline_key(
+pub const fn screen_space_specular_transmission_pipeline_key(
     screen_space_transmissive_blur_quality: ScreenSpaceTransmissionQuality,
 ) -> MeshPipelineKey {
     match screen_space_transmissive_blur_quality {
@@ -488,15 +488,15 @@ pub fn queue_material_meshes<M: Material>(
         Option<&DebandDither>,
         Option<&EnvironmentMapLight>,
         Option<&ShadowFilteringMethod>,
-        Option<&ScreenSpaceAmbientOcclusionSettings>,
+        Has<ScreenSpaceAmbientOcclusionSettings>,
         (
             Has<NormalPrepass>,
             Has<DepthPrepass>,
             Has<MotionVectorPrepass>,
             Has<DeferredPrepass>,
         ),
         Option<&Camera3d>,
-        Option<&TemporalJitter>,
+        Has<TemporalJitter>,
         Option<&Projection>,
         &mut RenderPhase<Opaque3d>,
         &mut RenderPhase<AlphaMask3d>,
@@ -535,20 +535,17 @@ pub fn queue_material_meshes<M: Material>(
         if normal_prepass {
             view_key |= MeshPipelineKey::NORMAL_PREPASS;
         }
-
         if depth_prepass {
             view_key |= MeshPipelineKey::DEPTH_PREPASS;
         }
-
         if motion_vector_prepass {
             view_key |= MeshPipelineKey::MOTION_VECTOR_PREPASS;
         }
-
         if deferred_prepass {
             view_key |= MeshPipelineKey::DEFERRED_PREPASS;
         }
 
-        if temporal_jitter.is_some() {
+        if temporal_jitter {
             view_key |= MeshPipelineKey::TEMPORAL_JITTER;
         }
 
@@ -586,7 +583,7 @@ pub fn queue_material_meshes<M: Material>(
                 view_key |= MeshPipelineKey::DEBAND_DITHER;
             }
         }
-        if ssao.is_some() {
+        if ssao {
             view_key |= MeshPipelineKey::SCREEN_SPACE_AMBIENT_OCCLUSION;
         }
 

crates/bevy_pbr/src/meshlet/material_draw_nodes.rs

@@ -1,6 +1,6 @@
 use super::{
     gpu_scene::{MeshletViewBindGroups, MeshletViewResources},
-    prepare_materials::MeshletViewMaterials,
+    material_draw_prepare::MeshletViewMaterials,
     MeshletGpuScene,
 };
 use crate::{MeshViewBindGroup, ViewFogUniformOffset, ViewLightsUniformOffset};

crates/bevy_pbr/src/meshlet/prepare_materials.rs → crates/bevy_pbr/src/meshlet/material_draw_prepare.rs

@@ -2,11 +2,13 @@ use super::{MeshletGpuScene, MESHLET_MESH_MATERIAL_SHADER_HANDLE};
 use crate::*;
 use bevy_asset::AssetServer;
 use bevy_core_pipeline::{
+    core_3d::Camera3d,
     prepass::{DeferredPrepass, DepthPrepass, MotionVectorPrepass, NormalPrepass},
     tonemapping::{DebandDither, Tonemapping},
 };
 use bevy_render::{
     camera::{Projection, TemporalJitter},
+    mesh::{InnerMeshVertexBufferLayout, Mesh, MeshVertexBufferLayout},
     render_asset::RenderAssets,
     render_resource::*,
     texture::BevyDefault,
@@ -21,9 +23,6 @@ pub struct MeshletViewMaterials {
     pub prepass: Vec<(u32, CachedRenderPipelineId, BindGroup)>,
 }
 
-// TODO: This whole thing is cursed
-// TODO: How to differentiate between main/prepass? Also, check not missing any keys/shaderdefs
-// TODO: Allow material specialization
 #[allow(clippy::too_many_arguments)]
 pub fn prepare_material_meshlet_meshes<M: Material>(
     mut gpu_scene: ResMut<MeshletGpuScene>,
@@ -43,10 +42,13 @@ pub fn prepare_material_meshlet_meshes<M: Material>(
         Option<&EnvironmentMapLight>,
         Option<&ShadowFilteringMethod>,
         Has<ScreenSpaceAmbientOcclusionSettings>,
-        Has<NormalPrepass>,
-        Has<DepthPrepass>,
-        Has<MotionVectorPrepass>,
-        Has<DeferredPrepass>,
+        (
+            Has<NormalPrepass>,
+            Has<DepthPrepass>,
+            Has<MotionVectorPrepass>,
+            Has<DeferredPrepass>,
+        ),
+        Option<&Camera3d>,
         Has<TemporalJitter>,
         Option<&Projection>,
     )>,
@@ -62,44 +64,84 @@ pub fn prepare_material_meshlet_meshes<M: Material>(
         environment_map,
         shadow_filter_method,
         ssao,
-        normal_prepass,
-        depth_prepass,
-        motion_vector_prepass,
-        deferred_prepass,
+        (normal_prepass, depth_prepass, motion_vector_prepass, deferred_prepass),
+        camera_3d,
         temporal_jitter,
         projection,
     ) in &views
     {
         let mut opaque_pass_material_map = Vec::new();
         let mut prepass_material_map = Vec::new();
 
-        let mut view_key = MeshPipelineKey::from_msaa_samples(1);
+        let fake_vertex_buffer_layout =
+            &MeshVertexBufferLayout::new(InnerMeshVertexBufferLayout::new(
+                vec![
+                    Mesh::ATTRIBUTE_POSITION.id,
+                    Mesh::ATTRIBUTE_NORMAL.id,
+                    Mesh::ATTRIBUTE_UV_0.id,
+                    Mesh::ATTRIBUTE_TANGENT.id,
+                ],
+                VertexBufferLayout {
+                    array_stride: 48,
+                    step_mode: VertexStepMode::Vertex,
+                    attributes: vec![
+                        VertexAttribute {
+                            format: Mesh::ATTRIBUTE_POSITION.format,
+                            offset: 0,
+                            shader_location: 0,
+                        },
+                        VertexAttribute {
+                            format: Mesh::ATTRIBUTE_NORMAL.format,
+                            offset: 12,
+                            shader_location: 1,
+                        },
+                        VertexAttribute {
+                            format: Mesh::ATTRIBUTE_UV_0.format,
+                            offset: 24,
+                            shader_location: 2,
+                        },
+                        VertexAttribute {
+                            format: Mesh::ATTRIBUTE_TANGENT.format,
+                            offset: 32,
+                            shader_location: 3,
+                        },
+                    ],
+                },
+            ));
 
-        if let Some(projection) = projection {
-            view_key |= match projection {
-                Projection::Perspective(_) => MeshPipelineKey::VIEW_PROJECTION_PERSPECTIVE,
-                Projection::Orthographic(_) => MeshPipelineKey::VIEW_PROJECTION_ORTHOGRAPHIC,
-            };
-        }
+        let mut view_key =
+            MeshPipelineKey::from_msaa_samples(1) | MeshPipelineKey::from_hdr(view.hdr);
 
-        if depth_prepass {
-            view_key |= MeshPipelineKey::DEPTH_PREPASS;
-        }
         if normal_prepass {
             view_key |= MeshPipelineKey::NORMAL_PREPASS;
         }
+        if depth_prepass {
+            view_key |= MeshPipelineKey::DEPTH_PREPASS;
+        }
         if motion_vector_prepass {
             view_key |= MeshPipelineKey::MOTION_VECTOR_PREPASS;
         }
         if deferred_prepass {
             view_key |= MeshPipelineKey::DEFERRED_PREPASS;
         }
 
+        if temporal_jitter {
+            view_key |= MeshPipelineKey::TEMPORAL_JITTER;
+        }
+
         let environment_map_loaded = environment_map.is_some_and(|map| map.is_loaded(&images));
+
         if environment_map_loaded {
             view_key |= MeshPipelineKey::ENVIRONMENT_MAP;
         }
 
+        if let Some(projection) = projection {
+            view_key |= match projection {
+                Projection::Perspective(_) => MeshPipelineKey::VIEW_PROJECTION_PERSPECTIVE,
+                Projection::Orthographic(_) => MeshPipelineKey::VIEW_PROJECTION_ORTHOGRAPHIC,
+            };
+        }
+
         match shadow_filter_method.unwrap_or(&ShadowFilteringMethod::default()) {
             ShadowFilteringMethod::Hardware2x2 => {
                 view_key |= MeshPipelineKey::SHADOW_FILTER_METHOD_HARDWARE_2X2;
@@ -121,15 +163,18 @@ pub fn prepare_material_meshlet_meshes<M: Material>(
                 view_key |= MeshPipelineKey::DEBAND_DITHER;
             }
         }
-
         if ssao {
             view_key |= MeshPipelineKey::SCREEN_SPACE_AMBIENT_OCCLUSION;
         }
 
-        if temporal_jitter {
-            view_key |= MeshPipelineKey::TEMPORAL_JITTER;
+        if let Some(camera_3d) = camera_3d {
+            view_key |= screen_space_specular_transmission_pipeline_key(
+                camera_3d.screen_space_specular_transmission_quality,
+            );
         }
 
+        view_key |= MeshPipelineKey::from_primitive_topology(PrimitiveTopology::TriangleList);
+
         for material_id in render_material_instances.values() {
             let Some(material) = render_materials.get(material_id) else {
                 continue;
@@ -138,6 +183,26 @@ pub fn prepare_material_meshlet_meshes<M: Material>(
                 continue;
             }
 
+            let Ok(material_pipeline_descriptor) = material_pipeline.specialize(
+                MaterialPipelineKey {
+                    mesh_key: view_key,
+                    bind_group_data: material.key.clone(),
+                },
+                fake_vertex_buffer_layout,
+            ) else {
+                continue;
+            };
+
+            let mut shader_defs = material_pipeline_descriptor
+                .fragment
+                .expect("TODO")
+                .shader_defs;
+            shader_defs.extend_from_slice(&[
+                "MESHLET_MESH_MATERIAL_PASS".into(),
+                ShaderDefVal::UInt("MESHLET_BIND_GROUP".into(), 1),
+                "MATERIAL_BIND_GROUP_2".into(),
+            ]);
+
             let pipeline_descriptor = RenderPipelineDescriptor {
                 label: Some("meshlet_material_draw".into()),
                 layout: vec![
@@ -148,10 +213,7 @@ pub fn prepare_material_meshlet_meshes<M: Material>(
                 push_constant_ranges: vec![],
                 vertex: VertexState {
                     shader: MESHLET_MESH_MATERIAL_SHADER_HANDLE,
-                    shader_defs: vec![
-                        "MESHLET_MESH_MATERIAL_PASS".into(),
-                        ShaderDefVal::UInt("MESHLET_BIND_GROUP".into(), 1),
-                    ],
+                    shader_defs: shader_defs.clone(),
                     entry_point: "vertex".into(),
                     buffers: Vec::new(),
                 },
@@ -170,10 +232,7 @@ pub fn prepare_material_meshlet_meshes<M: Material>(
                         ShaderRef::Handle(handle) => handle,
                         ShaderRef::Path(path) => asset_server.load(path),
                     },
-                    shader_defs: vec![
-                        "MESHLET_MESH_MATERIAL_PASS".into(),
-                        ShaderDefVal::UInt("MESHLET_BIND_GROUP".into(), 1),
-                    ],
+                    shader_defs,
                     entry_point: "fragment".into(),
                     targets: vec![Some(ColorTargetState {
                         format: if view.hdr {

crates/bevy_pbr/src/meshlet/mod.rs

@@ -2,15 +2,15 @@ mod asset;
 mod from_mesh;
 mod gpu_scene;
 mod material_draw_nodes;
+mod material_draw_prepare;
 mod persistent_buffer;
 mod persistent_buffer_impls;
 mod pipelines;
-mod prepare_materials;
 mod visibility_buffer_node;
 
 pub(crate) use self::{
     gpu_scene::{queue_material_meshlet_meshes, MeshletGpuScene},
-    prepare_materials::prepare_material_meshlet_meshes,
+    material_draw_prepare::prepare_material_meshlet_meshes,
 };
 
 pub use self::{

crates/bevy_pbr/src/meshlet/visibility_buffer_resolve.wgsl

@@ -55,7 +55,7 @@ fn compute_partial_derivatives(vertex_clip_positions: array<vec4<f32>, 3>, ndc_u
 }
 
 struct VertexOutput {
-    frag_coord: vec4<f32>,
+    position: vec4<f32>,
     world_position: vec4<f32>,
     world_normal: vec3<f32>,
     uv: vec2<f32>,

crates/bevy_pbr/src/pbr_material.rs

@@ -788,6 +788,10 @@ impl Material for StandardMaterial {
         PBR_SHADER_HANDLE.into()
     }
 
+    fn meshlet_mesh_fragment_shader() -> ShaderRef {
+        PBR_SHADER_HANDLE.into()
+    }
+
     #[inline]
     fn alpha_mode(&self) -> AlphaMode {
         self.alpha_mode

crates/bevy_pbr/src/render/pbr.wgsl

@@ -16,13 +16,25 @@
 }
 #endif
 
+#ifdef MESHLET_MESH_MATERIAL_PASS
+#import bevy_pbr::meshlet_visibility_buffer_resolve::resolve_vertex_output
+#endif
+
 @fragment
 fn fragment(
+#ifdef MESHLET_MESH_MATERIAL_PASS
+    @builtin(position) frag_coord: vec4<f32>,
+#else
     in: VertexOutput,
     @builtin(front_facing) is_front: bool,
+#endif
 ) -> FragmentOutput {
     // generate a PbrInput struct from the StandardMaterial bindings
+#ifdef MESHLET_MESH_MATERIAL_PASS
+    var pbr_input = pbr_input_from_standard_material(resolve_vertex_output(frag_coord), true);
+#else
     var pbr_input = pbr_input_from_standard_material(in, is_front);
+#endif
 
     // alpha discard
     pbr_input.material.base_color = alpha_discard(pbr_input.material, pbr_input.material.base_color);

crates/bevy_pbr/src/render/pbr_bindings.wgsl

@@ -2,6 +2,7 @@
 
 #import bevy_pbr::pbr_types::StandardMaterial
 
+#ifndef MATERIAL_BIND_GROUP_2
 @group(1) @binding(0) var<uniform> material: StandardMaterial;
 @group(1) @binding(1) var base_color_texture: texture_2d<f32>;
 @group(1) @binding(2) var base_color_sampler: sampler;
@@ -23,3 +24,26 @@
 @group(1) @binding(17) var diffuse_transmission_texture: texture_2d<f32>;
 @group(1) @binding(18) var diffuse_transmission_sampler: sampler;
 #endif
+#else
+@group(2) @binding(0) var<uniform> material: StandardMaterial;
+@group(2) @binding(1) var base_color_texture: texture_2d<f32>;
+@group(2) @binding(2) var base_color_sampler: sampler;
+@group(2) @binding(3) var emissive_texture: texture_2d<f32>;
+@group(2) @binding(4) var emissive_sampler: sampler;
+@group(2) @binding(5) var metallic_roughness_texture: texture_2d<f32>;
+@group(2) @binding(6) var metallic_roughness_sampler: sampler;
+@group(2) @binding(7) var occlusion_texture: texture_2d<f32>;
+@group(2) @binding(8) var occlusion_sampler: sampler;
+@group(2) @binding(9) var normal_map_texture: texture_2d<f32>;
+@group(2) @binding(10) var normal_map_sampler: sampler;
+@group(2) @binding(11) var depth_map_texture: texture_2d<f32>;
+@group(2) @binding(12) var depth_map_sampler: sampler;
+#ifdef PBR_TRANSMISSION_TEXTURES_SUPPORTED
+@group(2) @binding(13) var specular_transmission_texture: texture_2d<f32>;
+@group(2) @binding(14) var specular_transmission_sampler: sampler;
+@group(2) @binding(15) var thickness_texture: texture_2d<f32>;
+@group(2) @binding(16) var thickness_sampler: sampler;
+@group(2) @binding(17) var diffuse_transmission_texture: texture_2d<f32>;
+@group(2) @binding(18) var diffuse_transmission_sampler: sampler;
+#endif
+#endif