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