Implement MIS into the PT

crates/bevy_solari/src/pathtracer/pathtracer.wgsl

@@ -4,7 +4,7 @@
 #import bevy_render::maths::PI
 #import bevy_render::view::View
 #import bevy_solari::brdf::evaluate_brdf
-#import bevy_solari::sampling::{sample_random_light, sample_ggx_vndf, ggx_vndf_pdf}
+#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>;
@@ -37,30 +37,42 @@ 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;
 
-            // Use emissive only on the first ray (coming from the camera)
-            if ray_t_min == 0.0 { radiance = ray_hit.material.emissive; }
+            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;
 
             // Sample direct lighting
             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 += throughput * direct_lighting.radiance * direct_lighting.inverse_pdf * direct_lighting_brdf;
+            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
             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 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;
+            throughput *= next_bounce.mis_weight * (brdf * cos_theta) / next_bounce.pdf;
 
             // Russian roulette for early termination
             let p = luminance(throughput);
@@ -80,12 +92,18 @@ fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
 struct NextBounce {
     wi: vec3<f32>,
+    mis_weight: f32,
     pdf: f32,
+    perfectly_specular_bounce: bool,
 }
 
 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 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, 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];
@@ -96,7 +114,8 @@ fn importance_sample_next_bounce(wo: vec3<f32>, ray_hit: ResolvedRayHitFull, rng
 
     var wi: vec3<f32>;
     var wi_tangent: vec3<f32>;
-    if rand_f(rng) < diffuse_weight {
+    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 {
@@ -107,6 +126,25 @@ fn importance_sample_next_bounce(wo: vec3<f32>, ray_hit: ResolvedRayHitFull, rng
     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);
+    let mis_weight = select(balance_heuristic(specular_pdf, diffuse_pdf), balance_heuristic(diffuse_pdf, specular_pdf), diffuse_selected);
+
+    return NextBounce(wi, mis_weight, 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;
 
-    return NextBounce(wi, pdf);
+    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

@@ -190,6 +190,7 @@ 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);
@@ -352,6 +353,7 @@ struct GpuInstanceGeometryIds {
     vertex_buffer_offset: u32,
     index_buffer_id: u32,
     index_buffer_offset: u32,
+    triangle_count: u32,
 }
 
 #[derive(ShaderType)]

crates/bevy_solari/src/scene/raytracing_scene_bindings.wgsl

@@ -8,6 +8,7 @@ struct InstanceGeometryIds {
     vertex_buffer_offset: u32,
     index_buffer_id: u32,
     index_buffer_offset: u32,
+    triangle_count: u32,
 }
 
 struct VertexBuffer { vertices: array<PackedVertex> }
@@ -115,6 +116,7 @@ struct ResolvedRayHitFull {
     world_tangent: vec4<f32>,
     uv: vec2<f32>,
     triangle_area: f32,
+    triangle_count: u32,
     material: ResolvedMaterial,
 }
 
@@ -192,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, world_tangent, 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);
 }

crates/bevy_solari/src/scene/sampling.wgsl

@@ -3,7 +3,15 @@
 #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, 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}
+#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, ResolvedRayHitFull}
+
+fn power_heuristic(f: f32, g: f32) -> f32 {
+    return f * f / (f * f + g * g);
+}
+
+fn balance_heuristic(f: f32, g: f32) -> f32 {
+    return f / (f + g);
+}
 
 // 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> {
@@ -80,6 +88,12 @@ fn sample_random_light(ray_origin: vec3<f32>, origin_world_normal: vec3<f32>, rn
     return light_contribution;
 }
 
+fn random_light_pdf(hit: ResolvedRayHitFull) -> f32 {
+    let light_count = arrayLength(&light_sources);
+    let p_light = 1.0 / f32(light_count);
+    return p_light / (hit.triangle_area * f32(hit.triangle_count));
+}
+
 fn generate_random_light_sample(rng: ptr<function, u32>) -> GenerateRandomLightSampleResult {
     let light_count = arrayLength(&light_sources);
     let light_id = rand_range_u(light_count, rng);

release-content/release-notes/bevy_solari.md

@@ -1,6 +1,6 @@
 ---
 title: Initial raytraced lighting progress (bevy_solari)
-authors: ["@JMS55"]
+authors: ["@JMS55", "@SparkyPotato"]
 pull_requests: [19058, 19620, 19790, 20020, 20113, 20213, 20242]
 ---