Bevy Shader Explorer
Search across packages, versions, files, declarations, and comments. Use multiple terms to narrow results (AND), and add @vertex, @fragment, or @compute to filter by shader stage.
Packages › bevy_solari › specular_gi.wgsl › source
← Back to declarations 
1enable wgpu_ray_query;
3#define_import_path bevy_solari::specular_gi
5#import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance
6#import bevy_pbr::pbr_functions::{calculate_diffuse_color, calculate_F0}
7#import bevy_pbr::utils::rand_f
8#import bevy_render::maths::{orthonormalize, PI}
9#import bevy_render::view::View
10#import bevy_solari::brdf::{evaluate_brdf, evaluate_specular_brdf}
11#import bevy_solari::gbuffer_utils::{gpixel_resolve, ResolvedGPixel}
12#import bevy_solari::sampling::{sample_random_light, random_emissive_light_pdf, sample_ggx_vndf, ggx_vndf_pdf, ggx_vndf_sample_invalid, power_heuristic}
13#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, ResolvedRayHitFull, RAY_T_MIN, RAY_T_MAX, MIRROR_ROUGHNESS_THRESHOLD}
14#import bevy_solari::world_cache::{query_world_cache, get_cell_size, WORLD_CACHE_CELL_LIFETIME}
15#import bevy_solari::realtime_bindings::{view_output, gi_reservoirs_a, gbuffer, depth_buffer, view, constants}
16#ifdef DLSS_RR_GUIDE_BUFFERS
17#import bevy_solari::realtime_bindings::{diffuse_albedo, specular_albedo, normal_roughness, specular_motion_vectors, previous_view}
18#import bevy_solari::resolve_dlss_rr_textures::env_brdf_approx2
19#endif
21const DIFFUSE_GI_REUSE_ROUGHNESS_THRESHOLD: f32 = 0.4;
22const SPECULAR_GI_FOR_DI_ROUGHNESS_THRESHOLD: f32 = 0.0225;
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; }
28    let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z);
29    var rng = pixel_index + constants.frame_index;
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);
37    let wo_unnormalized = view.world_position - surface.world_position;
38    let wo_length = length(wo_unnormalized);
39    let wo = wo_unnormalized / wo_length;
41    var radiance: vec3<f32>;
42    var wi: vec3<f32>;
43    if surface.material.roughness > DIFFUSE_GI_REUSE_ROUGHNESS_THRESHOLD {
44        // Surface is very rough, reuse the ReSTIR GI reservoir
45        let gi_reservoir = gi_reservoirs_a[pixel_index];
46        wi = normalize(gi_reservoir.sample_point_world_position - surface.world_position);
47        radiance = gi_reservoir.radiance * gi_reservoir.unbiased_contribution_weight;
48    } else {
49        // Surface is glossy or mirror-like, trace a new path
50        let TBN = orthonormalize(surface.world_normal);
51        let T = TBN[0];
52        let B = TBN[1];
53        let N = TBN[2];
54        let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
55        let wi_tangent = sample_ggx_vndf(wo_tangent, surface.material.roughness, &rng);
56        if ggx_vndf_sample_invalid(wi_tangent) {
57            wi = vec3(0.0);
58            radiance = vec3(0.0);
59        } else {
60            wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N;
61            let pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, surface.material.roughness);
63            radiance = trace_glossy_path(global_id.xy, surface, wo_length, wi, pdf, &rng);
64            if surface.material.roughness > MIRROR_ROUGHNESS_THRESHOLD {
65                radiance /= pdf;
66            }
67        }
68    }
70    let brdf = evaluate_specular_brdf(wo, wi, surface.world_normal, surface.material);
71    radiance *= brdf * view.exposure;
73    var pixel_color = textureLoad(view_output, global_id.xy);
74    pixel_color += vec4(radiance, 0.0);
75    textureStore(view_output, global_id.xy, pixel_color);
77#ifdef VISUALIZE_WORLD_CACHE
78    textureStore(view_output, global_id.xy, vec4(query_world_cache(surface.world_position, surface.world_normal, view.world_position, RAY_T_MAX, WORLD_CACHE_CELL_LIFETIME, &rng) * view.exposure, 1.0));
79#endif
82fn trace_glossy_path(pixel_id: vec2<u32>, primary_surface: ResolvedGPixel, initial_ray_t: f32, initial_wi: vec3<f32>, initial_p_bounce: f32, rng: ptr<function, u32>) -> vec3<f32> {
83    var radiance = vec3(0.0);
84    var throughput = vec3(1.0);
86    var ray_origin = primary_surface.world_position + (primary_surface.world_normal * RAY_T_MIN);
87    var wi = initial_wi;
88    var p_bounce = initial_p_bounce;
89    var path_roughness = primary_surface.material.roughness;
91#ifdef DLSS_RR_GUIDE_BUFFERS
92    var mirror_rotations = reflection_matrix(primary_surface.world_normal);
93    var psr_finished = false;
94#endif
96    // Trace up to three bounces
97    for (var i = 0u; i < 3u; i += 1u) {
98        // Trace ray
99        let ray = trace_ray(ray_origin, wi, RAY_T_MIN, RAY_T_MAX, RAY_FLAG_NONE);
100        if ray.kind == RAY_QUERY_INTERSECTION_NONE { break; }
101        let ray_hit = resolve_ray_hit_full(ray);
103        let TBN = orthonormalize(ray_hit.world_normal);
104        let T = TBN[0];
105        let B = TBN[1];
106        let N = TBN[2];
108        let wo = -wi;
109        let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N));
111        // Add emissive contribution
112        let mis_weight = emissive_mis_weight(i, primary_surface.material.roughness, p_bounce, ray_hit);
113        radiance += throughput * mis_weight * ray_hit.material.emissive;
115        // Should not perform NEE for mirror-like surfaces
116        let surface_perfect_mirror = ray_hit.material.roughness <= MIRROR_ROUGHNESS_THRESHOLD && ray_hit.material.metallic > 0.9999;
118        // Primary surface replacement for perfect mirrors
119        // https://developer.nvidia.com/blog/rendering-perfect-reflections-and-refractions-in-path-traced-games/#primary_surface_replacement
120#ifdef DLSS_RR_GUIDE_BUFFERS
121        if !psr_finished && primary_surface.material.roughness <= MIRROR_ROUGHNESS_THRESHOLD && primary_surface.material.metallic > 0.9999 {
122            if surface_perfect_mirror {
123                mirror_rotations = mirror_rotations * reflection_matrix(ray_hit.world_normal);
124            } else {
125                psr_finished = true;
126                replace_primary_surface(pixel_id, ray_hit, mirror_rotations, primary_surface.world_position);
127            }
128        }
129#endif
131        // Terminate path in the world cache if the ray is long enough and the path spread is large enough
132        var rng_copy = *rng;
133        let world_cache_cell_size = get_cell_size(ray_hit.world_position, view.world_position, &rng_copy);
134        let ray_longer_than_cell = ray.t > sqrt(3.0) * world_cache_cell_size;
136        if i == 2u || (ray_longer_than_cell && path_roughness > DIFFUSE_GI_REUSE_ROUGHNESS_THRESHOLD) {
137            let diffuse_brdf = ray_hit.material.base_color / PI;
138            radiance += throughput * diffuse_brdf * query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position, ray.t, WORLD_CACHE_CELL_LIFETIME, rng);
139            break;
140        } else if !surface_perfect_mirror {
141            // Sample direct lighting (NEE)
142            let direct_lighting = sample_random_light(ray_hit.world_position, ray_hit.world_normal, rng);
143            let direct_lighting_brdf = evaluate_brdf(wo, direct_lighting.wi, ray_hit.world_normal, ray_hit.material);
144            let mis_weight = nee_mis_weight(direct_lighting.inverse_pdf, direct_lighting.brdf_rays_can_hit, wo_tangent, direct_lighting.wi, ray_hit, TBN);
145            radiance += throughput * mis_weight * direct_lighting.radiance * direct_lighting.inverse_pdf * direct_lighting_brdf;
146        }
148        // Sample new ray direction from the GGX BRDF for next bounce
149        let wi_tangent = sample_ggx_vndf(wo_tangent, ray_hit.material.roughness, rng);
150        if ggx_vndf_sample_invalid(wi_tangent) { break; }
151        wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N;
152        ray_origin = ray_hit.world_position + (ray_hit.geometric_world_normal * RAY_T_MIN);
154        // Update throughput for next bounce
155        p_bounce = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness);
156        throughput *= evaluate_brdf(wo, wi, N, ray_hit.material);
157        if ray_hit.material.roughness > MIRROR_ROUGHNESS_THRESHOLD {
158            throughput /= p_bounce;
159        }
160        path_roughness += ray_hit.material.roughness;
162        // Russian roulette for early termination
163        let p = luminance(throughput);
164        if rand_f(rng) > p { break; }
165        throughput /= p;
166    }
168    return radiance;
171fn emissive_mis_weight(i: u32, initial_roughness: f32, p_bounce: f32, ray_hit: ResolvedRayHitFull) -> f32 {
172    if i != 0u {
173        let p_light = random_emissive_light_pdf(ray_hit);
174        return power_heuristic(p_bounce, p_light);
175    } else {
176        // The first bounce gets MIS weight 0.0 or 1.0 depending on if ReSTIR DI shaded using the specular lobe or not
177        if initial_roughness <= SPECULAR_GI_FOR_DI_ROUGHNESS_THRESHOLD {
178            return 1.0;
179        } else {
180            return 0.0;
181        }
182    }
185fn 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 {
186    if !brdf_rays_can_hit {
187        return 1.0;
188    }
190    let T = TBN[0];
191    let B = TBN[1];
192    let N = TBN[2];
193    let wi_tangent = vec3(dot(wi, T), dot(wi, B), dot(wi, N));
195    let p_light = 1.0 / inverse_p_light;
196    let p_bounce = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness);
197    return power_heuristic(p_light, p_bounce);
200#ifdef DLSS_RR_GUIDE_BUFFERS
201// https://en.wikipedia.org/wiki/Householder_transformation
202fn reflection_matrix(plane_normal: vec3f) -> mat3x3<f32> {
203    // N times Nᵀ.
204    let n_nt = mat3x3<f32>(
205        plane_normal * plane_normal.x,
206        plane_normal * plane_normal.y,
207        plane_normal * plane_normal.z,
208    );
209    let identity_matrix = mat3x3<f32>(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0);
210    return identity_matrix - n_nt * 2.0;
213fn replace_primary_surface(pixel_id: vec2<u32>, ray_hit: ResolvedRayHitFull, mirror_rotations: mat3x3<f32>, primary_surface_world_position: vec3<f32>) {
214    // Simplification: Apply all rotations in the chain around the first mirror, rather than applying each rotation around its respective mirror
215    let virtual_position = (mirror_rotations * (ray_hit.world_position - primary_surface_world_position)) + primary_surface_world_position;
216    let virtual_previous_frame_position = (mirror_rotations * (ray_hit.previous_frame_world_position - primary_surface_world_position)) + primary_surface_world_position;
217    let specular_motion_vector = calculate_motion_vector(virtual_position, virtual_previous_frame_position);
219    let F0 = calculate_F0(ray_hit.material.base_color, ray_hit.material.metallic, vec3(ray_hit.material.reflectance));
220    let wo = normalize(view.world_position - virtual_position);
221    let virtual_normal = normalize(mirror_rotations * ray_hit.world_normal);
223    textureStore(specular_motion_vectors, pixel_id, vec4(specular_motion_vector, vec2(0.0)));
224    textureStore(diffuse_albedo, pixel_id, vec4(calculate_diffuse_color(ray_hit.material.base_color, ray_hit.material.metallic, 0.0, 0.0), 0.0));
225    textureStore(specular_albedo, pixel_id, vec4(env_brdf_approx2(F0, ray_hit.material.roughness, virtual_normal, wo), 0.0));
226    textureStore(normal_roughness, pixel_id, vec4(virtual_normal, ray_hit.material.perceptual_roughness));
229fn calculate_motion_vector(world_position: vec3<f32>, previous_world_position: vec3<f32>) -> vec2<f32> {
230    let clip_position_t = view.unjittered_clip_from_world * vec4(world_position, 1.0);
231    let clip_position = clip_position_t.xy / clip_position_t.w;
232    let previous_clip_position_t = previous_view.clip_from_world * vec4(previous_world_position, 1.0);
233    let previous_clip_position = previous_clip_position_t.xy / previous_clip_position_t.w;
234    // These motion vectors are used as offsets to UV positions and are stored
235    // in the range -1,1 to allow offsetting from the one corner to the
236    // diagonally-opposite corner in UV coordinates, in either direction.
237    // A difference between diagonally-opposite corners of clip space is in the
238    // range -2,2, so this needs to be scaled by 0.5. And the V direction goes
239    // down where clip space y goes up, so y needs to be flipped.
240    return (clip_position - previous_clip_position) * vec2(0.5, -0.5);
242#endif
Contents