1// References:
2// https://www.elopezr.com/temporal-aa-and-the-quest-for-the-holy-trail
3// http://behindthepixels.io/assets/files/TemporalAA.pdf
4// http://leiy.cc/publications/TAA/TAA_EG2020_Talk.pdf
5// https://advances.realtimerendering.com/s2014/index.html#_HIGH-QUALITY_TEMPORAL_SUPERSAMPLING
6
7/// Controls how much to blend between the current and past samples
8/// Lower numbers = less of the current sample and more of the past sample = more smoothing
9/// Values chosen empirically
10const DEFAULT_HISTORY_BLEND_RATE: f32 = 0.1; // Default blend rate to use when no confidence in history
11const MIN_HISTORY_BLEND_RATE: f32 = 0.015; // Minimum blend rate allowed, to ensure at least some of the current sample is used
12
13@group(0) @binding(0) var view_target: texture_2d<f32>;
14@group(0) @binding(1) var history: texture_2d<f32>;
15@group(0) @binding(2) var motion_vectors: texture_2d<f32>;
16@group(0) @binding(3) var depth: texture_depth_2d;
17@group(0) @binding(4) var nearest_sampler: sampler;
18@group(0) @binding(5) var linear_sampler: sampler;
19
20struct Output {
21    @location(0) view_target: vec4<f32>,
22    @location(1) history: vec4<f32>,
23};
24
25// TAA is ideally applied after tonemapping, but before post processing
26// Post processing wants to go before tonemapping, which conflicts
27// Solution: Put TAA before tonemapping, tonemap TAA input, apply TAA, invert-tonemap TAA output
28// https://advances.realtimerendering.com/s2014/index.html#_HIGH-QUALITY_TEMPORAL_SUPERSAMPLING, slide 20
29// https://gpuopen.com/learn/optimized-reversible-tonemapper-for-resolve
30fn rcp(x: f32) -> f32 { return 1.0 / x; }
31fn max3(x: vec3<f32>) -> f32 { return max(x.r, max(x.g, x.b)); }
32fn tonemap(color: vec3<f32>) -> vec3<f32> { return color * rcp(max3(color) + 1.0); }
33fn reverse_tonemap(color: vec3<f32>) -> vec3<f32> { return color * rcp(1.0 - max3(color)); }
34
35// The following 3 functions are from Playdead (MIT-licensed)
36// https://github.com/playdeadgames/temporal/blob/master/Assets/Shaders/TemporalReprojection.shader
37fn RGB_to_YCoCg(rgb: vec3<f32>) -> vec3<f32> {
38    let y = (rgb.r / 4.0) + (rgb.g / 2.0) + (rgb.b / 4.0);
39    let co = (rgb.r / 2.0) - (rgb.b / 2.0);
40    let cg = (-rgb.r / 4.0) + (rgb.g / 2.0) - (rgb.b / 4.0);
41    return vec3(y, co, cg);
42}
43
44fn YCoCg_to_RGB(ycocg: vec3<f32>) -> vec3<f32> {
45    let r = ycocg.x + ycocg.y - ycocg.z;
46    let g = ycocg.x + ycocg.z;
47    let b = ycocg.x - ycocg.y - ycocg.z;
48    return saturate(vec3(r, g, b));
49}
50
51fn clip_towards_aabb_center(history_color: vec3<f32>, current_color: vec3<f32>, aabb_min: vec3<f32>, aabb_max: vec3<f32>) -> vec3<f32> {
52    let p_clip = 0.5 * (aabb_max + aabb_min);
53    let e_clip = 0.5 * (aabb_max - aabb_min) + 0.00000001;
54    let v_clip = history_color - p_clip;
55    let v_unit = v_clip / e_clip;
56    let a_unit = abs(v_unit);
57    let ma_unit = max3(a_unit);
58    if ma_unit > 1.0 {
59        return p_clip + (v_clip / ma_unit);
60    } else {
61        return history_color;
62    }
63}
64
65fn sample_history(u: f32, v: f32) -> vec3<f32> {
66    return textureSample(history, linear_sampler, vec2(u, v)).rgb;
67}
68
69fn sample_view_target(uv: vec2<f32>) -> vec3<f32> {
70    var sample = textureSample(view_target, nearest_sampler, uv).rgb;
71@if(TONEMAP)
72    sample = tonemap(sample);
73    return RGB_to_YCoCg(sample);
74}
75
76@fragment
77fn taa(@location(0) uv: vec2<f32>) -> Output {
78    let texture_size = vec2<f32>(textureDimensions(view_target));
79    let texel_size = 1.0 / texture_size;
80
81    // Fetch the current sample
82    let original_color = textureSample(view_target, nearest_sampler, uv);
83    var current_color = original_color.rgb;
84@if(TONEMAP)
85    current_color = tonemap(current_color);
86
87    var history_confidence: f32;
88@if(!RESET) {
89    // Pick the closest motion_vector from 5 samples (reduces aliasing on the edges of moving entities)
90    // https://advances.realtimerendering.com/s2014/index.html#_HIGH-QUALITY_TEMPORAL_SUPERSAMPLING, slide 27
91    let offset = texel_size * 2.0;
92    let d_uv_tl = uv + vec2(-offset.x, offset.y);
93    let d_uv_tr = uv + vec2(offset.x, offset.y);
94    let d_uv_bl = uv + vec2(-offset.x, -offset.y);
95    let d_uv_br = uv + vec2(offset.x, -offset.y);
96    var closest_uv = uv;
97    let d_tl = textureSample(depth, nearest_sampler, d_uv_tl);
98    let d_tr = textureSample(depth, nearest_sampler, d_uv_tr);
99    var closest_depth = textureSample(depth, nearest_sampler, uv);
100    let d_bl = textureSample(depth, nearest_sampler, d_uv_bl);
101    let d_br = textureSample(depth, nearest_sampler, d_uv_br);
102    if d_tl > closest_depth {
103        closest_uv = d_uv_tl;
104        closest_depth = d_tl;
105    }
106    if d_tr > closest_depth {
107        closest_uv = d_uv_tr;
108        closest_depth = d_tr;
109    }
110    if d_bl > closest_depth {
111        closest_uv = d_uv_bl;
112        closest_depth = d_bl;
113    }
114    if d_br > closest_depth {
115        closest_uv = d_uv_br;
116    }
117    let closest_motion_vector = textureSample(motion_vectors, nearest_sampler, closest_uv).rg;
118
119    // Reproject to find the equivalent sample from the past
120    // Uses 5-sample Catmull-Rom filtering (reduces blurriness)
121    // Catmull-Rom filtering: https://gist.github.com/TheRealMJP/c83b8c0f46b63f3a88a5986f4fa982b1
122    // Ignoring corners: https://www.activision.com/cdn/research/Dynamic_Temporal_Antialiasing_and_Upsampling_in_Call_of_Duty_v4.pdf#page=68
123    // Technically we should renormalize the weights since we're skipping the corners, but it's basically the same result
124    let history_uv = uv - closest_motion_vector;
125    let sample_position = history_uv * texture_size;
126    let texel_center = floor(sample_position - 0.5) + 0.5;
127    let f = sample_position - texel_center;
128    let w0 = f * (-0.5 + f * (1.0 - 0.5 * f));
129    let w1 = 1.0 + f * f * (-2.5 + 1.5 * f);
130    let w2 = f * (0.5 + f * (2.0 - 1.5 * f));
131    let w3 = f * f * (-0.5 + 0.5 * f);
132    let w12 = w1 + w2;
133    let texel_position_0 = (texel_center - 1.0) * texel_size;
134    let texel_position_3 = (texel_center + 2.0) * texel_size;
135    let texel_position_12 = (texel_center + (w2 / w12)) * texel_size;
136    var history_color = sample_history(texel_position_12.x, texel_position_0.y) * w12.x * w0.y;
137    history_color += sample_history(texel_position_0.x, texel_position_12.y) * w0.x * w12.y;
138    history_color += sample_history(texel_position_12.x, texel_position_12.y) * w12.x * w12.y;
139    history_color += sample_history(texel_position_3.x, texel_position_12.y) * w3.x * w12.y;
140    history_color += sample_history(texel_position_12.x, texel_position_3.y) * w12.x * w3.y;
141
142    // Constrain past sample with 3x3 YCoCg variance clipping (reduces ghosting)
143    // YCoCg: https://advances.realtimerendering.com/s2014/index.html#_HIGH-QUALITY_TEMPORAL_SUPERSAMPLING, slide 33
144    // Variance clipping: https://developer.download.nvidia.com/gameworks/events/GDC2016/msalvi_temporal_supersampling.pdf
145    let s_tl = sample_view_target(uv + vec2(-texel_size.x,  texel_size.y));
146    let s_tm = sample_view_target(uv + vec2( 0.0,           texel_size.y));
147    let s_tr = sample_view_target(uv + vec2( texel_size.x,  texel_size.y));
148    let s_ml = sample_view_target(uv + vec2(-texel_size.x,  0.0));
149    let s_mm = RGB_to_YCoCg(current_color);
150    let s_mr = sample_view_target(uv + vec2( texel_size.x,  0.0));
151    let s_bl = sample_view_target(uv + vec2(-texel_size.x, -texel_size.y));
152    let s_bm = sample_view_target(uv + vec2( 0.0,          -texel_size.y));
153    let s_br = sample_view_target(uv + vec2( texel_size.x, -texel_size.y));
154    let moment_1 = s_tl + s_tm + s_tr + s_ml + s_mm + s_mr + s_bl + s_bm + s_br;
155    let moment_2 = (s_tl * s_tl) + (s_tm * s_tm) + (s_tr * s_tr) + (s_ml * s_ml) + (s_mm * s_mm) + (s_mr * s_mr) + (s_bl * s_bl) + (s_bm * s_bm) + (s_br * s_br);
156    let mean = moment_1 / 9.0;
157    let variance = (moment_2 / 9.0) - (mean * mean);
158    let std_deviation = sqrt(max(variance, vec3(0.0)));
159    history_color = RGB_to_YCoCg(history_color);
160    history_color = clip_towards_aabb_center(history_color, s_mm, mean - std_deviation, mean + std_deviation);
161    history_color = YCoCg_to_RGB(history_color);
162
163    // How confident we are that the history is representative of the current frame
164    history_confidence = textureSample(history, nearest_sampler, uv).a;
165    let pixel_motion_vector = abs(closest_motion_vector) * texture_size;
166    if pixel_motion_vector.x < 0.01 && pixel_motion_vector.y < 0.01 {
167        // Increment when pixels are not moving
168        history_confidence += 10.0;
169    } else {
170        // Else reset
171        history_confidence = 1.0;
172    }
173
174    // Blend current and past sample
175    // Use more of the history if we're confident in it (reduces noise when there is no motion)
176    // https://hhoppe.com/supersample.pdf, section 4.1
177    var current_color_factor = clamp(1.0 / history_confidence, MIN_HISTORY_BLEND_RATE, DEFAULT_HISTORY_BLEND_RATE);
178
179    // Reject history when motion vectors point off screen
180    if any(saturate(history_uv) != history_uv) {
181        current_color_factor = 1.0;
182        history_confidence = 1.0;
183    }
184
185    current_color = mix(history_color, current_color, current_color_factor);
186} // !RESET
187
188
189    // Write output to history and view target
190    var out: Output;
191@if(RESET)
192    history_confidence = 1.0 / MIN_HISTORY_BLEND_RATE;
193    out.history = vec4(current_color, history_confidence);
194@if(TONEMAP)
195    current_color = reverse_tonemap(current_color);
196    out.view_target = vec4(current_color, original_color.a);
197    return out;
198}