1//! Auto exposure
2//!
3//! This shader computes an auto exposure value for the current frame,
4//! which is then used as an exposure correction in the tone mapping shader.
5//!
6//! The auto exposure value is computed in two passes:
7//! * The compute_histogram pass calculates a histogram of the luminance values in the scene,
8//! taking into account the metering mask texture. The metering mask is a grayscale texture
9//! that defines the areas of the screen that should be given more weight when calculating
10//! the average luminance value. For example, the middle area of the screen might be more important
11//! than the edges.
12//! * The compute_average pass calculates the average luminance value of the scene, taking
13//! into account the low_percent and high_percent settings. These settings define the
14//! percentage of the histogram that should be excluded when calculating the average. This
15//! is useful to avoid overexposure when you have a lot of shadows, or underexposure when you
16//! have a lot of bright specular reflections.
17//!
18//! The final target_exposure is finally used to smoothly adjust the exposure value over time.
19
20import bevy_render::view::View;
21import bevy_render::globals::Globals;
22
23/// Constant to convert RGB to luminance, taken from Real Time Rendering, Vol 4 pg. 278, 4th edition
24const RGB_TO_LUM = vec3<f32>(0.2125, 0.7154, 0.0721);
25
26struct AutoExposure {
27    min_log_lum: f32,
28    inv_log_lum_range: f32,
29    log_lum_range: f32,
30    low_percent: f32,
31    high_percent: f32,
32    speed_up: f32,
33    speed_down: f32,
34    exponential_transition_distance: f32,
35}
36
37struct CompensationCurve {
38    min_log_lum: f32,
39    inv_log_lum_range: f32,
40    min_compensation: f32,
41    compensation_range: f32,
42}
43
44@group(0) @binding(0) var<uniform> globals: Globals;
45
46@group(0) @binding(1) var<uniform> settings: AutoExposure;
47
48@group(0) @binding(2) var tex_color: texture_2d<f32>;
49
50@group(0) @binding(3) var tex_mask: texture_2d<f32>;
51
52@group(0) @binding(4) var tex_compensation: texture_1d<f32>;
53
54@group(0) @binding(5) var<uniform> compensation_curve: CompensationCurve;
55
56@group(0) @binding(6) var<storage, read_write> histogram: array<atomic<u32>, 64>;
57
58@group(0) @binding(7) var<storage, read_write> exposure: f32;
59
60@group(0) @binding(8) var<storage, read_write> view: View;
61
62var<workgroup> histogram_shared: array<atomic<u32>, 64>;
63
64/// For a given color, return the histogram bin index
65fn color_to_bin(hdr: vec3<f32>) -> u32 {
66    // Convert color to luminance
67    let lum = dot(hdr, RGB_TO_LUM);
68
69    if lum < exp2(settings.min_log_lum) {
70        return 0u;
71    }
72
73    // Calculate the log_2 luminance and express it as a value in [0.0, 1.0]
74    // where 0.0 represents the minimum luminance, and 1.0 represents the max.
75    let log_lum = saturate((log2(lum) - settings.min_log_lum) * settings.inv_log_lum_range);
76
77    // Map [0, 1] to [1, 63]. The zeroth bin is handled by the epsilon check above.
78    return u32(log_lum * 62.0 + 1.0);
79}
80
81/// Read the metering mask at the given UV coordinates, returning a weight for the histogram.
82///
83/// Since the histogram is summed in the compute_average step, there is a limit to the amount of
84/// distinct values that can be represented. When using the chosen value of 16, the maximum
85/// amount of pixels that can be weighted and summed is 2^32 / 16 = 16384^2.
86fn metering_weight(coords: vec2<f32>) -> u32 {
87    let pos = vec2<i32>(coords * vec2<f32>(textureDimensions(tex_mask)));
88    let mask = textureLoad(tex_mask, pos, 0).r;
89    return u32(mask * 16.0);
90}
91
92@compute @workgroup_size(16, 16, 1)
93fn compute_histogram(
94    @builtin(global_invocation_id) global_invocation_id: vec3<u32>,
95    @builtin(local_invocation_index) local_invocation_index: u32
96) {
97    // Clear the workgroup shared histogram
98    if local_invocation_index < 64 {
99        histogram_shared[local_invocation_index] = 0u;
100    }
101
102    // Wait for all workgroup threads to clear the shared histogram
103    workgroupBarrier();
104
105    let dim = vec2<u32>(textureDimensions(tex_color));
106    let uv = vec2<f32>(global_invocation_id.xy) / vec2<f32>(dim);
107
108    if global_invocation_id.x < dim.x && global_invocation_id.y < dim.y {
109        let col = textureLoad(tex_color, vec2<i32>(global_invocation_id.xy), 0).rgb;
110        let index = color_to_bin(col);
111        let weight = metering_weight(uv);
112
113        // Increment the shared histogram bin by the weight obtained from the metering mask
114        atomicAdd(&histogram_shared[index], weight);
115    }
116
117    // Wait for all workgroup threads to finish updating the workgroup histogram
118    workgroupBarrier();
119
120    // Accumulate the workgroup histogram into the global histogram.
121    // Note that the global histogram was not cleared at the beginning,
122    // as it will be cleared in compute_average.
123    var histogram_value = atomicLoad(&histogram_shared[local_invocation_index]);
124    atomicAdd(&histogram[local_invocation_index], histogram_value);
125}
126
127@compute @workgroup_size(1, 1, 1)
128fn compute_average(@builtin(local_invocation_index) local_index: u32) {
129    var histogram_sum = 0u;
130
131    // Calculate the cumulative histogram and clear the histogram bins.
132    // Each bin in the cumulative histogram contains the sum of all bins up to that point.
133    // This way we can quickly exclude the portion of lowest and highest samples as required by
134    // the low_percent and high_percent settings.
135    for (var i=0u; i<64u; i+=1u) {
136        histogram_sum += atomicLoad(&histogram[i]);
137        histogram_shared[i] = histogram_sum;
138
139        // Clear the histogram bin for the next frame
140        histogram[i] = 0u;
141    }
142
143    let first_index = u32(f32(histogram_sum) * settings.low_percent);
144    let last_index = u32(f32(histogram_sum) * settings.high_percent);
145
146    var count = 0u;
147    var sum = 0.0;
148    for (var i=1u; i<64u; i+=1u) {
149        // The number of pixels in the bin. The histogram values are clamped to
150        // first_index and last_index to exclude the lowest and highest samples.
151        let bin_count =
152            clamp(atomicLoad(&histogram_shared[i]), first_index, last_index) -
153            clamp(atomicLoad(&histogram_shared[i - 1u]), first_index, last_index);
154
155        sum += f32(bin_count) * f32(i);
156        count += bin_count;
157    }
158
159    var avg_lum = settings.min_log_lum;
160
161    if count > 0u {
162        // The average luminance of the included histogram samples.
163        avg_lum = sum / (f32(count) * 63.0)
164            * settings.log_lum_range
165            + settings.min_log_lum;
166    }
167
168    // The position in the compensation curve texture to sample for avg_lum.
169    let u = (avg_lum - compensation_curve.min_log_lum) * compensation_curve.inv_log_lum_range;
170
171    // The target exposure is the negative of the average log luminance.
172    // The compensation value is added to the target exposure to adjust the exposure for
173    // artistic purposes.
174    let target_exposure = textureLoad(tex_compensation, i32(saturate(u) * 255.0), 0).r
175        * compensation_curve.compensation_range
176        + compensation_curve.min_compensation
177        - avg_lum;
178
179    // Smoothly adjust the `exposure` towards the `target_exposure`
180    let delta = target_exposure - exposure;
181    if target_exposure > exposure {
182        let speed_down = settings.speed_down * globals.delta_time;
183        let exp_down = speed_down / settings.exponential_transition_distance;
184        exposure = exposure + min(speed_down, delta * exp_down);
185    } else {
186        let speed_up = settings.speed_up * globals.delta_time;
187        let exp_up = speed_up / settings.exponential_transition_distance;
188        exposure = exposure + max(-speed_up, delta * exp_up);
189    }
190
191    // Apply the exposure to the color grading settings, from where it will be used for the color
192    // grading pass.
193    view.color_grading.exposure += exposure;
194}