M0.37: specialize median histograms for software Vulkan

crates/gpu/src/context.rs

@@ -351,6 +351,7 @@ pub struct GpuContext {
     pub queue: wgpu::Queue,
     pub(crate) name: String,
     pub(crate) max_dim: u32,
+    pub(crate) median_kernel: crate::fx_noise::MedianKernel,
     bgl: wgpu::BindGroupLayout,
     bgl_ext: wgpu::BindGroupLayout,
     pipelines: HashMap<&'static str, wgpu::ComputePipeline>,
@@ -481,6 +482,9 @@ impl GpuContext {
     /// [`Self::device_failure_notifier`] for immediate readback cancellation.
     pub fn new(adapter: &wgpu::Adapter, device: wgpu::Device, queue: wgpu::Queue) -> Result<GpuContext, String> {
         let info = adapter.get_info();
+        let median_kernel = crate::fx_noise::MedianKernel::for_adapter(info.backend, info.device_type);
+        #[cfg(test)]
+        let median_kernel = median_kernel.test_override()?;
         if !adapter.get_downlevel_capabilities().flags.contains(wgpu::DownlevelFlags::COMPUTE_SHADERS) {
             return Err(format!("{} ({:?}): no compute shaders", info.name, info.backend));
         }
@@ -500,6 +504,7 @@ impl GpuContext {
             ("max_compute_workgroup_size_x", l.max_compute_workgroup_size_x, 64),
             ("max_compute_workgroup_size_y", l.max_compute_workgroup_size_y, 16),
             ("max_compute_invocations_per_workgroup", l.max_compute_invocations_per_workgroup, 256),
+            ("max_compute_workgroup_storage_size", l.max_compute_workgroup_storage_size, median_kernel.workgroup_storage()),
         ];
         if let Some((limit, have, need)) = needs.iter().find(|(_, have, need)| have < need) {
             return Err(format!("{} ({:?}): device limit {limit} is {have}, needs {need}", info.name, info.backend));
@@ -525,6 +530,7 @@ impl GpuContext {
             include_str!("shaders/fx_key.wgsl"),
             include_str!("shaders/fx_stylize.wgsl"),
             include_str!("shaders/fx_noise.wgsl"),
+            median_kernel.source(),
             include_str!("shaders/fx_tone.wgsl"),
             include_str!("shaders/fx_warp.wgsl"),
             include_str!("shaders/fx_extra.wgsl"),
@@ -706,6 +712,7 @@ impl GpuContext {
             queue,
             name,
             max_dim,
+            median_kernel,
             bgl,
             bgl_ext,
             pipelines,

crates/gpu/src/fx_noise.rs

@@ -73,6 +73,53 @@ pub(crate) const IDS: &[&str] = &[
 /// not its square.
 pub(crate) const MEDIAN_BISECT_R: usize = 4;
 
+/// Physical GPUs retain the original parallel private histogram. Software Vulkan
+/// uses bounded workgroup storage to avoid its expensive first-use private-array
+/// lowering. This trades some warmed software throughput for bounded first use.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) enum MedianKernel {
+    Private64,
+    Software2,
+}
+
+impl MedianKernel {
+    pub(crate) fn for_adapter(backend: wgpu::Backend, device_type: wgpu::DeviceType) -> Self {
+        if backend == wgpu::Backend::Vulkan && device_type == wgpu::DeviceType::Cpu { Self::Software2 } else { Self::Private64 }
+    }
+
+    pub(crate) fn rows_per_workgroup(self) -> u32 {
+        match self {
+            Self::Private64 => 64,
+            Self::Software2 => 2,
+        }
+    }
+
+    pub(crate) fn workgroup_storage(self) -> u32 {
+        match self {
+            Self::Private64 => 0,
+            Self::Software2 => 16_384,
+        }
+    }
+
+    pub(crate) fn source(self) -> &'static str {
+        match self {
+            Self::Private64 => include_str!("shaders/median_private.wgsl"),
+            Self::Software2 => include_str!("shaders/median_software.wgsl"),
+        }
+    }
+
+    /// Separate test processes can qualify both implementations against existing oracles.
+    #[cfg(test)]
+    pub(crate) fn test_override(self) -> Result<Self, String> {
+        match std::env::var("EFFECTCRAFT_TEST_MEDIAN_KERNEL") {
+            Err(std::env::VarError::NotPresent) => Ok(self),
+            Ok(value) if value == "private64" => Ok(Self::Private64),
+            Ok(value) if value == "software2" => Ok(Self::Software2),
+            value => Err(format!("invalid EFFECTCRAFT_TEST_MEDIAN_KERNEL {value:?}; expected private64 or software2")),
+        }
+    }
+}
+
 /// Run effect `id` (one of [`IDS`]); `None` = this parameter combination runs on the CPU.
 pub(crate) fn apply(e: &mut Enc, id: &str, ctx: &EffectCtx, b: GBuf) -> Option<GBuf> {
     match id {
@@ -166,7 +213,8 @@ 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, q.width.div_ceil(seg)));
+    let groups = (q.height.div_ceil(e.g.median_kernel.rows_per_workgroup()), q.width.div_ceil(seg));
+    e.dispatch("fxn_median_huang", &with_u0([r as u32, seg, 0, 0]), &q, None, &out, Some(&levels), groups);
     out
 }
 

crates/gpu/src/shaders/fx_noise.wgsl

@@ -82,72 +82,6 @@ fn fxn_median(@builtin(global_invocation_id) gid: vec3<u32>) {
     textureStore(out, p, vec4<f32>(data[level.x], data[level.y], data[level.z], data[level.w]));
 }
 
-// The same median for large radii, as the CPU computes it: a sliding 512-bin histogram per
-// 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.
-// 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(1, 1)
-fn fxn_median_huang(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let dims = out_dims();
-    let y = i32(gid.x);
-    let seg = i32(P.u[0].y);
-    let x0 = i32(gid.y) * seg;
-    if (y >= dims.y || x0 >= dims.x) {
-        return;
-    }
-    let x1 = min(x0 + seg, dims.x);
-    let r = i32(P.u[0].x);
-    let half = u32((2 * r + 1) * (2 * r + 1) / 2);
-    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));
-            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.
-    var med = vec4<u32>(0u);
-    var lt = vec4<u32>(0u);
-    for (var x = x0; x < x1; x++) {
-        if (x > x0) {
-            for (var dy = -r; dy <= r; dy++) {
-                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++) {
-                    median_hist[c * 512u + o[c]] -= 1u;
-                    if (o[c] < med[c]) {
-                        lt[c] -= 1u;
-                    }
-                    median_hist[c * 512u + n[c]] += 1u;
-                    if (n[c] < med[c]) {
-                        lt[c] += 1u;
-                    }
-                }
-            }
-        }
-        for (var c = 0u; c < 4u; c++) {
-            let base = c * 512u;
-            while (lt[c] > half) {
-                med[c] -= 1u;
-                lt[c] -= median_hist[base + med[c]];
-            }
-            while (lt[c] + median_hist[base + med[c]] <= half) {
-                lt[c] += median_hist[base + med[c]];
-                med[c] += 1u;
-            }
-        }
-        textureStore(out, vec2<i32>(x, y), vec4<f32>(data[med.x], data[med.y], data[med.z], data[med.w]));
-    }
-}
-
 // noise::median_px.
 fn fxn_median_px(o: vec4<f32>, m: vec4<f32>, on_alpha: bool) -> vec4<f32> {
     if (on_alpha) {

crates/gpu/src/shaders/median_private.wgsl

@@ -0,0 +1,63 @@
+// The same median for large radii, as the CPU computes it: a sliding 512-bin histogram per
+// 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.
+// data = the 512 CPU-rounded level / 511 values.
+@compute @workgroup_size(64, 1)
+fn fxn_median_huang(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let dims = out_dims();
+    let y = i32(gid.x);
+    let seg = i32(P.u[0].y);
+    let x0 = i32(gid.y) * seg;
+    if (y >= dims.y || x0 >= dims.x) {
+        return;
+    }
+    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;
+        }
+    }
+    // Per channel: the median level and the count of values below it.
+    var med = vec4<u32>(0u);
+    var lt = vec4<u32>(0u);
+    for (var x = x0; x < x1; x++) {
+        if (x > x0) {
+            for (var dy = -r; dy <= r; dy++) {
+                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;
+                    if (o[c] < med[c]) {
+                        lt[c] -= 1u;
+                    }
+                    hist[c * 512u + n[c]] += 1u;
+                    if (n[c] < med[c]) {
+                        lt[c] += 1u;
+                    }
+                }
+            }
+        }
+        for (var c = 0u; c < 4u; c++) {
+            let base = c * 512u;
+            while (lt[c] > half) {
+                med[c] -= 1u;
+                lt[c] -= hist[base + med[c]];
+            }
+            while (lt[c] + hist[base + med[c]] <= half) {
+                lt[c] += hist[base + med[c]];
+                med[c] += 1u;
+            }
+        }
+        textureStore(out, vec2<i32>(x, y), vec4<f32>(data[med.x], data[med.y], data[med.z], data[med.w]));
+    }
+}
+

