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