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

// TRACES: FR-DSP-7
// A reduction of the display frame into bin counts, run where the pixels are.
//
// FR-DSP-7 does not merely permit this shape, it names it: "these derive from a
// GPU-side reduction into a small buffer", because the alternative — dragging
// the frame back across the bus to count it — is the per-frame round-trip
// ARCH §6.1 forbids and the bottleneck darktable documents. What leaves the
// device here is 4104 bytes whatever the image size.
//
// **Counted per workgroup first, then merged.** A photograph is not noise: a
// clear sky puts tens of thousands of adjacent pixels in one bin, and having
// every invocation contend for that single global atomic serialises the whole
// dispatch. Each workgroup therefore tallies its own 256 pixels into workgroup
// memory — where the atomic is cheap and the contention is between 256 threads
// rather than two million — and contributes one add per non-empty bin at the
// end. Four kilobytes of workgroup storage against a 16 KB floor.

const BINS: u32 = 256u;
// Two counters past the four channels: how many pixels clip at each end. They
// live in the same buffer because they are gathered from the same texel read
// and would otherwise need a second reduction to answer a question the first
// one already had the data for.
const CLIPPED_HIGH: u32 = 4u * BINS;
const CLIPPED_LOW: u32 = 4u * BINS + 1u;
const SLOTS: u32 = 4u * BINS + 2u;
// 16x16. Stated as a constant because the clear and merge loops below stride by
// it, and a workgroup size that disagreed would leave slots uncleared.
const THREADS: u32 = 256u;

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 8-bit level a sampled texel came from.
///
/// The source is `Rgba8Unorm`, so the sampler hands back exactly n/255 and this
/// recovers n. `floor(x + 0.5)` rather than `round`, which ties to even in WGSL
/// and away from zero in Rust — a difference invisible except on an exact tie,
/// which is precisely the kind of disagreement that makes a cross-check against
/// a CPU reference fail once in a thousand runs and look like flakiness.
fn level(v: f32) -> u32 {
    return u32(clamp(floor(v * 255.0 + 0.5), 0.0, 255.0));
}

@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 = level(texel.r);
        let g = level(texel.g);
        let b = level(texel.b);

        // Rec.709 luma in 8.8 fixed point. 54 + 183 + 19 is exactly 256, so the
        // weights sum to unity and the shift can never carry past 255.
        //
        // Integer rather than float on purpose (ARCH §6.13): a float weighting
        // is reproducible only to within the vendor's rounding, and the whole
        // value of the CPU cross-check in the tests is that it is exact.
        //
        // Weighted on the *encoded* values, not on linear light. That is what
        // every histogram a photographer has read is: the axis is the output
        // level, so a mid-grey has to sit in the middle of it.
        let y = (54u * r + 183u * g + 19u * b) >> 8u;

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

        // Any channel, not all three: a blown red channel is detail that is
        // gone, whatever green and blue still hold. Counting only neutral white
        // would stay silent on exactly the saturated highlight — a sunset, a
        // red jersey — that clips first and recovers worst.
        if (r == 255u || g == 255u || b == 255u) {
            atomicAdd(&tile[CLIPPED_HIGH], 1u);
        }
        if (r == 0u || g == 0u || b == 0u) {
            atomicAdd(&tile[CLIPPED_LOW], 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);
        }
    }
}