1// Visibility Bitmask Ambient Occlusion (VBAO)
2// Paper: ttps://ar5iv.labs.arxiv.org/html/2301.11376
3
4// Source code heavily based on XeGTAO v1.30 from Intel
5// https://github.com/GameTechDev/XeGTAO/blob/0d177ce06bfa642f64d8af4de1197ad1bcb862d4/Source/Rendering/Shaders/XeGTAO.hlsli
6
7// Source code based on the existing XeGTAO implementation and
8// https://cdrinmatane.github.io/posts/ssaovb-code/
9
10// Source code base on SSRT3 implementation
11// https://github.com/cdrinmatane/SSRT3
12
13#import bevy_render::maths::fast_acos
14
15#import bevy_render::{
16    view::View,
17    globals::Globals,
18    maths::{PI, HALF_PI},
19}
20
21@group(0) @binding(0) var preprocessed_depth: texture_2d<f32>;
22@group(0) @binding(1) var normals: texture_2d<f32>;
23@group(0) @binding(2) var hilbert_index_lut: texture_2d<u32>;
24#ifdef USE_R16FLOAT
25@group(0) @binding(3) var ambient_occlusion: texture_storage_2d<r16float, write>;
26#else
27@group(0) @binding(3) var ambient_occlusion: texture_storage_2d<r32float, write>;
28#endif
29@group(0) @binding(4) var depth_differences: texture_storage_2d<r32uint, write>;
30@group(0) @binding(5) var<uniform> globals: Globals;
31@group(0) @binding(6) var<uniform> thickness: f32;
32@group(1) @binding(0) var point_clamp_sampler: sampler;
33@group(1) @binding(1) var linear_clamp_sampler: sampler;
34@group(1) @binding(2) var<uniform> view: View;
35
36fn load_noise(pixel_coordinates: vec2<i32>) -> vec2<f32> {
37    var index = textureLoad(hilbert_index_lut, pixel_coordinates % 64, 0).r;
38
39#ifdef TEMPORAL_JITTER
40    index += 288u * (globals.frame_count % 64u);
41#endif
42
43    // R2 sequence - http://extremelearning.com.au/unreasonable-effectiveness-of-quasirandom-sequences
44    return fract(0.5 + f32(index) * vec2<f32>(0.75487766624669276005, 0.5698402909980532659114));
45}
46
47// Calculate differences in depth between neighbor pixels (later used by the spatial denoiser pass to preserve object edges)
48fn calculate_neighboring_depth_differences(pixel_coordinates: vec2<i32>) -> f32 {
49    // Sample the pixel's depth and 4 depths around it
50    let uv = vec2<f32>(pixel_coordinates) / view.viewport.zw;
51    let depths_upper_left = textureGather(0, preprocessed_depth, point_clamp_sampler, uv);
52    let depths_bottom_right = textureGather(0, preprocessed_depth, point_clamp_sampler, uv, vec2<i32>(1i, 1i));
53    let depth_center = depths_upper_left.y;
54    let depth_left = depths_upper_left.x;
55    let depth_top = depths_upper_left.z;
56    let depth_bottom = depths_bottom_right.x;
57    let depth_right = depths_bottom_right.z;
58
59    // Calculate the depth differences (large differences represent object edges)
60    var edge_info = vec4<f32>(depth_left, depth_right, depth_top, depth_bottom) - depth_center;
61    let slope_left_right = (edge_info.y - edge_info.x) * 0.5;
62    let slope_top_bottom = (edge_info.w - edge_info.z) * 0.5;
63    let edge_info_slope_adjusted = edge_info + vec4<f32>(slope_left_right, -slope_left_right, slope_top_bottom, -slope_top_bottom);
64    edge_info = min(abs(edge_info), abs(edge_info_slope_adjusted));
65    let bias = 0.25; // Using the bias and then saturating nudges the values a bit
66    let scale = depth_center * 0.011; // Weight the edges by their distance from the camera
67    edge_info = saturate((1.0 + bias) - edge_info / scale); // Apply the bias and scale, and invert edge_info so that small values become large, and vice versa
68
69    // Pack the edge info into the texture
70    let edge_info_packed = vec4<u32>(pack4x8unorm(edge_info), 0u, 0u, 0u);
71    textureStore(depth_differences, pixel_coordinates, edge_info_packed);
72
73    return depth_center;
74}
75
76fn load_normal_view_space(uv: vec2<f32>) -> vec3<f32> {
77    var world_normal = textureSampleLevel(normals, point_clamp_sampler, uv, 0.0).xyz;
78    world_normal = (world_normal * 2.0) - 1.0;
79    let view_from_world = mat3x3<f32>(
80        view.view_from_world[0].xyz,
81        view.view_from_world[1].xyz,
82        view.view_from_world[2].xyz,
83    );
84    return view_from_world * world_normal;
85}
86
87fn reconstruct_view_space_position(depth: f32, uv: vec2<f32>) -> vec3<f32> {
88    let clip_xy = vec2<f32>(uv.x * 2.0 - 1.0, 1.0 - 2.0 * uv.y);
89    let t = view.view_from_clip * vec4<f32>(clip_xy, depth, 1.0);
90    let view_xyz = t.xyz / t.w;
91    return view_xyz;
92}
93
94fn load_and_reconstruct_view_space_position(uv: vec2<f32>, sample_mip_level: f32) -> vec3<f32> {
95    let depth = textureSampleLevel(preprocessed_depth, linear_clamp_sampler, uv, sample_mip_level).r;
96    return reconstruct_view_space_position(depth, uv);
97}
98
99fn updateSectors(
100    min_horizon: f32,
101    max_horizon: f32,
102    samples_per_slice: f32,
103    bitmask: u32,
104) -> u32 {
105    let start_horizon = u32(min_horizon * samples_per_slice);
106    let angle_horizon = u32(ceil((max_horizon - min_horizon) * samples_per_slice));
107
108    return insertBits(bitmask, 0xFFFFFFFFu, start_horizon, angle_horizon);
109}
110
111fn processSample(
112    delta_position: vec3<f32>,
113    view_vec: vec3<f32>,
114    sampling_direction: f32,
115    n: vec2<f32>,
116    samples_per_slice: f32,
117    bitmask: ptr<function, u32>,
118) {
119    let delta_position_back_face = delta_position - view_vec * thickness;
120
121    var front_back_horizon = vec2(
122        fast_acos(dot(normalize(delta_position), view_vec)),
123        fast_acos(dot(normalize(delta_position_back_face), view_vec)),
124    );
125
126    front_back_horizon = saturate(fma(vec2(sampling_direction), -front_back_horizon, n));
127    front_back_horizon = select(front_back_horizon.xy, front_back_horizon.yx, sampling_direction >= 0.0);
128
129    *bitmask = updateSectors(front_back_horizon.x, front_back_horizon.y, samples_per_slice, *bitmask);
130}
131
132@compute
133@workgroup_size(8, 8, 1)
134fn ssao(@builtin(global_invocation_id) global_id: vec3<u32>) {
135    let slice_count = f32(#SLICE_COUNT);
136    let samples_per_slice_side = f32(#SAMPLES_PER_SLICE_SIDE);
137    let effect_radius = 0.5 * 1.457;
138    let falloff_range = 0.615 * effect_radius;
139    let falloff_from = effect_radius * (1.0 - 0.615);
140    let falloff_mul = -1.0 / falloff_range;
141    let falloff_add = falloff_from / falloff_range + 1.0;
142
143    let pixel_coordinates = vec2<i32>(global_id.xy);
144    let uv = (vec2<f32>(pixel_coordinates) + 0.5) / view.viewport.zw;
145
146    var pixel_depth = calculate_neighboring_depth_differences(pixel_coordinates);
147    pixel_depth += 0.00001; // Avoid depth precision issues
148
149    let pixel_position = reconstruct_view_space_position(pixel_depth, uv);
150    let pixel_normal = load_normal_view_space(uv);
151    let view_vec = normalize(-pixel_position);
152
153    let noise = load_noise(pixel_coordinates);
154    let sample_scale = (-0.5 * effect_radius * view.clip_from_view[0][0]) / pixel_position.z;
155
156    var visibility = 0.0;
157    var occluded_sample_count = 0u;
158    for (var slice_t = 0.0; slice_t < slice_count; slice_t += 1.0) {
159        let slice = slice_t + noise.x;
160        let phi = (PI / slice_count) * slice;
161        let omega = vec2<f32>(cos(phi), sin(phi));
162
163        let direction = vec3<f32>(omega.xy, 0.0);
164        let orthographic_direction = direction - (dot(direction, view_vec) * view_vec);
165        let axis = cross(direction, view_vec);
166        let projected_normal = pixel_normal - axis * dot(pixel_normal, axis);
167        let projected_normal_length = length(projected_normal);
168
169        let sign_norm = sign(dot(orthographic_direction, projected_normal));
170        let cos_norm = saturate(dot(projected_normal, view_vec) / projected_normal_length);
171        let n = vec2((HALF_PI - sign_norm * fast_acos(cos_norm)) * (1.0 / PI));
172
173        var bitmask = 0u;
174
175        let sample_mul = vec2<f32>(omega.x, -omega.y) * sample_scale;
176        for (var sample_t = 0.0; sample_t < samples_per_slice_side; sample_t += 1.0) {
177            var sample_noise = (slice_t + sample_t * samples_per_slice_side) * 0.6180339887498948482;
178            sample_noise = fract(noise.y + sample_noise);
179
180            var s = (sample_t + sample_noise) / samples_per_slice_side;
181            s *= s; // https://github.com/GameTechDev/XeGTAO#sample-distribution
182            let sample = s * sample_mul;
183
184            // * view.viewport.zw gets us from [0, 1] to [0, viewport_size], which is needed for this to get the correct mip levels
185            let sample_mip_level = clamp(log2(length(sample * view.viewport.zw)) - 3.3, 0.0, 5.0); // https://github.com/GameTechDev/XeGTAO#memory-bandwidth-bottleneck
186            let sample_position_1 = load_and_reconstruct_view_space_position(uv + sample, sample_mip_level);
187            let sample_position_2 = load_and_reconstruct_view_space_position(uv - sample, sample_mip_level);
188
189            let sample_difference_1 = sample_position_1 - pixel_position;
190            let sample_difference_2 = sample_position_2 - pixel_position;
191
192            processSample(sample_difference_1, view_vec, -1.0, n, samples_per_slice_side * 2.0, &bitmask);
193            processSample(sample_difference_2, view_vec, 1.0, n, samples_per_slice_side * 2.0, &bitmask);
194        }
195
196        occluded_sample_count += countOneBits(bitmask);
197    }
198
199    visibility = 1.0 - f32(occluded_sample_count) / (slice_count * 2.0 * samples_per_slice_side);
200
201    visibility = clamp(visibility, 0.03, 1.0);
202
203    textureStore(ambient_occlusion, pixel_coordinates, vec4<f32>(visibility, 0.0, 0.0, 0.0));
204}