1#import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
2#import bevy_pbr::pbr_functions::calculate_tbn_mikktspace
3#import bevy_pbr::utils::{rand_f, rand_vec2f, sample_cosine_hemisphere}
4#import bevy_render::maths::PI
5#import bevy_render::view::View
6#import bevy_solari::brdf::evaluate_brdf
7#import bevy_solari::sampling::{sample_random_light, random_emissive_light_pdf, sample_ggx_vndf, ggx_vndf_pdf, power_heuristic}
8#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, ResolvedRayHitFull, RAY_T_MIN, RAY_T_MAX}
9
10@group(1) @binding(0) var accumulation_texture: texture_storage_2d<rgba32float, read_write>;
11@group(1) @binding(1) var view_output: texture_storage_2d<rgba16float, write>;
12@group(1) @binding(2) var<uniform> view: View;
13
14@compute @workgroup_size(8, 8, 1)
15fn pathtrace(@builtin(global_invocation_id) global_id: vec3<u32>) {
16    if any(global_id.xy >= vec2u(view.viewport.zw)) {
17        return;
18    }
19
20    let old_color = textureLoad(accumulation_texture, global_id.xy);
21
22    // Setup RNG
23    let pixel_index = global_id.x + global_id.y * u32(view.viewport.z);
24    let frame_index = u32(old_color.a) * 5782582u;
25    var rng = pixel_index + frame_index;
26
27    // Shoot the first ray from the camera
28    let pixel_center = vec2<f32>(global_id.xy) + 0.5;
29    let jitter = rand_vec2f(&rng) - 0.5;
30    let pixel_uv = (pixel_center + jitter) / view.viewport.zw;
31    let pixel_ndc = (pixel_uv * 2.0) - 1.0;
32    let primary_ray_target = view.world_from_clip * vec4(pixel_ndc.x, -pixel_ndc.y, 1.0, 1.0);
33    var ray_origin = view.world_position;
34    var ray_direction = normalize((primary_ray_target.xyz / primary_ray_target.w) - ray_origin);
35    var ray_t_min = 0.0;
36
37    // Path trace
38    var radiance = vec3(0.0);
39    var throughput = vec3(1.0);
40    var p_bounce = 0.0;
41    var bounce_was_perfect_reflection = true;
42    var previous_normal = vec3(0.0);
43    loop {
44        let ray_hit = trace_ray(ray_origin, ray_direction, ray_t_min, RAY_T_MAX, RAY_FLAG_NONE);
45        if ray_hit.kind != RAY_QUERY_INTERSECTION_NONE {
46            let ray_hit = resolve_ray_hit_full(ray_hit);
47            let wo = -ray_direction;
48
49            var mis_weight = 1.0;
50            if !bounce_was_perfect_reflection {
51                let p_light = random_emissive_light_pdf(ray_hit);
52                mis_weight = power_heuristic(p_bounce, p_light);
53            }
54            radiance += mis_weight * throughput * ray_hit.material.emissive;
55
56            // Sample direct lighting, but only if the surface is not mirror-like
57            let is_perfectly_specular = ray_hit.material.roughness <= 0.001 && ray_hit.material.metallic > 0.9999;
58            if !is_perfectly_specular {
59                let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, &rng);
60
61                mis_weight = 1.0;
62                if direct_lighting.brdf_rays_can_hit {
63                    let pdf_of_bounce = brdf_pdf(wo, direct_lighting.wi, ray_hit);
64                    mis_weight = power_heuristic(1.0 / direct_lighting.inverse_pdf, pdf_of_bounce);
65                }
66
67                let direct_lighting_brdf = evaluate_brdf(ray_hit.world_normal, wo, direct_lighting.wi, ray_hit.material);
68                radiance += mis_weight * throughput * direct_lighting.radiance * direct_lighting.inverse_pdf * direct_lighting_brdf;
69            }
70
71            // Sample new ray direction from the material BRDF for next bounce
72            let next_bounce = importance_sample_next_bounce(wo, ray_hit, &rng);
73            ray_direction = next_bounce.wi;
74            ray_origin = ray_hit.world_position;
75            ray_t_min = RAY_T_MIN;
76            p_bounce = next_bounce.pdf;
77            bounce_was_perfect_reflection = next_bounce.perfectly_specular_bounce;
78            previous_normal = ray_hit.world_normal;
79
80            // Update throughput for next bounce
81            let brdf = evaluate_brdf(ray_hit.world_normal, wo, next_bounce.wi, ray_hit.material);
82            let cos_theta = dot(next_bounce.wi, ray_hit.world_normal);
83            throughput *= (brdf * cos_theta) / next_bounce.pdf;
84
85            // Russian roulette for early termination
86            let p = luminance(throughput);
87            if rand_f(&rng) > p { break; }
88            throughput /= p;
89        } else { break; }
90    }
91
92    // Camera exposure
93    radiance *= view.exposure;
94
95    // Accumulation over time via running average
96    let new_color = mix(old_color.rgb, radiance, 1.0 / (old_color.a + 1.0));
97    textureStore(accumulation_texture, global_id.xy, vec4(new_color, old_color.a + 1.0));
98    textureStore(view_output, global_id.xy, vec4(new_color, 1.0));
99}
100
101struct NextBounce {
102    wi: vec3<f32>,
103    pdf: f32,
104    perfectly_specular_bounce: bool,
105}
106
107fn importance_sample_next_bounce(wo: vec3<f32>, ray_hit: ResolvedRayHitFull, rng: ptr<function, u32>) -> NextBounce {
108    let is_perfectly_specular = ray_hit.material.roughness <= 0.001 && ray_hit.material.metallic > 0.9999;
109    if is_perfectly_specular {
110        return NextBounce(reflect(-wo, ray_hit.world_normal), 1.0, true);
111    }
112    let diffuse_weight = mix(mix(0.4, 0.9, ray_hit.material.perceptual_roughness), 0.0, ray_hit.material.metallic);
113    let specular_weight = 1.0 - diffuse_weight;
114
115    let TBN = calculate_tbn_mikktspace(ray_hit.world_normal, ray_hit.world_tangent);
116    let T = TBN[0];
117    let B = TBN[1];
118    let N = TBN[2];
119
120    let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
121
122    var wi: vec3<f32>;
123    var wi_tangent: vec3<f32>;
124    let diffuse_selected = rand_f(rng) < diffuse_weight;
125    if diffuse_selected {
126        wi = sample_cosine_hemisphere(ray_hit.world_normal, rng);
127        wi_tangent = vec3(dot(wi, T), dot(wi, B), dot(wi, N));
128    } else {
129        wi_tangent = sample_ggx_vndf(wo_tangent, ray_hit.material.roughness, rng);
130        wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N;
131    }
132
133    let diffuse_pdf = dot(wi, ray_hit.world_normal) / PI;
134    let specular_pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness);
135    let pdf = (diffuse_weight * diffuse_pdf) + (specular_weight * specular_pdf);
136
137    return NextBounce(wi, pdf, false);
138}
139
140fn brdf_pdf(wo: vec3<f32>, wi: vec3<f32>, ray_hit: ResolvedRayHitFull) -> f32 {
141    let diffuse_weight = mix(mix(0.4, 0.9, ray_hit.material.roughness), 0.0, ray_hit.material.metallic);
142    let specular_weight = 1.0 - diffuse_weight;
143
144    let TBN = calculate_tbn_mikktspace(ray_hit.world_normal, ray_hit.world_tangent);
145    let T = TBN[0];
146    let B = TBN[1];
147    let N = TBN[2];
148
149    let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
150    let wi_tangent = vec3(dot(wi, T), dot(wi, B), dot(wi, N));
151
152    let diffuse_pdf = wi_tangent.z / PI;
153    let specular_pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness);
154    let pdf = (diffuse_weight * diffuse_pdf) + (specular_weight * specular_pdf);
155    return pdf;
156}