1import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance;
2import bevy_pbr::render::utils::{rand_f, rand_vec2f};
3import bevy_render::view::View;
4import package::scene::brdf::{evaluate_brdf, evaluate_and_sample_brdf, brdf_pdf, F_AB};
5import package::scene::sampling::{sample_random_light, random_emissive_light_pdf, power_heuristic, sample_environment_map_light};
6import package::scene::bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX, MIRROR_ROUGHNESS_THRESHOLD};
7
8enable wgpu_ray_query;
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    loop {
42        let ray = trace_ray(ray_origin, ray_direction, ray_t_min, RAY_T_MAX, RAY_FLAG_NONE);
43        if ray.kind != RAY_QUERY_INTERSECTION_NONE {
44            let ray_hit = resolve_ray_hit_full(ray);
45            let wo = -ray_direction;
46            let NdotV = max(dot(ray_hit.world_normal, wo), 0.0001);
47            let F_ab = F_AB(ray_hit.material.perceptual_roughness, NdotV);
48
49            // Emissive contribution
50            var mis_weight = 1.0;
51            if p_bounce != 0.0 { // Not first bounce
52                let p_light = random_emissive_light_pdf(ray_hit, ray.t, NdotV);
53                mis_weight = power_heuristic(p_bounce, p_light);
54            }
55            if dot(ray_hit.world_normal, wo) > 0.0 {
56                radiance += mis_weight * throughput * ray_hit.material.emissive;
57            }
58
59            // Sample direct lighting, but only if the surface is not mirror-like
60            let is_perfectly_specular = ray_hit.material.roughness <= MIRROR_ROUGHNESS_THRESHOLD && ray_hit.material.metallic > 0.9999;
61            if !is_perfectly_specular {
62                let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, &rng);
63
64                mis_weight = 1.0;
65                if direct_lighting.brdf_rays_can_hit {
66                    let pdf_of_bounce = brdf_pdf(wo, direct_lighting.wi, ray_hit.world_normal, ray_hit.material, F_ab);
67                    mis_weight = power_heuristic(1.0 / direct_lighting.inverse_solid_angle_pdf, pdf_of_bounce);
68                }
69
70                let direct_lighting_brdf = evaluate_brdf(wo, direct_lighting.wi, ray_hit.world_normal, ray_hit.material, F_ab);
71                radiance += mis_weight * throughput * direct_lighting.radiance * direct_lighting.inverse_pdf * direct_lighting_brdf;
72            }
73
74            // Sample new ray direction from the material BRDF for next bounce and apply BRDF
75            let next_bounce = evaluate_and_sample_brdf(wo, ray_hit.world_normal, ray_hit.material, F_ab, &rng);
76            if next_bounce.pdf == 0.0 { break; }
77            ray_direction = next_bounce.wi;
78            ray_origin = ray_hit.world_position + (ray_hit.geometric_world_normal * RAY_T_MIN);
79            ray_t_min = RAY_T_MIN;
80            p_bounce = next_bounce.pdf;
81            throughput *= next_bounce.throughput;
82
83            // Russian roulette for early termination
84            let rr = saturate(luminance(throughput));
85            if rand_f(&rng) >= rr { break; }
86            throughput /= rr;
87        } else {
88            radiance += throughput * sample_environment_map_light(ray_direction);
89            break;
90        }
91    }
92
93    // Camera exposure
94    radiance *= view.exposure;
95
96    // Accumulation over time via running average
97    let new_color = mix(old_color.rgb, radiance, 1.0 / (old_color.a + 1.0));
98    textureStore(accumulation_texture, global_id.xy, vec4(new_color, old_color.a + 1.0));
99    textureStore(view_output, global_id.xy, vec4(new_color, 1.0));
100}