1enable wgpu_ray_query;
2
3#define_import_path bevy_solari::sampling
4
5#import bevy_pbr::lighting::D_GGX
6#import bevy_pbr::utils::{rand_f, rand_vec2f, rand_u, rand_range_u}
7#import bevy_render::maths::{PI_2, orthonormalize}
8#import bevy_solari::scene_bindings::{trace_ray, RAY_T_MIN, RAY_T_MAX, light_sources, directional_lights, LightSource, LIGHT_SOURCE_KIND_DIRECTIONAL, resolve_triangle_data_full, ResolvedRayHitFull, MIRROR_ROUGHNESS_THRESHOLD}
9
10fn power_heuristic(f: f32, g: f32) -> f32 {
11    return balance_heuristic(f * f, g * g);
12}
13
14fn balance_heuristic(f: f32, g: f32) -> f32 {
15    // Need to guard against NaNs since ReSTIR reservoirs can have UCW=0
16    if f == 0.0 {
17        return 0.0;
18    }
19    return max(0.0, 1.0 / (1.0 + (g / f)));
20}
21
22// https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (Listing 1)
23// Result is invalid when output.z <= 0.0, and must be discarded
24fn sample_ggx_vndf(wi_tangent: vec3<f32>, roughness: f32, rng: ptr<function, u32>) -> vec3<f32> {
25    // Mirror BRDF case
26    if roughness <= MIRROR_ROUGHNESS_THRESHOLD {
27        return vec3(-wi_tangent.xy, wi_tangent.z);
28    }
29
30    let i = wi_tangent;
31    let rand = rand_vec2f(rng);
32    let i_std = normalize(vec3(i.xy * roughness, i.z));
33    let phi = PI_2 * rand.x;
34    let a = roughness;
35    let s = 1.0 + length(vec2(i.xy));
36    let a2 = a * a;
37    let s2 = s * s;
38    let k = (1.0 - a2) * s2 / (s2 + a2 * i.z * i.z);
39    let b = select(i_std.z, k * i_std.z, i.z > 0.0);
40    let z = fma(1.0 - rand.y, 1.0 + b, -b);
41    let sin_theta = sqrt(saturate(1.0 - z * z));
42    let o_std = vec3(sin_theta * cos(phi), sin_theta * sin(phi), z);
43    let m_std = i_std + o_std;
44    let m = normalize(vec3(m_std.xy * roughness, m_std.z));
45    return 2.0 * dot(i, m) * m - i;
46}
47
48fn ggx_vndf_sample_invalid(ray_tangent: vec3<f32>) -> bool {
49    return !(ray_tangent.z > 0.0);
50}
51
52// https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf (Listing 2)
53fn ggx_vndf_pdf(wi_tangent: vec3<f32>, wo_tangent: vec3<f32>, roughness: f32) -> f32 {
54    // Mirror BRDF case
55    if roughness <= MIRROR_ROUGHNESS_THRESHOLD {
56        let mirror_wo = vec3(-wi_tangent.xy, wi_tangent.z);
57        if all(abs(mirror_wo - wo_tangent) < vec3(0.0001)) {
58            return bitcast<f32>(0x7F800000u); // INF
59        } else {
60            return 0.0;
61        }
62    }
63
64    let i = wi_tangent;
65    let o = wo_tangent;
66    let m = normalize(i + o);
67    let ndf = D_GGX(roughness, saturate(m.z));
68    let ai = roughness * i.xy;
69    let len2 = dot(ai, ai);
70    let t = sqrt(len2 + i.z * i.z);
71    var pdf: f32;
72    if i.z >= 0.0 {
73        let a = roughness;
74        let s = 1.0 + length(i.xy);
75        let a2 = a * a;
76        let s2 = s * s;
77        let k = (1.0 - a2) * s2 / (s2 + a2 * i.z * i.z);
78        pdf = ndf / (2.0 * (k * i.z + t));
79    } else {
80        pdf = ndf * (t - i.z) / (2.0 * len2);
81    }
82
83    return select(pdf, 0.0, isnan(pdf));
84}
85
86fn isnan(x: f32) -> bool {
87    return (bitcast<u32>(x) & 0x7fffffffu) > 0x7f800000u;
88}
89
90const NULL_LIGHT_ID = 0xFFFFFFFFu;
91
92struct LightSample {
93    light_id: u32,
94    seed: u32,
95}
96
97struct ResolvedLightSample {
98    world_position: vec4<f32>,
99    world_normal: vec3<f32>,
100    radiance: vec3<f32>,
101    inverse_pdf: f32,
102}
103
104struct LightContribution {
105    radiance: vec3<f32>,
106    inverse_pdf: f32,
107    wi: vec3<f32>,
108    brdf_rays_can_hit: bool,
109}
110
111struct LightContributionNoPdf {
112    radiance: vec3<f32>,
113    wi: vec3<f32>,
114}
115
116struct GenerateRandomLightSampleResult {
117    light_sample: LightSample,
118    resolved_light_sample: ResolvedLightSample,
119}
120
121fn sample_random_light(ray_origin: vec3<f32>, origin_world_normal: vec3<f32>, rng: ptr<function, u32>) -> LightContribution {
122    let sample = generate_random_light_sample(rng);
123    var light_contribution = calculate_resolved_light_contribution(sample.resolved_light_sample, ray_origin, origin_world_normal);
124    light_contribution.radiance *= trace_light_visibility(ray_origin, sample.resolved_light_sample.world_position);
125    return light_contribution;
126}
127
128fn random_emissive_light_pdf(hit: ResolvedRayHitFull) -> f32 {
129    let light_count = arrayLength(&light_sources);
130    return 1.0 / (f32(light_count) * f32(hit.triangle_count) * hit.triangle_area);
131}
132
133fn generate_random_light_sample(rng: ptr<function, u32>) -> GenerateRandomLightSampleResult {
134    let light_count = arrayLength(&light_sources);
135    let light_id = rand_range_u(light_count, rng);
136
137    let light_source = light_sources[light_id];
138
139    var triangle_id = 0u;
140    if light_source.kind != LIGHT_SOURCE_KIND_DIRECTIONAL {
141        let triangle_count = light_source.kind >> 1u;
142        triangle_id = rand_range_u(triangle_count, rng);
143    }
144
145    let seed = rand_u(rng);
146    let light_sample = LightSample((light_id << 16u) | triangle_id, seed);
147
148    var resolved_light_sample = resolve_light_sample(light_sample, light_source);
149    resolved_light_sample.inverse_pdf *= f32(light_count);
150
151    return GenerateRandomLightSampleResult(light_sample, resolved_light_sample);
152}
153
154fn resolve_light_sample(light_sample: LightSample, light_source: LightSource) -> ResolvedLightSample {
155    if light_source.kind == LIGHT_SOURCE_KIND_DIRECTIONAL {
156        let directional_light = directional_lights[light_source.id];
157
158#ifndef NO_DIRECTIONAL_LIGHT_SOFT_SHADOWS
159        // Sample a random direction within a cone whose base is the sun approximated as a disk
160        // https://www.realtimerendering.com/raytracinggems/unofficial_RayTracingGems_v1.9.pdf#0004286901.INDD%3ASec30%3A305
161        var rng = light_sample.seed;
162        let random = rand_vec2f(&rng);
163        let cos_theta = (1.0 - random.x) + random.x * directional_light.cos_theta_max;
164        let sin_theta = sqrt(1.0 - cos_theta * cos_theta);
165        let phi = random.y * PI_2;
166        let x = cos(phi) * sin_theta;
167        let y = sin(phi) * sin_theta;
168        var direction_to_light = vec3(x, y, cos_theta);
169
170        // Rotate the ray so that the cone it was sampled from is aligned with the light direction
171        direction_to_light = orthonormalize(directional_light.direction_to_light) * direction_to_light;
172#else
173        let direction_to_light = directional_light.direction_to_light;
174#endif
175
176        return ResolvedLightSample(
177        vec4(direction_to_light, 0.0),
178        -direction_to_light,
179        directional_light.luminance,
180        directional_light.inverse_pdf,
181    );
182} else {
183        let triangle_count = light_source.kind >> 1u;
184        let triangle_id = light_sample.light_id & 0xFFFFu;
185        let barycentrics = triangle_barycentrics(light_sample.seed);
186        let triangle_data = resolve_triangle_data_full(light_source.id, triangle_id, barycentrics);
187
188        return ResolvedLightSample(
189            vec4(triangle_data.world_position, 1.0),
190            triangle_data.world_normal,
191            triangle_data.material.emissive.rgb,
192            f32(triangle_count) * triangle_data.triangle_area,
193        );
194    }
195}
196
197fn calculate_resolved_light_contribution(resolved_light_sample: ResolvedLightSample, ray_origin: vec3<f32>, origin_world_normal: vec3<f32>) -> LightContribution {
198    let ray = resolved_light_sample.world_position.xyz - (resolved_light_sample.world_position.w * ray_origin);
199    let light_distance = length(ray);
200    let wi = ray / light_distance;
201
202    let cos_theta_light = saturate(dot(-wi, resolved_light_sample.world_normal));
203    let light_distance_squared = light_distance * light_distance;
204
205    let radiance = resolved_light_sample.radiance * (cos_theta_light / light_distance_squared);
206
207    return LightContribution(radiance, resolved_light_sample.inverse_pdf, wi, resolved_light_sample.world_position.w == 1.0);
208}
209
210fn resolve_and_calculate_light_contribution(light_sample: LightSample, ray_origin: vec3<f32>, origin_world_normal: vec3<f32>) -> LightContributionNoPdf {
211    let resolved_light_sample = resolve_light_sample(light_sample, light_sources[light_sample.light_id >> 16u]);
212    let light_contribution = calculate_resolved_light_contribution(resolved_light_sample, ray_origin, origin_world_normal);
213    return LightContributionNoPdf(light_contribution.radiance, light_contribution.wi);
214}
215
216fn trace_light_visibility(ray_origin: vec3<f32>, light_sample_world_position: vec4<f32>) -> f32 {
217    var ray_direction = light_sample_world_position.xyz;
218    var ray_t_max = RAY_T_MAX;
219
220    if light_sample_world_position.w == 1.0 {
221        let ray = ray_direction - ray_origin;
222        let dist = length(ray);
223        ray_direction = ray / dist;
224        ray_t_max = dist - RAY_T_MIN;
225    }
226
227    if ray_t_max < RAY_T_MIN { return 0.0; }
228
229    let ray_hit = trace_ray(ray_origin, ray_direction, RAY_T_MIN, ray_t_max, RAY_FLAG_TERMINATE_ON_FIRST_HIT);
230    return f32(ray_hit.kind == RAY_QUERY_INTERSECTION_NONE);
231}
232
233fn trace_point_visibility(ray_origin: vec3<f32>, point: vec3<f32>) -> f32 {
234    let ray = point - ray_origin;
235    let dist = length(ray);
236    let ray_direction = ray / dist;
237
238    let ray_t_max = dist - RAY_T_MIN;
239    if ray_t_max < RAY_T_MIN { return 0.0; }
240
241    let ray_hit = trace_ray(ray_origin, ray_direction, RAY_T_MIN, ray_t_max, RAY_FLAG_TERMINATE_ON_FIRST_HIT);
242    return f32(ray_hit.kind == RAY_QUERY_INTERSECTION_NONE);
243}
244
245// https://www.realtimerendering.com/raytracinggems/unofficial_RayTracingGems_v1.9.pdf#0004286901.INDD%3ASec22%3A297
246fn triangle_barycentrics(seed: u32) -> vec3<f32> {
247    var rng = seed;
248    var barycentrics = rand_vec2f(&rng);
249    if barycentrics.x + barycentrics.y > 1.0 { barycentrics = 1.0 - barycentrics; }
250    return vec3(1.0 - barycentrics.x - barycentrics.y, barycentrics);
251}