|
| 1 | +#import bevy_pbr::pbr_functions::calculate_tbn_mikktspace |
| 2 | +#import bevy_render::maths::{orthonormalize, PI} |
| 3 | +#import bevy_render::view::View |
| 4 | +#import bevy_solari::brdf::{evaluate_brdf, evaluate_specular_brdf} |
| 5 | +#import bevy_solari::gbuffer_utils::gpixel_resolve |
| 6 | +#import bevy_solari::sampling::{sample_ggx_vndf, ggx_vndf_pdf} |
| 7 | +#import bevy_solari::scene_bindings::{trace_ray, resolve_ray_hit_full, RAY_T_MIN, RAY_T_MAX} |
| 8 | +#import bevy_solari::world_cache::query_world_cache |
| 9 | + |
| 10 | +@group(1) @binding(0) var view_output: texture_storage_2d<rgba16float, read_write>; |
| 11 | +@group(1) @binding(5) var<storage, read_write> gi_reservoirs_a: array<Reservoir>; |
| 12 | +@group(1) @binding(7) var gbuffer: texture_2d<u32>; |
| 13 | +@group(1) @binding(8) var depth_buffer: texture_depth_2d; |
| 14 | +@group(1) @binding(12) var<uniform> view: View; |
| 15 | +struct PushConstants { frame_index: u32, reset: u32 } |
| 16 | +var<push_constant> constants: PushConstants; |
| 17 | + |
| 18 | +@compute @workgroup_size(8, 8, 1) |
| 19 | +fn specular_gi(@builtin(global_invocation_id) global_id: vec3<u32>) { |
| 20 | + if any(global_id.xy >= vec2u(view.main_pass_viewport.zw)) { return; } |
| 21 | + |
| 22 | + let pixel_index = global_id.x + global_id.y * u32(view.main_pass_viewport.z); |
| 23 | + var rng = pixel_index + constants.frame_index; |
| 24 | + |
| 25 | + let depth = textureLoad(depth_buffer, global_id.xy, 0); |
| 26 | + if depth == 0.0 { |
| 27 | + return; |
| 28 | + } |
| 29 | + let surface = gpixel_resolve(textureLoad(gbuffer, global_id.xy, 0), depth, global_id.xy, view.main_pass_viewport.zw, view.world_from_clip); |
| 30 | + |
| 31 | + let wo = normalize(view.world_position - surface.world_position); |
| 32 | + |
| 33 | + var radiance: vec3<f32>; |
| 34 | + var wi: vec3<f32>; |
| 35 | + if surface.material.roughness > 0.04 { |
| 36 | + // Surface is very rough, reuse the ReSTIR GI reservoir |
| 37 | + let gi_reservoir = gi_reservoirs_a[pixel_index]; |
| 38 | + wi = normalize(gi_reservoir.sample_point_world_position - surface.world_position); |
| 39 | + radiance = gi_reservoir.radiance * gi_reservoir.unbiased_contribution_weight; |
| 40 | + } else { |
| 41 | + // Surface is glossy or mirror-like, trace a new path |
| 42 | + let TBN = orthonormalize(surface.world_normal); |
| 43 | + let T = TBN[0]; |
| 44 | + let B = TBN[1]; |
| 45 | + let N = TBN[2]; |
| 46 | + let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N)); |
| 47 | + let wi_tangent = sample_ggx_vndf(wo_tangent, surface.material.roughness, &rng); |
| 48 | + wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N; |
| 49 | + let pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, surface.material.roughness); |
| 50 | + |
| 51 | + radiance = trace_glossy_path(surface.world_position, wi, &rng) / pdf; |
| 52 | + } |
| 53 | + |
| 54 | + let brdf = evaluate_specular_brdf(surface.world_normal, wo, wi, surface.material.base_color, surface.material.metallic, |
| 55 | + surface.material.reflectance, surface.material.perceptual_roughness, surface.material.roughness); |
| 56 | + let cos_theta = saturate(dot(wi, surface.world_normal)); |
| 57 | + radiance *= brdf * cos_theta * view.exposure; |
| 58 | + |
| 59 | + var pixel_color = textureLoad(view_output, global_id.xy); |
| 60 | + pixel_color += vec4(radiance, 0.0); |
| 61 | + textureStore(view_output, global_id.xy, pixel_color); |
| 62 | +} |
| 63 | + |
| 64 | +fn trace_glossy_path(initial_ray_origin: vec3<f32>, initial_wi: vec3<f32>, rng: ptr<function, u32>) -> vec3<f32> { |
| 65 | + var ray_origin = initial_ray_origin; |
| 66 | + var wi = initial_wi; |
| 67 | + |
| 68 | + // Trace up to three bounces, getting the net throughput from them |
| 69 | + var throughput = vec3(1.0); |
| 70 | + for (var i = 0u; i < 3u; i += 1u) { |
| 71 | + // Trace ray |
| 72 | + let ray = trace_ray(ray_origin, wi, RAY_T_MIN, RAY_T_MAX, RAY_FLAG_NONE); |
| 73 | + if ray.kind == RAY_QUERY_INTERSECTION_NONE { break; } |
| 74 | + let ray_hit = resolve_ray_hit_full(ray); |
| 75 | + |
| 76 | + // Surface is very rough, terminate path in the world cache |
| 77 | + if ray_hit.material.roughness > 0.04 || i == 2u { |
| 78 | + let diffuse_brdf = ray_hit.material.base_color / PI; |
| 79 | + return throughput * diffuse_brdf * query_world_cache(ray_hit.world_position, ray_hit.geometric_world_normal, view.world_position); |
| 80 | + } |
| 81 | + |
| 82 | + // Sample new ray direction from the GGX BRDF for next bounce |
| 83 | + let TBN = calculate_tbn_mikktspace(ray_hit.world_normal, ray_hit.world_tangent); |
| 84 | + let T = TBN[0]; |
| 85 | + let B = TBN[1]; |
| 86 | + let N = TBN[2]; |
| 87 | + let wo = -wi; |
| 88 | + let wo_tangent = vec3(dot(wo, T), dot(wo, B), dot(wo, N)); |
| 89 | + let wi_tangent = sample_ggx_vndf(wo_tangent, ray_hit.material.roughness, rng); |
| 90 | + wi = wi_tangent.x * T + wi_tangent.y * B + wi_tangent.z * N; |
| 91 | + ray_origin = ray_hit.world_position; |
| 92 | + |
| 93 | + // Update throughput for next bounce |
| 94 | + let pdf = ggx_vndf_pdf(wo_tangent, wi_tangent, ray_hit.material.roughness); |
| 95 | + let brdf = evaluate_brdf(N, wo, wi, ray_hit.material); |
| 96 | + let cos_theta = dot(wi, N); |
| 97 | + throughput *= (brdf * cos_theta) / pdf; |
| 98 | + } |
| 99 | + |
| 100 | + return vec3(0.0); |
| 101 | +} |
| 102 | + |
| 103 | +// Don't adjust the size of this struct without also adjusting GI_RESERVOIR_STRUCT_SIZE. |
| 104 | +struct Reservoir { |
| 105 | + sample_point_world_position: vec3<f32>, |
| 106 | + weight_sum: f32, |
| 107 | + radiance: vec3<f32>, |
| 108 | + confidence_weight: f32, |
| 109 | + sample_point_world_normal: vec3<f32>, |
| 110 | + unbiased_contribution_weight: f32, |
| 111 | +} |
0 commit comments