1enable wgpu_ray_query;
2#define_import_path bevy_solari::resolve_dlss_rr_textures
3
4#import bevy_pbr::pbr_functions::{calculate_diffuse_color, calculate_F0}
5#import bevy_render::view::View
6#import bevy_solari::gbuffer_utils::gpixel_resolve
7#import bevy_solari::realtime_bindings::{gbuffer, depth_buffer, motion_vectors, view, diffuse_albedo, specular_albedo, normal_roughness, specular_motion_vectors}
8
9@compute @workgroup_size(8, 8, 1)
10fn resolve_dlss_rr_textures(@builtin(global_invocation_id) global_id: vec3<u32>) {
11    let pixel_id = global_id.xy;
12    if any(pixel_id >= vec2u(view.main_pass_viewport.zw)) { return; }
13
14    textureStore(specular_motion_vectors, pixel_id, textureLoad(motion_vectors, pixel_id, 0));
15
16    let depth = textureLoad(depth_buffer, global_id.xy, 0);
17    if depth == 0.0 {
18        textureStore(diffuse_albedo, pixel_id, vec4(0.0));
19        textureStore(specular_albedo, pixel_id, vec4(0.5));
20        textureStore(normal_roughness, pixel_id, vec4(0.0, 0.0, 1.0, 0.0));
21        return;
22    }
23
24    let surface = gpixel_resolve(textureLoad(gbuffer, pixel_id, 0), depth, pixel_id, view.main_pass_viewport.zw, view.world_from_clip);
25    let F0 = calculate_F0(surface.material.base_color, surface.material.metallic, vec3(surface.material.reflectance));
26    let wo = normalize(view.world_position - surface.world_position);
27
28    textureStore(diffuse_albedo, pixel_id, vec4(calculate_diffuse_color(surface.material.base_color, surface.material.metallic, 0.0, 0.0), 0.0));
29    textureStore(specular_albedo, pixel_id, vec4(env_brdf_approx2(F0, surface.material.roughness, surface.world_normal, wo), 0.0));
30    textureStore(normal_roughness, pixel_id, vec4(surface.world_normal, surface.material.perceptual_roughness));
31}
32
33fn env_brdf_approx2(specular_color: vec3<f32>, alpha: f32, N: vec3<f32>, V: vec3<f32>) -> vec3<f32> {
34    let NoV = abs(dot(N, V));
35
36    var X: vec4<f32>;
37    X.x = 1.0;
38    X.y = NoV;
39    X.z = NoV * NoV;
40    X.w = NoV * X.z;
41
42    var Y: vec4<f32>;
43    Y.x = 1.0;
44    Y.y = alpha;
45    Y.z = alpha * alpha;
46    Y.w = alpha * Y.z;
47
48    let M1 = mat2x2<f32>(0.99044, 1.29678, -1.28514, -0.755907);
49    let M2 = mat3x3<f32>(1.0, 20.3225, 121.563, 2.92338, -27.0302, 626.13, 59.4188, 222.592, 316.627);
50    let M3 = mat2x2<f32>(0.0365463, 9.0632, 3.32707, -9.04756);
51    let M4 = mat3x3<f32>(1.0, 9.04401, 5.56589, 3.59685, -16.3174, 19.7886, -1.36772, 9.22949, -20.2123);
52
53    var bias = dot(M1 * X.xy, Y.xy) / dot(M2 * X.xyw, Y.xyw);
54    let scale = dot(M3 * X.xy, Y.xy) / dot(M4 * X.xzw, Y.xyw);
55
56    bias *= saturate(specular_color.g * 50.0);
57
58    return fma(specular_color, vec3(max(0.0, scale)), vec3(max(0.0, bias)));
59}