DarkRoom/core/dr-gpu/src/shaders/raw_histogram.wgsl

// TRACES: FR-CULL-3
// A reduction of the *scene-linear* frame onto a stops-below-saturation axis.
//
// The shader beside this one, `histogram.wgsl`, counts the frame the display
// is about to show: an 8-bit code value, after white balance, the camera
// matrix, the tone curve, the view transform and the output transform. This one
// counts the texture the demosaic wrote, before any of that. The two differ in
// exactly one place — the axis — and everything else here is deliberately the
// same construction, because the two reductions have the same shape and any
// divergence between them would be a difference nobody chose.
//
// **Why the axis is logarithmic and measured downward from 1.0.** The texture
// is normalised by the sensor's own black and white levels (see
// `demosaic.wgsl`), so 1.0 *is* saturation by construction and there is no
// other meaningful origin. And the question a photographer asks of raw data is
// "how much headroom is left", which is a question in stops: a linear axis
// spends half its width on the top stop and crushes the eleven below it into
// the leftmost pixel, which is why nobody has ever drawn a useful linear raw
// histogram.
//
// So bin 255 is the top 1/16 stop below saturation, bin 239 is one stop below,
// bin 0 is 15.9375 stops below **and everything darker**. A bin is 1/16 of a
// stop, which is about 4.4% in value — fine enough that a clipped highlight is
// visibly its own column rather than smeared over the top quarter of the plot.
//
// **Counted per workgroup first, then merged**, for the reason
// `histogram.wgsl` gives at length: a photograph is not noise, tens of
// thousands of adjacent pixels land in one bin, and having every invocation
// contend for one global atomic serialises the dispatch. Four kilobytes of
// workgroup storage against a 16 KB floor.

const BINS: u32 = 256u;
// Bins per stop. Must equal `BINS_PER_STOP` on the Rust side, which turns a
// bin index back into a stops figure for the readout — the two disagreeing
// would put a correct plot under a wrong number.
const BINS_PER_STOP: f32 = 16.0;

// Two counters past the four series, in the same buffer and for the same
// reason `histogram.wgsl` puts its there: they are gathered from the texel
// read the binning already did, and a second reduction to answer them would be
// a second pass over the whole image for two integers.
const SATURATED: u32 = 4u * BINS;
const AT_BLACK: u32 = 4u * BINS + 1u;
const SLOTS: u32 = 4u * BINS + 2u;
// 16x16. Stated as a constant because the clear and merge loops stride by it,
// and a workgroup size that disagreed would leave slots uncleared.
const THREADS: u32 = 256u;

// How close to 1.0 counts as saturated.
//
// **Not `>= 1.0`, and the margin is arithmetic rather than taste.** A photosite
// at or above the declared white level leaves `demosaic.wgsl` clamped to
// exactly 1.0, but it then travels through an f16 texture and, at a green
// site, through an interpolation kernel whose terms are added in a different
// order than a reader would expect. f16 spaces its values 2^-11 apart just
// below 1.0, so a value that was 1.0 can arrive an ulp or two under it. This
// margin is four of those, and 0.0014 of a stop — a hundredth of the width of
// one bin — so it cannot move a column of the plot, only stop the counter
// missing a genuinely blown pixel.
const SATURATION: f32 = 0.999;

struct Dims {
    width: u32,
    height: u32,
    // std140 rounds a uniform block up to 16 bytes; named rather than left
    // implicit so the Rust side's padding is visibly the same shape.
    pad_0: u32,
    pad_1: u32,
}

@group(0) @binding(0) var source: texture_2d<f32>;
@group(0) @binding(1) var<storage, read_write> bins: array<atomic<u32>>;
@group(0) @binding(2) var<uniform> dims: Dims;

var<workgroup> tile: array<atomic<u32>, SLOTS>;

// The bin a scene-linear value belongs in.
//
// Zero and below land in bin 0: `log2` of them is not a number, and a
// photosite at or under the black level has no signal to place on a
// logarithmic axis anyway — it is the bottom of the scale, which is where the
// bottom bin is.
//
// Above 1.0 lands in bin 255. Such values exist: the demosaic clamps each
// *photosite* at the white level, but the Malvar correction term can overshoot
// upward between two clipped ones, so the interpolated output can pass 1.0 by
// a little. That is an artefact of the reconstruction and not headroom above
// saturation, and folding it into the top bin says exactly that.
fn bin_of(v: f32) -> u32 {
    if (v <= 0.0) {
        return 0u;
    }
    // Stops below saturation, in sixteenths, floored. `-log2` because the axis
    // increases downward from 1.0 and a bin index increases upward.
    let steps = floor(-log2(v) * BINS_PER_STOP);
    return (BINS - 1u) - u32(clamp(steps, 0.0, f32(BINS - 1u)));
}

@compute @workgroup_size(16, 16, 1)
fn main(
    @builtin(global_invocation_id) gid: vec3<u32>,
    @builtin(local_invocation_index) lid: u32,
) {
    for (var i = lid; i < SLOTS; i = i + THREADS) {
        atomicStore(&tile[i], 0u);
    }
    workgroupBarrier();

    // Guarded rather than dispatched exactly: the workgroup is 16x16 and an
    // image is not, so the last row and column of workgroups run off the edge.
    if (gid.x < dims.width && gid.y < dims.height) {
        let texel = textureLoad(source, vec2<i32>(i32(gid.x), i32(gid.y)), 0);
        let r = texel.r;
        let g = texel.g;
        let b = texel.b;

        // The fourth series is the **brightest** channel at each pixel, where
        // the display histogram's fourth series is Rec.709 luma. Luma would be
        // meaningless here: these are camera-native values with the as-shot
        // white balance still un-applied, so red and blue sit a stop or more
        // below green on a neutral subject and any weighted sum of them is a
        // number about nothing. The brightest channel, by contrast, is exactly
        // the quantity the headroom question is about — it is the one that
        // reaches saturation first and decides whether the pixel is
        // recoverable.
        let brightest = max(r, max(g, b));

        atomicAdd(&tile[bin_of(r)], 1u);
        atomicAdd(&tile[BINS + bin_of(g)], 1u);
        atomicAdd(&tile[2u * BINS + bin_of(b)], 1u);
        atomicAdd(&tile[3u * BINS + bin_of(brightest)], 1u);

        // Any channel, not all three, exactly as the display histogram counts
        // it: a saturated red photosite has no gradation left in it however
        // much headroom green and blue still hold, and a sunset or a red
        // jersey is precisely the subject that clips one channel first.
        if (brightest >= SATURATION) {
            atomicAdd(&tile[SATURATED], 1u);
        }
        // And at the other end: a channel that reached the black level has no
        // signal, only the read noise the normalisation clamped away.
        if (min(r, min(g, b)) <= 0.0) {
            atomicAdd(&tile[AT_BLACK], 1u);
        }
    }
    workgroupBarrier();

    for (var i = lid; i < SLOTS; i = i + THREADS) {
        let count = atomicLoad(&tile[i]);
        // Most bins of most workgroups are empty — a 16x16 tile can touch 256
        // of 1026 slots at the very most, and usually far fewer.
        if (count != 0u) {
            atomicAdd(&bins[i], count);
        }
    }
}