M0.37: keep median histograms in per-workgroup memory

This commit is contained in:
skylar taylor-barrick
2026-10-26 18:49:54 +01:00
parent af354e49e8
commit 5d60eb4449
2 changed files with 16 additions and 13 deletions
+1 -1
View File
@@ -166,7 +166,7 @@ pub(crate) fn median_image(e: &mut Enc, img: &GpuImage, r: usize, straight: bool
// sliding along the rest.
let seg = (4 * r as u32).clamp(32, 512).min(q.width);
let out = e.scratch(q.width, q.height);
e.dispatch("fxn_median_huang", &with_u0([r as u32, seg, 0, 0]), &q, None, &out, Some(&levels), (q.height.div_ceil(64), q.width.div_ceil(seg)));
e.dispatch("fxn_median_huang", &with_u0([r as u32, seg, 0, 0]), &q, None, &out, Some(&levels), (q.height, q.width.div_ceil(seg)));
out
}
+15 -12
View File
@@ -86,9 +86,13 @@ fn fxn_median(@builtin(global_invocation_id) gid: vec3<u32>) {
// channel (Huang). One invocation per row segment of u[0].y pixels: it fills the histogram for
// the segment's first window, then slides right one column at a time (remove the column
// leaving, add the one entering) and moves each channel's median by the counts below it.
// u[0] = (r, segment length); dispatched over (rows, segments) with 64 × 1 workgroups.
// One invocation owns the workgroup's 8 KiB histogram. Keeping the indirectly indexed
// bins in workgroup memory avoids a large per-invocation private array. WGSL initializes
// this memory to zero before each workgroup; no invocation shares its bins with another.
var<workgroup> median_hist: array<u32, 2048>;
// u[0] = (r, segment length); dispatched over (rows, segments) with 1 × 1 workgroups.
// data = the 512 CPU-rounded level / 511 values.
@compute @workgroup_size(64, 1)
@compute @workgroup_size(1, 1)
fn fxn_median_huang(@builtin(global_invocation_id) gid: vec3<u32>) {
let dims = out_dims();
let y = i32(gid.x);
@@ -100,14 +104,13 @@ fn fxn_median_huang(@builtin(global_invocation_id) gid: vec3<u32>) {
let x1 = min(x0 + seg, dims.x);
let r = i32(P.u[0].x);
let half = u32((2 * r + 1) * (2 * r + 1) / 2);
var hist: array<u32, 2048>;
for (var dy = -r; dy <= r; dy++) {
for (var dx = -r; dx <= r; dx++) {
let v = vec4<u32>(tex_get_clamped(src, x0 + dx, y + dy));
hist[v.x] += 1u;
hist[512u + v.y] += 1u;
hist[1024u + v.z] += 1u;
hist[1536u + v.w] += 1u;
median_hist[v.x] += 1u;
median_hist[512u + v.y] += 1u;
median_hist[1024u + v.z] += 1u;
median_hist[1536u + v.w] += 1u;
}
}
// Per channel: the median level and the count of values below it.
@@ -119,11 +122,11 @@ fn fxn_median_huang(@builtin(global_invocation_id) gid: vec3<u32>) {
let o = vec4<u32>(tex_get_clamped(src, x - 1 - r, y + dy));
let n = vec4<u32>(tex_get_clamped(src, x + r, y + dy));
for (var c = 0u; c < 4u; c++) {
hist[c * 512u + o[c]] -= 1u;
median_hist[c * 512u + o[c]] -= 1u;
if (o[c] < med[c]) {
lt[c] -= 1u;
}
hist[c * 512u + n[c]] += 1u;
median_hist[c * 512u + n[c]] += 1u;
if (n[c] < med[c]) {
lt[c] += 1u;
}
@@ -134,10 +137,10 @@ fn fxn_median_huang(@builtin(global_invocation_id) gid: vec3<u32>) {
let base = c * 512u;
while (lt[c] > half) {
med[c] -= 1u;
lt[c] -= hist[base + med[c]];
lt[c] -= median_hist[base + med[c]];
}
while (lt[c] + hist[base + med[c]] <= half) {
lt[c] += hist[base + med[c]];
while (lt[c] + median_hist[base + med[c]] <= half) {
lt[c] += median_hist[base + med[c]];
med[c] += 1u;
}
}