crates/gpu/src/shaders/median_software.wgsl

@@ -0,0 +1,65 @@
+// The same median for large radii, as the CPU computes it: a sliding 512-bin histogram per
+// 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.
+// Each invocation owns one 8 KiB slice of the workgroup's zero-initialized histogram.
+// Two slices fit the guaranteed 16 KiB workgroup storage limit. No bins are shared.
+var<workgroup> median_hist: array<array<u32, 2048>, 2>;
+// u[0] = (r, segment length); dispatched over (rows, segments) with 2 × 1 workgroups.
+// data = the 512 CPU-rounded level / 511 values.
+@compute @workgroup_size(2, 1)
+fn fxn_median_huang(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(local_invocation_index) lane: u32) {
+    let dims = out_dims();
+    let y = i32(gid.x);
+    let seg = i32(P.u[0].y);
+    let x0 = i32(gid.y) * seg;
+    if (y >= dims.y || x0 >= dims.x) {
+        return;
+    }
+    let x1 = min(x0 + seg, dims.x);
+    let r = i32(P.u[0].x);
+    let half = u32((2 * r + 1) * (2 * r + 1) / 2);
+    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));
+            median_hist[lane][v.x] += 1u;
+            median_hist[lane][512u + v.y] += 1u;
+            median_hist[lane][1024u + v.z] += 1u;
+            median_hist[lane][1536u + v.w] += 1u;
+        }
+    }
+    // Per channel: the median level and the count of values below it.
+    var med = vec4<u32>(0u);
+    var lt = vec4<u32>(0u);
+    for (var x = x0; x < x1; x++) {
+        if (x > x0) {
+            for (var dy = -r; dy <= r; dy++) {
+                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++) {
+                    median_hist[lane][c * 512u + o[c]] -= 1u;
+                    if (o[c] < med[c]) {
+                        lt[c] -= 1u;
+                    }
+                    median_hist[lane][c * 512u + n[c]] += 1u;
+                    if (n[c] < med[c]) {
+                        lt[c] += 1u;
+                    }
+                }
+            }
+        }
+        for (var c = 0u; c < 4u; c++) {
+            let base = c * 512u;
+            while (lt[c] > half) {
+                med[c] -= 1u;
+                lt[c] -= median_hist[lane][base + med[c]];
+            }
+            while (lt[c] + median_hist[lane][base + med[c]] <= half) {
+                lt[c] += median_hist[lane][base + med[c]];
+                med[c] += 1u;
+            }
+        }
+        textureStore(out, vec2<i32>(x, y), vec4<f32>(data[med.x], data[med.y], data[med.z], data[med.w]));
+    }
+}
+

