Solari specular scene/PT support

crates/bevy_solari/src/pathtracer/pathtracer.wgsl

@@ -1,9 +1,11 @@
 #import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
+#import bevy_pbr::pbr_functions::calculate_tbn_mikktspace
 #import bevy_pbr::utils::{rand_f, rand_vec2f, sample_cosine_hemisphere}
 #import bevy_render::maths::PI
 #import bevy_render::view::View
-#import bevy_solari::sampling::sample_random_light
-#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX}
+#import bevy_solari::brdf::evaluate_brdf
+#import bevy_solari::sampling::{sample_random_light, sample_ggx_vndf, ggx_vndf_pdf}
+#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, ResolvedRayHitFull, RAY_T_MIN, RAY_T_MAX}
 
 @group(1) @binding(0) var accumulation_texture: texture_storage_2d<rgba32float, read_write>;
 @group(1) @binding(1) var view_output: texture_storage_2d<rgba16float, write>;
@@ -39,28 +41,26 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
         let ray_hit = trace_ray(ray_origin, ray_direction, ray_t_min, RAY_T_MAX, RAY_FLAG_NONE);
         if ray_hit.kind != RAY_QUERY_INTERSECTION_NONE {
             let ray_hit = resolve_ray_hit_full(ray_hit);
-
-            // Evaluate material BRDF
-            let diffuse_brdf = ray_hit.material.base_color / PI;
+            let wo = -ray_direction;
 
             // Use emissive only on the first ray (coming from the camera)
             if ray_t_min == 0.0 { radiance = ray_hit.material.emissive; }
 
             // Sample direct lighting
             let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, &rng);
-            radiance += throughput * diffuse_brdf * direct_lighting.radiance * direct_lighting.inverse_pdf;
+            let direct_lighting_brdf = evaluate_brdf(ray_hit.world_normal, wo, direct_lighting.wi, ray_hit.material);
+            radiance += throughput * direct_lighting.radiance * direct_lighting.inverse_pdf * direct_lighting_brdf;
 
             // Sample new ray direction from the material BRDF for next bounce
-            ray_direction = sample_cosine_hemisphere(ray_hit.world_normal, &rng);
-
-            // Update other variables for next bounce
+            let next_bounce = importance_sample_next_bounce(wo, ray_hit, &rng);
+            ray_direction = next_bounce.wi;
             ray_origin = ray_hit.world_position;
             ray_t_min = RAY_T_MIN;
 
             // Update throughput for next bounce
-            let cos_theta = dot(-ray_direction, ray_hit.world_normal);
-            let cosine_hemisphere_pdf = cos_theta / PI; // Weight for the next bounce because we importance sampled the diffuse BRDF for the next ray direction
-            throughput *= (diffuse_brdf * cos_theta) / cosine_hemisphere_pdf;
+            let brdf = evaluate_brdf(ray_hit.world_normal, wo, next_bounce.wi, ray_hit.material);
+            let cos_theta = dot(next_bounce.wi, ray_hit.world_normal);
+            throughput *= (brdf * cos_theta) / next_bounce.pdf;
 
             // Russian roulette for early termination
             let p = luminance(throughput);
@@ -77,3 +77,36 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
     textureStore(accumulation_texture, global_id.xy, vec4(new_color, old_color.a + 1.0));
     textureStore(view_output, global_id.xy, vec4(new_color, 1.0));
 }
+
+struct NextBounce {
+    wi: vec3<f32>,
+    pdf: f32,
+}
+
+fn importance_sample_next_bounce(wo: vec3<f32>, ray_hit: ResolvedRayHitFull, rng: ptr<function, u32>) -> NextBounce {
+    let diffuse_weight = 1.0 - ray_hit.material.metallic;
+    let specular_weight = ray_hit.material.metallic;
+
+    let TBN = calculate_tbn_mikktspace(ray_hit.world_normal, ray_hit.world_tangent);
+    let T = TBN[0];
+    let B = TBN[1];
+    let N = TBN[2];
+
+    let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
+
+    var wi: vec3<f32>;
+    var wi_tangent: vec3<f32>;
+    if rand_f(rng) < diffuse_weight {
+        wi = sample_cosine_hemisphere(ray_hit.world_normal, rng);
+        wi_tangent = vec3(dot(wi, T), dot(wi, B), dot(wi, N));
+    } else {
+        wi_tangent = sample_ggx_vndf(wo_tangent, ray_hit.material.roughness, rng);
+        wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N;
+    }
+
+    let diffuse_pdf = dot(wi, ray_hit.world_normal) / PI;
+    let specular_pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness);
+    let pdf = (diffuse_weight * diffuse_pdf) + (specular_weight * specular_pdf);
+
+    return NextBounce(wi, pdf);
+}

