diff --git a/crates/gpu/src/fx_noise.rs b/crates/gpu/src/fx_noise.rs index 1a20aa453..e3bec6fcd 100644 --- a/crates/gpu/src/fx_noise.rs +++ b/crates/gpu/src/fx_noise.rs @@ -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 } diff --git a/crates/gpu/src/shaders/fx_noise.wgsl b/crates/gpu/src/shaders/fx_noise.wgsl index 5e68e10c1..24b0d1e0b 100644 --- a/crates/gpu/src/shaders/fx_noise.wgsl +++ b/crates/gpu/src/shaders/fx_noise.wgsl @@ -86,9 +86,13 @@ fn fxn_median(@builtin(global_invocation_id) gid: vec3) { // 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 median_hist: array; +// 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) { let dims = out_dims(); let y = i32(gid.x); @@ -100,14 +104,13 @@ fn fxn_median_huang(@builtin(global_invocation_id) gid: vec3) { 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; for (var dy = -r; dy <= r; dy++) { for (var dx = -r; dx <= r; dx++) { let v = vec4(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) { let o = vec4(tex_get_clamped(src, x - 1 - r, y + dy)); let n = vec4(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) { 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; } }