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#ifdef TONEMAP
72    sample = tonemap(sample);
73#endif
74    return RGB_to_YCoCg(sample);
75}
76
77@fragment
78fn taa(@location(0) uv: vec2<f32>) -> Output {
79    let texture_size = vec2<f32>(textureDimensions(view_target));
80    let texel_size = 1.0 / texture_size;
81
82    // Fetch the current sample
83    let original_color = textureSample(view_target, nearest_sampler, uv);
84    var current_color = original_color.rgb;
85#ifdef TONEMAP
86    current_color = tonemap(current_color);
87#endif
88
89#ifndef RESET
90    // Pick the closest motion_vector from 5 samples (reduces aliasing on the edges of moving entities)
91    // https://advances.realtimerendering.com/s2014/index.html#_HIGH-QUALITY_TEMPORAL_SUPERSAMPLING, slide 27
92    let offset = texel_size * 2.0;
93    let d_uv_tl = uv + vec2(-offset.x, offset.y);
94    let d_uv_tr = uv + vec2(offset.x, offset.y);
95    let d_uv_bl = uv + vec2(-offset.x, -offset.y);
96    let d_uv_br = uv + vec2(offset.x, -offset.y);
97    var closest_uv = uv;
98    let d_tl = textureSample(depth, nearest_sampler, d_uv_tl);
99    let d_tr = textureSample(depth, nearest_sampler, d_uv_tr);
100    var closest_depth = textureSample(depth, nearest_sampler, uv);
101    let d_bl = textureSample(depth, nearest_sampler, d_uv_bl);
102    let d_br = textureSample(depth, nearest_sampler, d_uv_br);
103    if d_tl > closest_depth {
104        closest_uv = d_uv_tl;
105        closest_depth = d_tl;
106    }
107    if d_tr > closest_depth {
108        closest_uv = d_uv_tr;
109        closest_depth = d_tr;
110    }
111    if d_bl > closest_depth {
112        closest_uv = d_uv_bl;
113        closest_depth = d_bl;
114    }
115    if d_br > closest_depth {
116        closest_uv = d_uv_br;
117    }
118    let closest_motion_vector = textureSample(motion_vectors, nearest_sampler, closest_uv).rg;
119
120    // Reproject to find the equivalent sample from the past
121    // Uses 5-sample Catmull-Rom filtering (reduces blurriness)
122    // Catmull-Rom filtering: https://gist.github.com/TheRealMJP/c83b8c0f46b63f3a88a5986f4fa982b1
123    // Ignoring corners: https://www.activision.com/cdn/research/Dynamic_Temporal_Antialiasing_and_Upsampling_in_Call_of_Duty_v4.pdf#page=68
124    // Technically we should renormalize the weights since we're skipping the corners, but it's basically the same result
125    let history_uv = uv - closest_motion_vector;
126    let sample_position = history_uv * texture_size;
127    let texel_center = floor(sample_position - 0.5) + 0.5;
128    let f = sample_position - texel_center;
129    let w0 = f * (-0.5 + f * (1.0 - 0.5 * f));
130    let w1 = 1.0 + f * f * (-2.5 + 1.5 * f);
131    let w2 = f * (0.5 + f * (2.0 - 1.5 * f));
132    let w3 = f * f * (-0.5 + 0.5 * f);
133    let w12 = w1 + w2;
134    let texel_position_0 = (texel_center - 1.0) * texel_size;
135    let texel_position_3 = (texel_center + 2.0) * texel_size;
136    let texel_position_12 = (texel_center + (w2 / w12)) * texel_size;
137    var history_color = sample_history(texel_position_12.x, texel_position_0.y) * w12.x * w0.y;
138    history_color += sample_history(texel_position_0.x, texel_position_12.y) * w0.x * w12.y;
139    history_color += sample_history(texel_position_12.x, texel_position_12.y) * w12.x * w12.y;
140    history_color += sample_history(texel_position_3.x, texel_position_12.y) * w3.x * w12.y;
141    history_color += sample_history(texel_position_12.x, texel_position_3.y) * w12.x * w3.y;
142
143    // Constrain past sample with 3x3 YCoCg variance clipping (reduces ghosting)
144    // YCoCg: https://advances.realtimerendering.com/s2014/index.html#_HIGH-QUALITY_TEMPORAL_SUPERSAMPLING, slide 33
145    // Variance clipping: https://developer.download.nvidia.com/gameworks/events/GDC2016/msalvi_temporal_supersampling.pdf
146    let s_tl = sample_view_target(uv + vec2(-texel_size.x,  texel_size.y));
147    let s_tm = sample_view_target(uv + vec2( 0.0,           texel_size.y));
148    let s_tr = sample_view_target(uv + vec2( texel_size.x,  texel_size.y));
149    let s_ml = sample_view_target(uv + vec2(-texel_size.x,  0.0));
150    let s_mm = RGB_to_YCoCg(current_color);
151    let s_mr = sample_view_target(uv + vec2( texel_size.x,  0.0));
152    let s_bl = sample_view_target(uv + vec2(-texel_size.x, -texel_size.y));
153    let s_bm = sample_view_target(uv + vec2( 0.0,          -texel_size.y));
154    let s_br = sample_view_target(uv + vec2( texel_size.x, -texel_size.y));
155    let moment_1 = s_tl + s_tm + s_tr + s_ml + s_mm + s_mr + s_bl + s_bm + s_br;
156    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);
157    let mean = moment_1 / 9.0;
158    let variance = (moment_2 / 9.0) - (mean * mean);
159    let std_deviation = sqrt(max(variance, vec3(0.0)));
160    history_color = RGB_to_YCoCg(history_color);
161    history_color = clip_towards_aabb_center(history_color, s_mm, mean - std_deviation, mean + std_deviation);
162    history_color = YCoCg_to_RGB(history_color);
163
164    // How confident we are that the history is representative of the current frame
165    var history_confidence = textureSample(history, nearest_sampler, uv).a;
166    let pixel_motion_vector = abs(closest_motion_vector) * texture_size;
167    if pixel_motion_vector.x < 0.01 && pixel_motion_vector.y < 0.01 {
168        // Increment when pixels are not moving
169        history_confidence += 10.0;
170    } else {
171        // Else reset
172        history_confidence = 1.0;
173    }
174
175    // Blend current and past sample
176    // Use more of the history if we're confident in it (reduces noise when there is no motion)
177    // https://hhoppe.com/supersample.pdf, section 4.1
178    var current_color_factor = clamp(1.0 / history_confidence, MIN_HISTORY_BLEND_RATE, DEFAULT_HISTORY_BLEND_RATE);
179
180    // Reject history when motion vectors point off screen
181    if any(saturate(history_uv) != history_uv) {
182        current_color_factor = 1.0;
183        history_confidence = 1.0;
184    }
185
186    current_color = mix(history_color, current_color, current_color_factor);
187#endif // #ifndef RESET
188
189
190    // Write output to history and view target
191    var out: Output;
192#ifdef RESET
193    let history_confidence = 1.0 / MIN_HISTORY_BLEND_RATE;
194#endif
195    out.history = vec4(current_color, history_confidence);
196#ifdef TONEMAP
197    current_color = reverse_tonemap(current_color);
198#endif
199    out.view_target = vec4(current_color, original_color.a);
200    return out;
201}