1// NVIDIA FXAA 3.11
2// Original source code by TIMOTHY LOTTES
3// https://gist.github.com/kosua20/0c506b81b3812ac900048059d2383126
4//
5// Cleaned version - https://github.com/kosua20/Rendu/blob/master/resources/common/shaders/screens/fxaa.frag
6//
7// Tweaks by mrDIMAS - https://github.com/FyroxEngine/Fyrox/blob/master/src/renderer/shaders/fxaa_fs.glsl
8
9import bevy_core_pipeline::fullscreen_vertex_shader::FullscreenVertexOutput;
10
11@group(0) @binding(0) var screenTexture: texture_2d<f32>;
12@group(0) @binding(1) var samp: sampler;
13
14/// Trims the algorithm from processing darks.
15    @if(EDGE_THRESH_MIN_LOW)
16    const EDGE_THRESHOLD_MIN: f32 = 0.0833;
17
18    @if(EDGE_THRESH_MIN_MEDIUM)
19    const EDGE_THRESHOLD_MIN: f32 = 0.0625;
20
21    @if(EDGE_THRESH_MIN_HIGH)
22    const EDGE_THRESHOLD_MIN: f32 = 0.0312;
23
24    @if(EDGE_THRESH_MIN_ULTRA)
25    const EDGE_THRESHOLD_MIN: f32 = 0.0156;
26
27    @if(EDGE_THRESH_MIN_EXTREME)
28    const EDGE_THRESHOLD_MIN: f32 = 0.0078;
29
30/// The minimum amount of local contrast required to apply algorithm.
31    @if(EDGE_THRESH_LOW)
32    const EDGE_THRESHOLD_MAX: f32 = 0.250;
33
34    @if(EDGE_THRESH_MEDIUM)
35    const EDGE_THRESHOLD_MAX: f32 = 0.166;
36
37    @if(EDGE_THRESH_HIGH)
38    const EDGE_THRESHOLD_MAX: f32 = 0.125;
39
40    @if(EDGE_THRESH_ULTRA)
41    const EDGE_THRESHOLD_MAX: f32 = 0.063;
42
43    @if(EDGE_THRESH_EXTREME)
44    const EDGE_THRESHOLD_MAX: f32 = 0.031;
45
46const ITERATIONS: i32 = 12; //default is 12
47const SUBPIXEL_QUALITY: f32 = 0.75;
48// #define QUALITY(q) ((q) < 5 ? 1.0 : ((q) > 5 ? ((q) < 10 ? 2.0 : ((q) < 11 ? 4.0 : 8.0)) : 1.5))
49fn QUALITY(q: i32) -> f32 {
50    switch (q) {
51        //case 0, 1, 2, 3, 4: { return 1.0; }
52        default:              { return 1.0; }
53        case 5:               { return 1.5; }
54        case 6, 7, 8, 9:      { return 2.0; }
55        case 10:              { return 4.0; }
56        case 11:              { return 8.0; }
57    }
58}
59
60fn rgb2luma(rgb: vec3<f32>) -> f32 {
61    return sqrt(dot(rgb, vec3<f32>(0.299, 0.587, 0.114)));
62}
63
64/// Performs FXAA post-process anti-aliasing as described in the Nvidia FXAA white paper and the associated shader code.
65@fragment
66fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
67    let resolution = vec2<f32>(textureDimensions(screenTexture));
68    let inverseScreenSize = 1.0 / resolution.xy;
69    let texCoord = in.position.xy * inverseScreenSize;
70
71    let centerSample = textureSampleLevel(screenTexture, samp, texCoord, 0.0);
72    let colorCenter = centerSample.rgb;
73
74    // Luma at the current fragment
75    let lumaCenter = rgb2luma(colorCenter);
76
77    // Luma at the four direct neighbors of the current fragment.
78    let lumaDown = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(0, -1)).rgb);
79    let lumaUp = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(0, 1)).rgb);
80    let lumaLeft = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(-1, 0)).rgb);
81    let lumaRight = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(1, 0)).rgb);
82
83    // Find the maximum and minimum luma around the current fragment.
84    let lumaMin = min(lumaCenter, min(min(lumaDown, lumaUp), min(lumaLeft, lumaRight)));
85    let lumaMax = max(lumaCenter, max(max(lumaDown, lumaUp), max(lumaLeft, lumaRight)));
86
87    // Compute the delta.
88    let lumaRange = lumaMax - lumaMin;
89
90    // If the luma variation is lower that a threshold (or if we are in a really dark area), we are not on an edge, don't perform any AA.
91    if (lumaRange < max(EDGE_THRESHOLD_MIN, lumaMax * EDGE_THRESHOLD_MAX)) {
92        return centerSample;
93    }
94
95    // Query the 4 remaining corners lumas.
96    let lumaDownLeft  = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(-1, -1)).rgb);
97    let lumaUpRight   = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(1, 1)).rgb);
98    let lumaUpLeft    = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(-1, 1)).rgb);
99    let lumaDownRight = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(1, -1)).rgb);
100
101    // Combine the four edges lumas (using intermediary variables for future computations with the same values).
102    let lumaDownUp = lumaDown + lumaUp;
103    let lumaLeftRight = lumaLeft + lumaRight;
104
105    // Same for corners
106    let lumaLeftCorners = lumaDownLeft + lumaUpLeft;
107    let lumaDownCorners = lumaDownLeft + lumaDownRight;
108    let lumaRightCorners = lumaDownRight + lumaUpRight;
109    let lumaUpCorners = lumaUpRight + lumaUpLeft;
110
111    // Compute an estimation of the gradient along the horizontal and vertical axis.
112    let edgeHorizontal = abs(-2.0 * lumaLeft   + lumaLeftCorners)  +
113                         abs(-2.0 * lumaCenter + lumaDownUp) * 2.0 +
114                         abs(-2.0 * lumaRight  + lumaRightCorners);
115
116    let edgeVertical =   abs(-2.0 * lumaUp     + lumaUpCorners)       +
117                         abs(-2.0 * lumaCenter + lumaLeftRight) * 2.0 +
118                         abs(-2.0 * lumaDown   + lumaDownCorners);
119
120    // Is the local edge horizontal or vertical ?
121    let isHorizontal = (edgeHorizontal >= edgeVertical);
122
123    // Choose the step size (one pixel) accordingly.
124    var stepLength = select(inverseScreenSize.x, inverseScreenSize.y, isHorizontal);
125
126    // Select the two neighboring texels lumas in the opposite direction to the local edge.
127    var luma1 = select(lumaLeft, lumaDown, isHorizontal);
128    var luma2 = select(lumaRight, lumaUp, isHorizontal);
129
130    // Compute gradients in this direction.
131    let gradient1 = luma1 - lumaCenter;
132    let gradient2 = luma2 - lumaCenter;
133
134    // Which direction is the steepest ?
135    let is1Steepest = abs(gradient1) >= abs(gradient2);
136
137    // Gradient in the corresponding direction, normalized.
138    let gradientScaled = 0.25 * max(abs(gradient1), abs(gradient2));
139
140    // Average luma in the correct direction.
141    var lumaLocalAverage = 0.0;
142    if (is1Steepest) {
143        // Switch the direction
144        stepLength = -stepLength;
145        lumaLocalAverage = 0.5 * (luma1 + lumaCenter);
146    } else {
147        lumaLocalAverage = 0.5 * (luma2 + lumaCenter);
148    }
149
150    // Shift UV in the correct direction by half a pixel.
151    // Compute offset (for each iteration step) in the right direction.
152    var currentUv = texCoord;
153    var offset = vec2<f32>(0.0, 0.0);
154    if (isHorizontal) {
155        currentUv.y = currentUv.y + stepLength * 0.5;
156        offset.x = inverseScreenSize.x;
157    } else {
158        currentUv.x = currentUv.x + stepLength * 0.5;
159        offset.y = inverseScreenSize.y;
160    }
161
162    // Compute UVs to explore on each side of the edge, orthogonally. The QUALITY allows us to step faster.
163    var uv1 = currentUv - offset; // * QUALITY(0); // (quality 0 is 1.0)
164    var uv2 = currentUv + offset; // * QUALITY(0); // (quality 0 is 1.0)
165
166    // Read the lumas at both current extremities of the exploration segment, and compute the delta wrt to the local average luma.
167    var lumaEnd1 = rgb2luma(textureSampleLevel(screenTexture, samp, uv1, 0.0).rgb);
168    var lumaEnd2 = rgb2luma(textureSampleLevel(screenTexture, samp, uv2, 0.0).rgb);
169    lumaEnd1 = lumaEnd1 - lumaLocalAverage;
170    lumaEnd2 = lumaEnd2 - lumaLocalAverage;
171
172    // If the luma deltas at the current extremities is larger than the local gradient, we have reached the side of the edge.
173    var reached1 = abs(lumaEnd1) >= gradientScaled;
174    var reached2 = abs(lumaEnd2) >= gradientScaled;
175    var reachedBoth = reached1 && reached2;
176
177    // If the side is not reached, we continue to explore in this direction.
178    uv1 = select(uv1 - offset, uv1, reached1); // * QUALITY(1); // (quality 1 is 1.0)
179    uv2 = select(uv2 - offset, uv2, reached2); // * QUALITY(1); // (quality 1 is 1.0)
180
181    // If both sides have not been reached, continue to explore.
182    if (!reachedBoth) {
183        for (var i: i32 = 2; i < ITERATIONS; i = i + 1) {
184            // If needed, read luma in 1st direction, compute delta.
185            if (!reached1) {
186                lumaEnd1 = rgb2luma(textureSampleLevel(screenTexture, samp, uv1, 0.0).rgb);
187                lumaEnd1 = lumaEnd1 - lumaLocalAverage;
188            }
189            // If needed, read luma in opposite direction, compute delta.
190            if (!reached2) {
191                lumaEnd2 = rgb2luma(textureSampleLevel(screenTexture, samp, uv2, 0.0).rgb);
192                lumaEnd2 = lumaEnd2 - lumaLocalAverage;
193            }
194            // If the luma deltas at the current extremities is larger than the local gradient, we have reached the side of the edge.
195            reached1 = abs(lumaEnd1) >= gradientScaled;
196            reached2 = abs(lumaEnd2) >= gradientScaled;
197            reachedBoth = reached1 && reached2;
198
199            // If the side is not reached, we continue to explore in this direction, with a variable quality.
200            if (!reached1) {
201                uv1 = uv1 - offset * QUALITY(i);
202            }
203            if (!reached2) {
204                uv2 = uv2 + offset * QUALITY(i);
205            }
206
207            // If both sides have been reached, stop the exploration.
208            if (reachedBoth) {
209                break;
210            }
211        }
212    }
213
214    // Compute the distances to each side edge of the edge (!).
215    var distance1 = select(texCoord.y - uv1.y, texCoord.x - uv1.x, isHorizontal);
216    var distance2 = select(uv2.y - texCoord.y, uv2.x - texCoord.x, isHorizontal);
217
218    // In which direction is the side of the edge closer ?
219    let isDirection1 = distance1 < distance2;
220    let distanceFinal = min(distance1, distance2);
221
222    // Thickness of the edge.
223    let edgeThickness = (distance1 + distance2);
224
225    // Is the luma at center smaller than the local average ?
226    let isLumaCenterSmaller = lumaCenter < lumaLocalAverage;
227
228    // If the luma at center is smaller than at its neighbor, the delta luma at each end should be positive (same variation).
229    let correctVariation1 = (lumaEnd1 < 0.0) != isLumaCenterSmaller;
230    let correctVariation2 = (lumaEnd2 < 0.0) != isLumaCenterSmaller;
231
232    // Only keep the result in the direction of the closer side of the edge.
233    var correctVariation = select(correctVariation2, correctVariation1, isDirection1);
234
235    // UV offset: read in the direction of the closest side of the edge.
236    let pixelOffset = - distanceFinal / edgeThickness + 0.5;
237
238    // If the luma variation is incorrect, do not offset.
239    var finalOffset = select(0.0, pixelOffset, correctVariation);
240
241    // Sub-pixel shifting
242    // Full weighted average of the luma over the 3x3 neighborhood.
243    let lumaAverage = (1.0 / 12.0) * (2.0 * (lumaDownUp + lumaLeftRight) + lumaLeftCorners + lumaRightCorners);
244    // Ratio of the delta between the global average and the center luma, over the luma range in the 3x3 neighborhood.
245    let subPixelOffset1 = clamp(abs(lumaAverage - lumaCenter) / lumaRange, 0.0, 1.0);
246    let subPixelOffset2 = (-2.0 * subPixelOffset1 + 3.0) * subPixelOffset1 * subPixelOffset1;
247    // Compute a sub-pixel offset based on this delta.
248    let subPixelOffsetFinal = subPixelOffset2 * subPixelOffset2 * SUBPIXEL_QUALITY;
249
250    // Pick the biggest of the two offsets.
251    finalOffset = max(finalOffset, subPixelOffsetFinal);
252
253    // Compute the final UV coordinates.
254    var finalUv = texCoord;
255    if (isHorizontal) {
256        finalUv.y = finalUv.y + finalOffset * stepLength;
257    } else {
258        finalUv.x = finalUv.x + finalOffset * stepLength;
259    }
260
261    // Read the color at the new UV coordinates, and use it.
262    var finalColor = textureSampleLevel(screenTexture, samp, finalUv, 0.0).rgb;
263    return vec4<f32>(finalColor, centerSample.a);
264}