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
20#import bevy_render::view::View
21#import 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    atomicAdd(&histogram[local_invocation_index], histogram_shared[local_invocation_index]);
124}
125
126@compute @workgroup_size(1, 1, 1)
127fn compute_average(@builtin(local_invocation_index) local_index: u32) {
128    var histogram_sum = 0u;
129
130    // Calculate the cumulative histogram and clear the histogram bins.
131    // Each bin in the cumulative histogram contains the sum of all bins up to that point.
132    // This way we can quickly exclude the portion of lowest and highest samples as required by
133    // the low_percent and high_percent settings.
134    for (var i=0u; i<64u; i+=1u) {
135        histogram_sum += histogram[i];
136        histogram_shared[i] = histogram_sum;
137
138        // Clear the histogram bin for the next frame
139        histogram[i] = 0u;
140    }
141
142    let first_index = u32(f32(histogram_sum) * settings.low_percent);
143    let last_index = u32(f32(histogram_sum) * settings.high_percent);
144
145    var count = 0u;
146    var sum = 0.0;
147    for (var i=1u; i<64u; i+=1u) {
148        // The number of pixels in the bin. The histogram values are clamped to
149        // first_index and last_index to exclude the lowest and highest samples.
150        let bin_count =
151            clamp(histogram_shared[i], first_index, last_index) -
152            clamp(histogram_shared[i - 1u], first_index, last_index);
153
154        sum += f32(bin_count) * f32(i);
155        count += bin_count;
156    }
157
158    var avg_lum = settings.min_log_lum;
159
160    if count > 0u {
161        // The average luminance of the included histogram samples.
162        avg_lum = sum / (f32(count) * 63.0)
163            * settings.log_lum_range
164            + settings.min_log_lum;
165    }
166
167    // The position in the compensation curve texture to sample for avg_lum.
168    let u = (avg_lum - compensation_curve.min_log_lum) * compensation_curve.inv_log_lum_range;
169
170    // The target exposure is the negative of the average log luminance.
171    // The compensation value is added to the target exposure to adjust the exposure for
172    // artistic purposes.
173    let target_exposure = textureLoad(tex_compensation, i32(saturate(u) * 255.0), 0).r
174        * compensation_curve.compensation_range
175        + compensation_curve.min_compensation
176        - avg_lum;
177
178    // Smoothly adjust the `exposure` towards the `target_exposure`
179    let delta = target_exposure - exposure;
180    if target_exposure > exposure {
181        let speed_down = settings.speed_down * globals.delta_time;
182        let exp_down = speed_down / settings.exponential_transition_distance;
183        exposure = exposure + min(speed_down, delta * exp_down);
184    } else {
185        let speed_up = settings.speed_up * globals.delta_time;
186        let exp_up = speed_up / settings.exponential_transition_distance;
187        exposure = exposure + max(-speed_up, delta * exp_up);
188    }
189
190    // Apply the exposure to the color grading settings, from where it will be used for the color
191    // grading pass.
192    view.color_grading.exposure += exposure;
193}