1import bevy_core_pipeline::tonemapping::tonemapping_luminance as luminance;
2import bevy_pbr::deferred::types::unpack_24bit_normal;
3import bevy_pbr::render::pbr_functions::{calculate_diffuse_color, calculate_F0, calculate_F0_dielectric};
4import bevy_render::utils::octahedral_decode;
5import package::realtime::gbuffer_utils::gpixel_resolve;
6import package::realtime::bindings::{gbuffer, depth_buffer, load_motion_vector, motion_vectors, previous_view, view, diffuse_albedo, specular_albedo, normal_roughness, dlss_rr_depth, specular_motion_vectors};
7import package::scene::brdf::{evaluate_specular_brdf, F_AB, lobe_reflectances};
8import package::scene::bindings::{MIRROR_ROUGHNESS_THRESHOLD, ResolvedMaterial, ResolvedRayHitFull};
9
10enable wgpu_ray_query;
11
12/// A surface is replaced outright only when nearly all of its reflectance leaves through the delta
13/// lobe.
14const FULL_REPLACEMENT_FRACTION = 0.9;
15
16/// How far a reflector's virtual image may shift between neighboring pixels, as a fraction of the
17/// virtual distance itself, before the surface is treated as too curved to replace.
18const MAX_VIRTUAL_DEPTH_JITTER = 0.002;
19
20/// Even a perfectly flat surface measures a small nonzero normal change, because G-buffer normals are
21/// 24-bit octahedral and neighboring texels round differently.
22const NORMAL_QUANTIZATION_FLOOR = 0.002;
23
24/// Distance, in meters, standing in for "infinitely far" when a mirror ray finds nothing.
25const ENVIRONMENT_VIRTUAL_DISTANCE = 10000.0;
26
27@compute @workgroup_size(8, 8, 1)
28fn resolve_dlss_rr_textures(@builtin(global_invocation_id) global_id: vec3<u32>) {
29    let pixel_id = global_id.xy;
30    if any(pixel_id >= vec2u(view.main_pass_viewport.zw)) { return; }
31
32    let depth = textureLoad(depth_buffer, global_id.xy, 0);
33    let surface_motion_vector = load_motion_vector(pixel_id);
34
35    textureStore(dlss_rr_depth, pixel_id, vec4(depth));
36
37    if depth == 0.0 {
38        textureStore(specular_motion_vectors, pixel_id, vec4(surface_motion_vector, vec2(0.0)));
39        textureStore(diffuse_albedo, pixel_id, vec4(0.0));
40        textureStore(specular_albedo, pixel_id, vec4(0.5));
41        textureStore(normal_roughness, pixel_id, vec4(0.0, 0.0, 1.0, 0.0));
42        return;
43    }
44
45    let surface = gpixel_resolve(textureLoad(gbuffer, pixel_id, 0), depth, pixel_id, view.main_pass_viewport.zw, view.world_from_clip);
46    let F0 = calculate_F0(surface.material.base_color, surface.material.metallic, vec3(surface.material.reflectance));
47    let wo = normalize(view.world_position - surface.world_position);
48
49    textureStore(diffuse_albedo, pixel_id, vec4(calculate_diffuse_color(surface.material.base_color, surface.material.metallic, 0.0, 0.0), 0.0));
50    textureStore(specular_albedo, pixel_id, vec4(env_brdf_approx2(F0, surface.material.roughness, surface.world_normal, wo), 0.0));
51    textureStore(normal_roughness, pixel_id, vec4(surface.world_normal, surface.material.perceptual_roughness));
52
53    // Resetting to the specular motion vector on the frames we draw a diffuse lobe means we
54    // wouldn't be writing the motion vectors that frame, so it would alternate the guide between
55    // to values, and the denoiser blurs the reflection rather than reprojecting it.
56    if !needs_specular_motion_vector(surface.material.roughness) {
57        textureStore(specular_motion_vectors, pixel_id, vec4(surface_motion_vector, vec2(0.0)));
58    }
59}
60
61fn needs_specular_motion_vector(roughness: f32) -> bool {
62    return roughness <= 0.0625;
63}
64
65/// Carried across every bounce of the path walk, so anything derivable from a value the walk already
66/// holds is derived at the end instead of stored. What is left is only what the walk itself produces.
67struct PsrState {
68    finished: bool,
69    /// Whether this pixel's surface is being swapped out for the reflected one, or is only having its
70    /// reflection's motion described. Keyed on the primary surface only, and false for anything glossy.
71    replace_fully: bool,
72    /// The chain's cumulative transform - the product of one Householder reflection per mirror, in
73    /// chain order - stored as a rotor plus the parity of the reflection count rather than as the 3x3
74    /// matrix it denotes.
75    mirror_rotor: vec4<f32>,
76    mirror_parity_odd: bool,
77    /// Distance traveled since the primary surface, camera through every mirror. Not the total path
78    /// length: keeping the reflection's own share means the primary distance is only ever added, never
79    /// subtracted back out, so it can be recovered at the end without the two disagreeing.
80    reflection_length: f32,
81    /// Product of the Fresnel term at each mirror in the chain: what the reflected surface's color is
82    /// multiplied by on its way back to the eye.
83    throughput: vec3<f32>,
84}
85
86fn psr_init(world_normal: vec3<f32>, material: ResolvedMaterial) -> PsrState {
87    var psr: PsrState;
88    psr.finished = !needs_specular_motion_vector(material.roughness);
89    psr.replace_fully = is_delta_mirror(material);
90    // The chain opens with the primary surface's own reflection, so parity starts odd.
91    psr.mirror_rotor = mirror_chain_rotor(world_normal, false);
92    psr.mirror_parity_odd = true;
93    psr.reflection_length = 0.0;
94    psr.throughput = vec3(1.0);
95    return psr;
96}
97
98/// Fold one more mirror into the chain, equivalent to right-multiplying the chain's matrix by this
99/// mirror's Householder reflection.
100/// https://en.wikipedia.org/wiki/Householder_transformation
101///
102/// A product of k Householder reflections is orthogonal with determinant (-1)^k, so it is a rotation
103/// when k is even, and a rotation composed with any one fixed reflection when k is odd.
104fn psr_chain_mirror(psr: PsrState, plane_normal: vec3<f32>) -> PsrState {
105    var chained = psr;
106    chained.mirror_rotor = quaternion_multiply(psr.mirror_rotor, mirror_chain_rotor(plane_normal, psr.mirror_parity_odd));
107    chained.mirror_parity_odd = !psr.mirror_parity_odd;
108    return chained;
109}
110
111fn psr_environment_miss(pixel_id: vec2<u32>, primary_wo: vec3<f32>) {
112    let far_position = view.world_position - (primary_wo * ENVIRONMENT_VIRTUAL_DISTANCE);
113    textureStore(specular_motion_vectors, pixel_id, vec4(calculate_motion_vector(far_position, far_position), vec2(0.0)));
114}
115
116/// Primary surface replacement for perfect mirrors. The reflection chain has been followed to its
117/// first non-mirror hit; write that surface's attributes, reflected into the mirror's virtual space,
118/// to the DLSS RR guide buffers so the denoiser treats this pixel as directly seeing it.
119/// https://developer.nvidia.com/blog/rendering-perfect-reflections-and-refractions-in-path-traced-games/#primary_surface_replacement
120fn replace_primary_surface(
121    pixel_id: vec2<u32>,
122    psr: PsrState,
123    ray_hit: ResolvedRayHitFull,
124    primary_material: ResolvedMaterial,
125    primary_world_position: vec3<f32>,
126    primary_normal: vec3<f32>,
127    primary_wo: vec3<f32>,
128    primary_F_ab: vec2<f32>,
129) {
130    let primary_distance = length(primary_world_position - view.world_position);
131    let path_length = primary_distance + psr.reflection_length;
132
133    // Too curved to track the reflection; fall back to the surface's own motion so the guide is not stale
134    if psr.replace_fully && 2.0 * reflector_curvature(pixel_id, primary_normal) * psr.reflection_length > MAX_VIRTUAL_DEPTH_JITTER * path_length {
135        textureStore(specular_motion_vectors, pixel_id, vec4(load_motion_vector(pixel_id), vec2(0.0)));
136        return;
137    }
138
139    let virtual_position = primary_world_position - (primary_wo * psr.reflection_length);
140    let world_motion = ray_hit.previous_frame_world_position - ray_hit.world_position;
141    let virtual_previous_position = virtual_position + apply_mirror_chain(psr, world_motion);
142    let specular_motion_vector = calculate_motion_vector(virtual_position, virtual_previous_position);
143
144    textureStore(specular_motion_vectors, pixel_id, vec4(specular_motion_vector, vec2(0.0)));
145
146    if !psr.replace_fully {
147        return;
148    }
149
150    let wo = normalize(view.world_position - virtual_position);
151    let virtual_normal = normalize(apply_mirror_chain(psr, ray_hit.world_normal));
152
153    let primary_depth = textureLoad(depth_buffer, pixel_id, 0);
154    textureStore(dlss_rr_depth, pixel_id, vec4(primary_depth * primary_distance / path_length));
155    @if(MOTION_VECTORS_READ_WRITE)
156    textureStore(motion_vectors, pixel_id, vec4(specular_motion_vector, vec2(0.0)));
157
158    let reflected_diffuse = calculate_diffuse_color(ray_hit.material.base_color, ray_hit.material.metallic, 0.0, 0.0);
159    let reflected_F0 = calculate_F0(ray_hit.material.base_color, ray_hit.material.metallic, vec3(ray_hit.material.reflectance));
160    let reflected_specular = env_brdf_approx2(reflected_F0, ray_hit.material.roughness, virtual_normal, wo);
161
162    let own_diffuse = calculate_diffuse_color(primary_material.base_color, primary_material.metallic, 0.0, 0.0);
163    let own_F0 = calculate_F0(primary_material.base_color, primary_material.metallic, vec3(primary_material.reflectance));
164    let own_specular = env_brdf_approx2(own_F0, primary_material.roughness, primary_normal, primary_wo);
165
166    let specular_fraction = specular_fraction(primary_material, primary_F_ab);
167
168    textureStore(diffuse_albedo, pixel_id, vec4(psr.throughput * reflected_diffuse + own_diffuse, 0.0));
169    textureStore(specular_albedo, pixel_id, vec4(mix(own_specular, psr.throughput * reflected_specular, specular_fraction), 0.0));
170    textureStore(normal_roughness, pixel_id, vec4(virtual_normal, ray_hit.material.perceptual_roughness));
171}
172
173fn primary_surface_glossy(material: ResolvedMaterial) -> bool {
174    return material.roughness > MIRROR_ROUGHNESS_THRESHOLD;
175}
176
177fn is_delta_mirror(material: ResolvedMaterial) -> bool {
178    return material.roughness <= MIRROR_ROUGHNESS_THRESHOLD
179        && specular_fraction(material, F_AB(material.perceptual_roughness, 1.0)) >= FULL_REPLACEMENT_FRACTION;
180}
181
182/// Share of a surface's total reflectance that leaves through the specular lobe, in 0..1.
183fn specular_fraction(material: ResolvedMaterial, F_ab: vec2<f32>) -> f32 {
184    let rho = lobe_reflectances(material.base_color, calculate_F0_dielectric(vec3(material.reflectance)), material, F_ab);
185    let specular_luminance = luminance(rho.specular);
186    return specular_luminance / max(specular_luminance + luminance(rho.diffuse), 0.0001);
187}
188
189fn psr_bounce_reflectance(wo: vec3<f32>, wi: vec3<f32>, normal: vec3<f32>, material: ResolvedMaterial, F_ab: vec2<f32>, delta: bool) -> vec3<f32> {
190    if delta {
191        return evaluate_specular_brdf(wo, wi, normal, material, F_ab);
192    }
193    return lobe_reflectances(material.base_color, calculate_F0_dielectric(vec3(material.reflectance)), material, F_ab).specular;
194}
195
196/// Angle, in radians, that the shading normal changes over one pixel in screen space.
197fn reflector_curvature(pixel_id: vec2<u32>, normal: vec3<f32>) -> f32 {
198    let max_pixel_id = vec2<u32>(view.main_pass_viewport.zw) - vec2(1u);
199    let right = load_gbuffer_normal(min(pixel_id + vec2(1u, 0u), max_pixel_id));
200    let down = load_gbuffer_normal(min(pixel_id + vec2(0u, 1u), max_pixel_id));
201    let min_cos = min(dot(normal, right), dot(normal, down));
202    let angle = sqrt(max(0.0, 2.0 * (1.0 - min_cos)));
203    return max(0.0, angle - NORMAL_QUANTIZATION_FLOOR);
204}
205
206fn load_gbuffer_normal(pixel_id: vec2<u32>) -> vec3<f32> {
207    return octahedral_decode(unpack_24bit_normal(textureLoad(gbuffer, pixel_id, 0).a));
208}
209
210fn apply_mirror_chain(psr: PsrState, v: vec3<f32>) -> vec3<f32> {
211    var reflected = v;
212    if psr.mirror_parity_odd {
213        reflected.z = -reflected.z;
214    }
215    let t = 2.0 * cross(psr.mirror_rotor.xyz, reflected);
216    return reflected + psr.mirror_rotor.w * t + cross(psr.mirror_rotor.xyz, t);
217}
218
219/// The rotation half of `H(n) * F` at even parity, or of `F * H(n)` at odd, where `H(n)` is the
220/// Householder reflection about `plane_normal` and `F` the fixed z-axis one.
221fn mirror_chain_rotor(plane_normal: vec3<f32>, parity_odd: bool) -> vec4<f32> {
222    if parity_odd {
223        return vec4(plane_normal.y, -plane_normal.x, 0.0, plane_normal.z);
224    }
225    return vec4(-plane_normal.y, plane_normal.x, 0.0, plane_normal.z);
226}
227
228fn quaternion_multiply(a: vec4<f32>, b: vec4<f32>) -> vec4<f32> {
229    return vec4(a.w * b.xyz + b.w * a.xyz + cross(a.xyz, b.xyz), a.w * b.w - dot(a.xyz, b.xyz));
230}
231
232fn calculate_motion_vector(world_position: vec3<f32>, previous_world_position: vec3<f32>) -> vec2<f32> {
233    let clip_position_t = view.unjittered_clip_from_world * vec4(world_position, 1.0);
234    let clip_position = clip_position_t.xy / clip_position_t.w;
235    let previous_clip_position_t = previous_view.unjittered_clip_from_world * vec4(previous_world_position, 1.0);
236    let previous_clip_position = previous_clip_position_t.xy / previous_clip_position_t.w;
237    // Motion vectors are UV-space offsets in [-1, 1], from one corner to the diagonally-opposite one.
238    // A clip-space diagonal difference is in [-2, 2], so scale by 0.5, and flip y since V goes down
239    // where clip-space y goes up.
240    return (clip_position - previous_clip_position) * vec2(0.5, -0.5);
241}
242
243fn env_brdf_approx2(specular_color: vec3<f32>, alpha: f32, N: vec3<f32>, V: vec3<f32>) -> vec3<f32> {
244    let NoV = abs(dot(N, V));
245
246    var X: vec4<f32>;
247    X.x = 1.0;
248    X.y = NoV;
249    X.z = NoV * NoV;
250    X.w = NoV * X.z;
251
252    var Y: vec4<f32>;
253    Y.x = 1.0;
254    Y.y = alpha;
255    Y.z = alpha * alpha;
256    Y.w = alpha * Y.z;
257
258    let M1 = mat2x2<f32>(0.99044, 1.29678, -1.28514, -0.755907);
259    let M2 = mat3x3<f32>(1.0, 20.3225, 121.563, 2.92338, -27.0302, 626.13, 59.4188, 222.592, 316.627);
260    let M3 = mat2x2<f32>(0.0365463, 9.0632, 3.32707, -9.04756);
261    let M4 = mat3x3<f32>(1.0, 9.04401, 5.56589, 3.59685, -16.3174, 19.7886, -1.36772, 9.22949, -20.2123);
262
263    var bias = dot(M1 * X.xy, Y.xy) / dot(M2 * X.xyw, Y.xyw);
264    let scale = dot(M3 * X.xy, Y.xy) / dot(M4 * X.xzw, Y.xyw);
265
266    bias *= saturate(specular_color.g * 50.0);
267
268    return fma(specular_color, vec3(max(0.0, scale)), vec3(max(0.0, bias)));
269}