1#define_import_path bevy_solari::specular_gi
2
3#import bevy_pbr::pbr_functions::calculate_tbn_mikktspace
4#import bevy_render::maths::{orthonormalize, PI}
5#import bevy_render::view::View
6#import bevy_solari::brdf::{evaluate_brdf, evaluate_specular_brdf}
7#import bevy_solari::gbuffer_utils::gpixel_resolve
8#import bevy_solari::sampling::{sample_random_light, random_emissive_light_pdf, sample_ggx_vndf, ggx_vndf_pdf, power_heuristic}
9#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, ResolvedRayHitFull, RAY_T_MIN, RAY_T_MAX}
10#import bevy_solari::world_cache::{query_world_cache, get_cell_size, WORLD_CACHE_CELL_LIFETIME}
11
12@group(1) @binding(0) var view_output: texture_storage_2d<rgba16float, read_write>;
13@group(1) @binding(5) var<storage, read_write> gi_reservoirs_a: array<Reservoir>;
14@group(1) @binding(7) var gbuffer: texture_2d<u32>;
15@group(1) @binding(8) var depth_buffer: texture_depth_2d;
16@group(1) @binding(12) var<uniform> view: View;
17struct PushConstants { frame_index: u32, reset: u32 }
18var<push_constant> constants: PushConstants;
19
20const DIFFUSE_GI_REUSE_ROUGHNESS_THRESHOLD: f32 = 0.4;
21const SPECULAR_GI_FOR_DI_ROUGHNESS_THRESHOLD: f32 = 0.0225;
22const TERMINATE_IN_WORLD_CACHE_THRESHOLD: f32 = 0.03;
23
24@compute @workgroup_size(8, 8, 1)
25fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) {
26    if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; }
27
28    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
29    var rng = pixel_index + constants.frame_index;
30
31    let depth = textureLoad(depth_buffer, global_id.xy, 0);
32    if depth == 0.0 {
33        return;
34    }
35    let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip);
36
37    let wo_unnormalized = view.world_position - surface.world_position;
38    let wo = normalize(wo_unnormalized);
39
40    var radiance: vec3<f32>;
41    var wi: vec3<f32>;
42    if surface.material.roughness > DIFFUSE_GI_REUSE_ROUGHNESS_THRESHOLD {
43        // Surface is very rough, reuse the ReSTIR GI reservoir
44        let gi_reservoir = gi_reservoirs_a[pixel_index];
45        wi = normalize(gi_reservoir.sample_point_world_position - surface.world_position);
46        radiance = gi_reservoir.radiance * gi_reservoir.unbiased_contribution_weight;
47    } else {
48        // Surface is glossy or mirror-like, trace a new path
49        let TBN = orthonormalize(surface.world_normal);
50        let T = TBN[0];
51        let B = TBN[1];
52        let N = TBN[2];
53        let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
54        let wi_tangent = sample_ggx_vndf(wo_tangent, surface.material.roughness, &rng);
55        wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N;
56        let pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, surface.material.roughness);
57
58        // https://d1qx31qr3h6wln.cloudfront.net/publications/mueller21realtime.pdf#subsection.3.4, equation (4)
59        let cos_theta = saturate(dot(wo, surface.world_normal));
60        var a0 = dot(wo_unnormalized, wo_unnormalized) / (4.0 * PI * cos_theta);
61        a0 *= TERMINATE_IN_WORLD_CACHE_THRESHOLD;
62
63        radiance = trace_glossy_path(surface.world_position, wi, surface.material.roughness, pdf, a0, &rng) / pdf;
64    }
65
66    let brdf = evaluate_specular_brdf(surface.world_normal, wo, wi, surface.material.base_color, surface.material.metallic,
67        surface.material.reflectance, surface.material.perceptual_roughness, surface.material.roughness);
68    let cos_theta = saturate(dot(wi, surface.world_normal));
69    radiance *= brdf * cos_theta * view.exposure;
70
71    var pixel_color = textureLoad(view_output, global_id.xy);
72    pixel_color += vec4(radiance, 0.0);
73    textureStore(view_output, global_id.xy, pixel_color);
74
75#ifdef VISUALIZE_WORLD_CACHE
76    textureStore(view_output, global_id.xy, vec4(query_world_cache(surface.world_position, surface.world_normal, view.world_position, WORLD_CACHE_CELL_LIFETIME, &rng) * view.exposure, 1.0));
77#endif
78}
79
80fn trace_glossy_path(initial_ray_origin: vec3<f32>, initial_wi: vec3<f32>, initial_roughness: f32, initial_p_bounce: f32, a0: f32, rng: ptr<function, u32>) -> vec3<f32> {
81    var ray_origin = initial_ray_origin;
82    var wi = initial_wi;
83    var p_bounce = initial_p_bounce;
84    var surface_perfectly_specular = false;
85    var path_spread = 0.0;
86
87    // Trace up to three bounces, getting the net throughput from them
88    var radiance = vec3(0.0);
89    var throughput = vec3(1.0);
90    for (var i = 0u; i < 3u; i += 1u) {
91        // Trace ray
92        let ray = trace_ray(ray_origin, wi, RAY_T_MIN, RAY_T_MAX, RAY_FLAG_NONE);
93        if ray.kind == RAY_QUERY_INTERSECTION_NONE { break; }
94        let ray_hit = resolve_ray_hit_full(ray);
95
96        let TBN = calculate_tbn_mikktspace(ray_hit.world_normal, ray_hit.world_tangent);
97        let T = TBN[0];
98        let B = TBN[1];
99        let N = TBN[2];
100
101        let wo = -wi;
102        let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
103
104        // Add emissive contribution
105        let mis_weight = emissive_mis_weight(i, initial_roughness, p_bounce, ray_hit, surface_perfectly_specular);
106        radiance += throughput * mis_weight * ray_hit.material.emissive;
107
108        // Should not perform NEE for mirror-like surfaces
109        surface_perfectly_specular = ray_hit.material.roughness <= 0.001 && ray_hit.material.metallic > 0.9999;
110
111        // https://d1qx31qr3h6wln.cloudfront.net/publications/mueller21realtime.pdf#subsection.3.4, equation (3)
112        path_spread += sqrt((ray.t * ray.t) / (p_bounce * wo_tangent.z));
113
114        if path_spread * path_spread > a0 * get_cell_size(ray_hit.world_position, view.world_position) {
115            // Path spread is wide enough, terminate path in the world cache
116            let diffuse_brdf = ray_hit.material.base_color / PI;
117            radiance += throughput * diffuse_brdf * query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position, WORLD_CACHE_CELL_LIFETIME, rng);
118            break;
119        } else if !surface_perfectly_specular {
120            // Sample direct lighting (NEE)
121            let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, rng);
122            let direct_lighting_brdf = evaluate_brdf(ray_hit.world_normal, wo, direct_lighting.wi, ray_hit.material);
123            let mis_weight = nee_mis_weight(direct_lighting.inverse_pdf, direct_lighting.brdf_rays_can_hit, wo_tangent, direct_lighting.wi, ray_hit, TBN);
124            radiance += throughput * mis_weight * direct_lighting.radiance * direct_lighting.inverse_pdf * direct_lighting_brdf;
125        }
126
127        // Sample new ray direction from the GGX BRDF for next bounce
128        let wi_tangent = sample_ggx_vndf(wo_tangent, ray_hit.material.roughness, rng);
129        wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N;
130        ray_origin = ray_hit.world_position;
131
132        // Update throughput for next bounce
133        p_bounce = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness);
134        let brdf = evaluate_brdf(N, wo, wi, ray_hit.material);
135        let cos_theta = saturate(dot(wi, N));
136        throughput *= (brdf * cos_theta) / p_bounce;
137    }
138
139    return radiance;
140}
141
142fn emissive_mis_weight(i: u32, initial_roughness: f32, p_bounce: f32, ray_hit: ResolvedRayHitFull, previous_surface_perfectly_specular: bool) -> f32 {
143    if i != 0u {
144        if previous_surface_perfectly_specular { return 1.0; }
145
146        let p_light = random_emissive_light_pdf(ray_hit);
147        return power_heuristic(p_bounce, p_light);
148    } else {
149        // The first bounce gets MIS weight 0.0 or 1.0 depending on if ReSTIR DI shaded using the specular lobe or not
150        if initial_roughness <= SPECULAR_GI_FOR_DI_ROUGHNESS_THRESHOLD {
151            return 1.0;
152        } else {
153            return 0.0;
154        }
155    }
156}
157
158fn nee_mis_weight(inverse_p_light: f32, brdf_rays_can_hit: bool, wo_tangent: vec3<f32>, wi: vec3<f32>, ray_hit: ResolvedRayHitFull, TBN: mat3x3<f32>) -> f32 {
159    if !brdf_rays_can_hit {
160        return 1.0;
161    }
162
163    let T = TBN[0];
164    let B = TBN[1];
165    let N = TBN[2];
166    let wi_tangent = vec3(dot(wi, T), dot(wi, B), dot(wi, N));
167
168    let p_light = 1.0 / inverse_p_light;
169    let p_bounce = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness);
170    return power_heuristic(p_light, p_bounce);
171}
172
173// Don't adjust the size of this struct without also adjusting GI_RESERVOIR_STRUCT_SIZE.
174struct Reservoir {
175    sample_point_world_position: vec3<f32>,
176    weight_sum: f32,
177    radiance: vec3<f32>,
178    confidence_weight: f32,
179    sample_point_world_normal: vec3<f32>,
180    unbiased_contribution_weight: f32,
181}