crates/gpu/src/tests.rs

@@ -74,7 +74,7 @@ pub(crate) fn gpu() -> Option<&'static Gpu> {
         }
         let g = Gpu::try_headless();
         match &g {
-            Ok(g) => eprintln!("effectcraft-gpu tests: adapter {}", g.ctx.name),
+            Ok(g) => eprintln!("effectcraft-gpu tests: adapter {}, median {:?}", g.ctx.name, g.ctx.median_kernel),
             Err(e) => eprintln!("effectcraft-gpu tests: no GPU ({e}), skipping GPU comparisons"),
         }
         g

crates/gpu/src/tests_fx_noise.rs

@@ -22,6 +22,24 @@ fn pt(x: f64, y: f64) -> Value {
     Value::Vec2([x, y])
 }
 
+#[test]
+fn median_kernel_selection_preserves_physical_gpu_path() {
+    use crate::fx_noise::MedianKernel;
+    for backend in [wgpu::Backend::Vulkan, wgpu::Backend::Metal, wgpu::Backend::Dx12, wgpu::Backend::Gl, wgpu::Backend::BrowserWebGpu] {
+        for device_type in
+            [wgpu::DeviceType::DiscreteGpu, wgpu::DeviceType::IntegratedGpu, wgpu::DeviceType::VirtualGpu, wgpu::DeviceType::Cpu, wgpu::DeviceType::Other]
+        {
+            let selected = MedianKernel::for_adapter(backend, device_type);
+            let expected =
+                if backend == wgpu::Backend::Vulkan && device_type == wgpu::DeviceType::Cpu { MedianKernel::Software2 } else { MedianKernel::Private64 };
+            assert_eq!(selected, expected, "{backend:?} {device_type:?}");
+        }
+    }
+    assert_eq!(MedianKernel::Private64.rows_per_workgroup(), 64);
+    assert_eq!(MedianKernel::Software2.rows_per_workgroup(), 2);
+    assert_eq!(MedianKernel::Software2.workgroup_storage(), 16_384);
+}
+
 #[test]
 fn medians() {
     effect_case("ec.noise.median", &[("radius", n(2.0))]);