crates/bevy_solari/src/scene/binder.rs

@@ -115,13 +115,23 @@ pub fn prepare_raytracing_scene_bindings(
         let Some(emissive_texture_id) = process_texture(&material.emissive_texture) else {
             continue;
         };
+        let Some(metallic_roughness_texture_id) =
+            process_texture(&material.metallic_roughness_texture)
+        else {
+            continue;
+        };
 
         materials.get_mut().push(GpuMaterial {
-            base_color: material.base_color.to_linear(),
-            emissive: material.emissive,
-            base_color_texture_id,
             normal_map_texture_id,
+            base_color_texture_id,
             emissive_texture_id,
+            metallic_roughness_texture_id,
+
+            base_color: LinearRgba::from(material.base_color).to_vec3(),
+            perceptual_roughness: material.perceptual_roughness,
+            emissive: material.emissive.to_vec3(),
+            metallic: material.metallic,
+            reflectance: LinearRgba::from(material.specular_tint).to_vec3() * material.reflectance,
             _padding: Default::default(),
         });
 
@@ -184,7 +194,7 @@ pub fn prepare_raytracing_scene_bindings(
 
         material_ids.get_mut().push(material_id);
 
-        if material.emissive != LinearRgba::BLACK {
+        if material.emissive != Vec3::ZERO {
             light_sources
                 .get_mut()
                 .push(GpuLightSource::new_emissive_mesh_light(
@@ -346,12 +356,17 @@ struct GpuInstanceGeometryIds {
 
 #[derive(ShaderType)]
 struct GpuMaterial {
-    base_color: LinearRgba,
-    emissive: LinearRgba,
-    base_color_texture_id: u32,
     normal_map_texture_id: u32,
+    base_color_texture_id: u32,
     emissive_texture_id: u32,
-    _padding: u32,
+    metallic_roughness_texture_id: u32,
+
+    base_color: Vec3,
+    perceptual_roughness: f32,
+    emissive: Vec3,
+    metallic: f32,
+    reflectance: Vec3,
+    _padding: f32,
 }
 
 #[derive(ShaderType)]

crates/bevy_solari/src/scene/brdf.wgsl

@@ -0,0 +1,56 @@
+#define_import_path bevy_solari::brdf
+
+#import bevy_pbr::lighting::{F_AB, D_GGX, V_SmithGGXCorrelated, fresnel, specular_multiscatter}
+#import bevy_pbr::pbr_functions::{calculate_diffuse_color, calculate_F0}
+#import bevy_render::maths::PI
+#import bevy_solari::scene_bindings::ResolvedMaterial
+
+fn evaluate_brdf(
+    world_normal: vec3<f32>,
+    wo: vec3<f32>,
+    wi: vec3<f32>,
+    material: ResolvedMaterial,
+) -> vec3<f32> {
+    let diffuse_brdf = diffuse_brdf(material.base_color, material.metallic);
+    let specular_brdf = specular_brdf(
+        world_normal,
+        wo,
+        wi,
+        material.base_color,
+        material.metallic,
+        material.reflectance,
+        material.perceptual_roughness,
+        material.roughness,
+    );
+    return diffuse_brdf + specular_brdf;
+}
+
+fn diffuse_brdf(base_color: vec3<f32>, metallic: f32) -> vec3<f32> {
+    let diffuse_color = calculate_diffuse_color(base_color, metallic, 0.0, 0.0);
+    return diffuse_color / PI;
+}
+
+fn specular_brdf(
+    N: vec3<f32>,
+    V: vec3<f32>,
+    L: vec3<f32>,
+    base_color: vec3<f32>,
+    metallic: f32,
+    reflectance: vec3<f32>,
+    perceptual_roughness: f32,
+    roughness: f32,
+) -> vec3<f32> {
+    let H = normalize(L + V);
+    let NdotL = saturate(dot(N, L));
+    let NdotH = saturate(dot(N, H));
+    let LdotH = saturate(dot(L, H));
+    let NdotV = max(dot(N, V), 0.0001);
+
+    let F0 = calculate_F0(base_color, metallic, reflectance);
+    let F_ab = F_AB(perceptual_roughness, NdotV);
+
+    let D = D_GGX(roughness, NdotH);
+    let Vs = V_SmithGGXCorrelated(roughness, NdotV, NdotL);
+    let F = fresnel(F0, LdotH);
+    return specular_multiscatter(D, Vs, F, F0, F_ab, 1.0);
+}

crates/bevy_solari/src/scene/mod.rs

@@ -31,6 +31,7 @@ pub struct RaytracingScenePlugin;
 
 impl Plugin for RaytracingScenePlugin {
     fn build(&self, app: &mut App) {
+        load_shader_library!(app, "brdf.wgsl");
         load_shader_library!(app, "raytracing_scene_bindings.wgsl");
         load_shader_library!(app, "sampling.wgsl");
 

crates/bevy_solari/src/scene/raytracing_scene_bindings.wgsl

@@ -1,5 +1,8 @@
 #define_import_path bevy_solari::scene_bindings
 
+#import bevy_pbr::lighting::perceptualRoughnessToRoughness
+#import bevy_pbr::pbr_functions::calculate_tbn_mikktspace
+
 struct InstanceGeometryIds {
     vertex_buffer_id: u32,
     vertex_buffer_offset: u32,
@@ -34,12 +37,17 @@ fn unpack_vertex(packed: PackedVertex) -> Vertex {
 }
 
 struct Material {
-    base_color: vec4<f32>,
-    emissive: vec4<f32>,
-    base_color_texture_id: u32,
     normal_map_texture_id: u32,
+    base_color_texture_id: u32,
     emissive_texture_id: u32,
-    _padding: u32,
+    metallic_roughness_texture_id: u32,
+
+    base_color: vec3<f32>,
+    perceptual_roughness: f32,
+    emissive: vec3<f32>,
+    metallic: f32,
+    reflectance: vec3<f32>,
+    _padding: f32,
 }
 
 const TEXTURE_MAP_NONE = 0xFFFFFFFFu;
@@ -94,12 +102,17 @@ fn sample_texture(id: u32, uv: vec2<f32>) -> vec3<f32> {
 struct ResolvedMaterial {
     base_color: vec3<f32>,
     emissive: vec3<f32>,
+    reflectance: vec3<f32>,
+    perceptual_roughness: f32,
+    roughness: f32,
+    metallic: f32,
 }
 
 struct ResolvedRayHitFull {
     world_position: vec3<f32>,
     world_normal: vec3<f32>,
     geometric_world_normal: vec3<f32>,
+    world_tangent: vec4<f32>,
     uv: vec2<f32>,
     triangle_area: f32,
     material: ResolvedMaterial,
@@ -118,6 +131,17 @@ fn resolve_material(material: Material, uv: vec2<f32>) -> ResolvedMaterial {
         m.emissive *= sample_texture(material.emissive_texture_id, uv);
     }
 
+    m.reflectance = material.reflectance;
+
+    m.perceptual_roughness = material.perceptual_roughness;
+    m.metallic = material.metallic;
+    if material.metallic_roughness_texture_id != TEXTURE_MAP_NONE {
+        let metallic_roughness = sample_texture(material.metallic_roughness_texture_id, uv);
+        m.perceptual_roughness *= metallic_roughness.g;
+        m.metallic *= metallic_roughness.b;
+    }
+    m.roughness = clamp(m.perceptual_roughness * m.perceptual_roughness, 0.001, 1.0);
+
     return m;
 }
 
@@ -144,15 +168,20 @@ fn resolve_triangle_data_full(instance_id: u32, triangle_id: u32, barycentrics:
 
     let uv = mat3x2(vertices[0].uv, vertices[1].uv, vertices[2].uv) * barycentrics;
 
+    let local_tangent = mat3x3(vertices[0].tangent.xyz, vertices[1].tangent.xyz, vertices[2].tangent.xyz) * barycentrics;
+    let world_tangent = vec4(
+        normalize(mat3x3(transform[0].xyz, transform[1].xyz, transform[2].xyz) * local_tangent),
+        vertices[0].tangent.w,
+    );
+
     let local_normal = mat3x3(vertices[0].normal, vertices[1].normal, vertices[2].normal) * barycentrics; // TODO: Use barycentric lerp, ray_hit.object_to_world, cross product geo normal
     var world_normal = normalize(mat3x3(transform[0].xyz, transform[1].xyz, transform[2].xyz) * local_normal);
     let geometric_world_normal = world_normal;
     if material.normal_map_texture_id != TEXTURE_MAP_NONE {
-        let local_tangent = mat3x3(vertices[0].tangent.xyz, vertices[1].tangent.xyz, vertices[2].tangent.xyz) * barycentrics;
-        let world_tangent = normalize(mat3x3(transform[0].xyz, transform[1].xyz, transform[2].xyz) * local_tangent);
-        let N = world_normal;
-        let T = world_tangent;
-        let B = vertices[0].tangent.w * cross(N, T);
+        let TBN = calculate_tbn_mikktspace(world_normal, world_tangent);
+        let T = TBN[0];
+        let B = TBN[1];
+        let N = TBN[2];
         let Nt = sample_texture(material.normal_map_texture_id, uv);
         world_normal = normalize(Nt.x * T + Nt.y * B + Nt.z * N);
     }
@@ -163,5 +192,5 @@ fn resolve_triangle_data_full(instance_id: u32, triangle_id: u32, barycentrics:
 
     let resolved_material = resolve_material(material, uv);
 
-    return ResolvedRayHitFull(world_position, world_normal, geometric_world_normal, uv, triangle_area, resolved_material);
+    return ResolvedRayHitFull(world_position, world_normal, geometric_world_normal, world_tangent, uv, triangle_area, resolved_material);
 }

crates/bevy_solari/src/scene/sampling.wgsl

@@ -1,9 +1,50 @@
 #define_import_path bevy_solari::sampling
 
+#import bevy_pbr::lighting::D_GGX
 #import bevy_pbr::utils::{rand_f, rand_vec2f, rand_u, rand_range_u}
-#import bevy_render::maths::PI_2
+#import bevy_render::maths::{PI_2, orthonormalize}
 #import bevy_solari::scene_bindings::{trace_ray, RAY_T_MIN, RAY_T_MAX, light_sources, directional_lights, LightSource, LIGHT_SOURCE_KIND_DIRECTIONAL, resolve_triangle_data_full}
 
+// https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (Listing 1)
+fn sample_ggx_vndf(wi_tangent: vec3<f32>, roughness: f32, rng: ptr<function, u32>) -> vec3<f32> {
+    let i = wi_tangent;
+    let rand = rand_vec2f(rng);
+    let i_std = normalize(vec3(i.xy * roughness, i.z));
+    let phi = PI_2 * rand.x;
+    let a = roughness;
+    let s = 1.0 + length(vec2(i.xy));
+    let a2 = a * a;
+    let s2 = s * s;
+    let k = (1.0 - a2) * s2 / (s2 + a2 * i.z * i.z);
+    let b = select(i_std.z, k * i_std.z, i.z > 0.0);
+    let z = fma(1.0 - rand.y, 1.0 + b, -b);
+    let sin_theta = sqrt(saturate(1.0 - z * z));
+    let o_std = vec3(sin_theta * cos(phi), sin_theta * sin(phi), z);
+    let m_std = i_std + o_std;
+    let m = normalize(vec3(m_std.xy * roughness, m_std.z));
+    return 2.0 * dot(i, m) * m - i;
+}
+
+// https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (Listing 2)
+fn ggx_vndf_pdf(wi_tangent: vec3<f32>, wo_tangent: vec3<f32>, roughness: f32) -> f32 {
+    let i = wi_tangent;
+    let o = wo_tangent;
+    let m = normalize(i + o);
+    let ndf = D_GGX(roughness, saturate(m.z));
+    let ai = roughness * i.xy;
+    let len2 = dot(ai, ai);
+    let t = sqrt(len2 + i.z * i.z);
+    if i.z >= 0.0 {
+        let a = roughness;
+        let s = 1.0 + length(i.xy);
+        let a2 = a * a;
+        let s2 = s * s;
+        let k = (1.0 - a2) * s2 / (s2 + a2 * i.z * i.z);
+        return ndf / (2.0 * (k * i.z + t));
+    }
+    return ndf * (t - i.z) / (2.0 * len2);
+}
+
 struct LightSample {
     light_id: u32,
     seed: u32,

examples/3d/solari.rs

@@ -87,7 +87,7 @@ fn add_raytracing_meshes_on_scene_load(
     mut commands: Commands,
     args: Res<Args>,
 ) {
-    // Ensure meshes are bevy_solari compatible
+    // Ensure meshes are Solari compatible
     for (_, mesh) in meshes.iter_mut() {
         mesh.remove_attribute(Mesh::ATTRIBUTE_UV_1.id);
         mesh.remove_attribute(Mesh::ATTRIBUTE_COLOR.id);
@@ -113,8 +113,13 @@ fn add_raytracing_meshes_on_scene_load(
         }
     }
 
-    // Increase material emissive intensity to make it prettier for the example
+    // Adjust scene materials to better demo Solari features
     for (_, material) in materials.iter_mut() {
         material.emissive *= 200.0;
+
+        if material.base_color.to_linear() == LinearRgba::new(0.5, 0.5, 0.5, 1.0) {
+            material.metallic = 1.0;
+            material.perceptual_roughness = 0.15;
+        }
     }
 }

release-content/release-notes/bevy_solari.md

@@ -1,7 +1,7 @@
 ---
 title: Initial raytraced lighting progress (bevy_solari)
 authors: ["@JMS55"]
-pull_requests: [19058, 19620, 19790, 20020, 20113, 20213]
+pull_requests: [19058, 19620, 19790, 20020, 20113, 20213, 20242]
 ---
 
 (TODO: Embed solari example screenshot here)