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, random_light_pdf, sample_ggx_vndf, ggx_vndf_pdf, balance_heuristic, power_heuristic}
+#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>;
@@ -35,32 +37,45 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
     // Path trace
     var radiance = vec3(0.0);
     var throughput = vec3(1.0);
+    var p_bounce = 0.0;
+    var bounce_was_perfect_reflection = true;
+    var previous_normal = vec3(0.0);
     loop {
         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);
+            let wo = -ray_direction;
 
-            // Evaluate material BRDF
-            let diffuse_brdf = ray_hit.material.base_color / PI;
+            var mis_weight = 1.0;
+            if !bounce_was_perfect_reflection {
+                let p_light = random_light_pdf(ray_hit);
+                mis_weight = power_heuristic(p_bounce, p_light);
+            }
+            radiance += mis_weight * throughput * ray_hit.material.emissive;
 
-            // 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;
+            // Sample direct lighting, but only if the surface is not mirror-like
+            let is_perfectly_specular = ray_hit.material.roughness < 0.0001 && ray_hit.material.metallic > 0.9999;
+            if !is_perfectly_specular {
+                let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, &rng);
+                let pdf_of_bounce = brdf_pdf(wo, direct_lighting.wi, ray_hit);
+                mis_weight = power_heuristic(1.0 / direct_lighting.inverse_pdf, pdf_of_bounce);
+                let direct_lighting_brdf = evaluate_brdf(ray_hit.world_normal, wo, direct_lighting.wi, ray_hit.material);
+                radiance += mis_weight * 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;
+            p_bounce = next_bounce.pdf;
+            bounce_was_perfect_reflection = next_bounce.perfectly_specular_bounce;
+            previous_normal = ray_hit.world_normal;
 
             // 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 +92,60 @@ 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,
+    perfectly_specular_bounce: bool,
+}
+
+fn importance_sample_next_bounce(wo: vec3<f32>, ray_hit: ResolvedRayHitFull, rng: ptr<function, u32>) -> NextBounce {
+    let is_perfectly_specular = ray_hit.material.roughness < 0.0001 && ray_hit.material.metallic > 0.9999;
+    if is_perfectly_specular {
+        return NextBounce(reflect(-wo, ray_hit.world_normal), 1.0, true);
+    }
+    let diffuse_weight = mix(mix(0.4f, 0.9f, ray_hit.material.perceptual_roughness), 0.f, ray_hit.material.metallic);
+    let specular_weight = 1.0 - diffuse_weight;
+
+    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>;
+    let diffuse_selected = rand_f(rng) < diffuse_weight;
+    if diffuse_selected {
+        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, false);
+}
+
+fn brdf_pdf(wo: vec3<f32>, wi: vec3<f32>, ray_hit: ResolvedRayHitFull) -> f32 {
+    let diffuse_weight = mix(mix(0.4f, 0.9f, ray_hit.material.roughness), 0.f, ray_hit.material.metallic);
+    let specular_weight = 1.0 - diffuse_weight;
+
+    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));
+    let wi_tangent = vec3(dot(wi, T), dot(wi, B), dot(wi, N));
+
+    let diffuse_pdf = wi_tangent.z / 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 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(),
         });
 
@@ -180,11 +190,12 @@ pub fn prepare_raytracing_scene_bindings(
             vertex_buffer_offset: vertex_slice.range.start,
             index_buffer_id,
             index_buffer_offset: index_slice.range.start,
+            triangle_count: (index_slice.range.len() / 3) as u32,
         });
 
         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(
@@ -342,16 +353,22 @@ struct GpuInstanceGeometryIds {
     vertex_buffer_offset: u32,
     index_buffer_id: u32,
     index_buffer_offset: u32,
+    triangle_count: u32,
 }
 
 #[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,10 +1,14 @@
 #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,
     index_buffer_id: u32,
     index_buffer_offset: u32,
+    triangle_count: u32,
 }
 
 struct VertexBuffer { vertices: array<PackedVertex> }
@@ -34,12 +38,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,14 +103,20 @@ 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,
+    triangle_count: u32,
     material: ResolvedMaterial,
 }
 
@@ -118,6 +133,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 +170,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 +194,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, instance_geometry_ids.triangle_count, resolved_material);
 }