Speed up retouch blur convolution with exact-output coverage (#4511)

Co-authored-by: jonathnwxh <[email protected]>

crates/algo/src/retouch.rs

@@ -117,49 +117,53 @@ pub fn sponge(c: [f32; 3], amount: f32, saturate: bool, vibrance: bool) -> [f32;
 pub fn local_blur(data: &[f32], w: usize, h: usize, ch: usize, alpha: Option<usize>, win: (usize, usize, usize, usize), sigma: f32) -> Vec<f32> {
     let (x0, y0, x1, y1) = win;
     let (ow, oh) = (x1 - x0, y1 - y0);
+    if ow == 0 || oh == 0 || ch == 0 {
+        return Vec::new();
+    }
     let r = (sigma * 3.0).ceil().max(1.0) as i32;
     let k: Vec<f32> = (-r..=r).map(|i| (-((i * i) as f32) / (2.0 * sigma * sigma).max(1e-6)).exp()).collect();
     let ks: f32 = k.iter().sum();
     let k: Vec<f32> = k.into_iter().map(|v| v / ks).collect();
-    let get = |x: i32, y: i32, out: &mut [f32]| {
-        let (xx, yy) = (x.clamp(0, w as i32 - 1) as usize, y.clamp(0, h as i32 - 1) as usize);
-        let i = (yy * w + xx) * ch;
-        out.copy_from_slice(&data[i..i + ch]);
-        if let Some(a) = alpha {
-            let av = out[a];
-            for (c, v) in out.iter_mut().enumerate() {
-                if c != a {
-                    *v *= av;
+    // Horizontal pass over the window rows extended vertically by r.
+    let th = oh + 2 * r as usize;
+    let stride = ow * ch;
+    let mut tmp = vec![0.0f32; stride * th];
+    let mut row = vec![0.0f32; (ow + 2 * r as usize) * ch];
+    for ty in 0..th {
+        let y = (y0 as i32 + ty as i32 - r).clamp(0, h as i32 - 1) as usize;
+        // Clamp and premultiply each source pixel once, rather than once for every kernel tap.
+        for (rx, px) in row.chunks_exact_mut(ch).enumerate() {
+            let x = (x0 as i32 + rx as i32 - r).clamp(0, w as i32 - 1) as usize;
+            let i = (y * w + x) * ch;
+            px.copy_from_slice(&data[i..i + ch]);
+            if let Some(a) = alpha {
+                let av = px[a];
+                for (c, v) in px.iter_mut().enumerate() {
+                    if c != a {
+                        *v *= av;
+                    }
                 }
             }
         }
-    };
-    // Horizontal pass over the window rows extended vertically by r.
-    let th = oh + 2 * r as usize;
-    let mut tmp = vec![0.0f32; ow * th * ch];
-    let mut px = vec![0.0f32; ch];
-    for ty in 0..th {
-        let y = y0 as i32 + ty as i32 - r;
-        for ox in 0..ow {
-            let o = (ty * ow + ox) * ch;
-            for (i, kv) in k.iter().enumerate() {
-                get(x0 as i32 + ox as i32 + i as i32 - r, y, &mut px);
-                for c in 0..ch {
-                    tmp[o + c] += px[c] * kv;
-                }
+        let dst = &mut tmp[ty * stride..(ty + 1) * stride];
+        // Contiguous accumulation can be vectorized for any channel layout. Each output still
+        // sums the same products in the same kernel order as the scalar convolution.
+        for (i, kv) in k.iter().enumerate() {
+            for (d, s) in dst.iter_mut().zip(&row[i * ch..i * ch + stride]) {
+                *d += s * kv;
             }
         }
     }
     let mut out = vec![0.0f32; ow * oh * ch];
     for oy in 0..oh {
+        let dst = &mut out[oy * stride..(oy + 1) * stride];
+        for (i, kv) in k.iter().enumerate() {
+            for (d, s) in dst.iter_mut().zip(&tmp[(oy + i) * stride..(oy + i + 1) * stride]) {
+                *d += s * kv;
+            }
+        }
         for ox in 0..ow {
             let o = (oy * ow + ox) * ch;
-            for (i, kv) in k.iter().enumerate() {
-                let s = ((oy + i) * ow + ox) * ch;
-                for c in 0..ch {
-                    out[o + c] += tmp[s + c] * kv;
-                }
-            }
             if let Some(a) = alpha {
                 let av = out[o + a];
                 for c in 0..ch {

crates/algo/tests/retouch_blur_exact.rs

@@ -0,0 +1,90 @@
+//! Frozen scalar convolution from main 5dac5c6, used as an exact-output oracle.
+use photocraft_algo::retouch::local_blur;
+
+#[test]
+fn matches_scalar_blur_bit_for_bit_across_layouts_edges_and_hdr() {
+    for (w, h) in [(1, 1), (3, 7), (37, 29)] {
+        for ch in 1..=5 {
+            for alpha in std::iter::once(None).chain((0..ch).map(Some)) {
+                let data: Vec<f32> = (0..w * h * ch)
+                    .map(|i| if Some(i % ch) == alpha { [0.0, 1e-8, 0.25, 0.75, 1.0][(i / ch) % 5] } else { ((i * 73 + 19) % 113) as f32 / 17.0 - 2.0 })
+                    .collect();
+                for win in [(0, 0, w, h), (w / 2, h / 2, w, h), (0, 0, 1, 1), (0, 0, 0, h)] {
+                    for sigma in [0.0, 0.25, 1.0, 3.0, 4.0, 12.0] {
+                        let expected = scalar_local_blur(&data, w, h, ch, alpha, win, sigma);
+                        let actual = local_blur(&data, w, h, ch, alpha, win, sigma);
+                        assert_eq!(actual.len(), expected.len());
+                        for (i, (a, b)) in actual.iter().zip(&expected).enumerate() {
+                            assert_eq!(a.to_bits(), b.to_bits(), "{w}x{h} ch={ch} alpha={alpha:?} win={win:?} sigma={sigma} sample={i}");
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[test]
+fn empty_channel_layout_remains_empty() {
+    assert!(local_blur(&[], 5, 5, 0, None, (0, 0, 5, 5), 4.0).is_empty());
+}
+
+#[allow(clippy::too_many_arguments)]
+fn scalar_local_blur(data: &[f32], w: usize, h: usize, ch: usize, alpha: Option<usize>, win: (usize, usize, usize, usize), sigma: f32) -> Vec<f32> {
+    let (x0, y0, x1, y1) = win;
+    let (ow, oh) = (x1 - x0, y1 - y0);
+    let r = (sigma * 3.0).ceil().max(1.0) as i32;
+    let k: Vec<f32> = (-r..=r).map(|i| (-((i * i) as f32) / (2.0 * sigma * sigma).max(1e-6)).exp()).collect();
+    let ks: f32 = k.iter().sum();
+    let k: Vec<f32> = k.into_iter().map(|v| v / ks).collect();
+    let get = |x: i32, y: i32, out: &mut [f32]| {
+        let (xx, yy) = (x.clamp(0, w as i32 - 1) as usize, y.clamp(0, h as i32 - 1) as usize);
+        let i = (yy * w + xx) * ch;
+        out.copy_from_slice(&data[i..i + ch]);
+        if let Some(a) = alpha {
+            let av = out[a];
+            for (c, v) in out.iter_mut().enumerate() {
+                if c != a {
+                    *v *= av;
+                }
+            }
+        }
+    };
+    // Horizontal pass over the window rows extended vertically by r.
+    let th = oh + 2 * r as usize;
+    let mut tmp = vec![0.0f32; ow * th * ch];
+    let mut px = vec![0.0f32; ch];
+    for ty in 0..th {
+        let y = y0 as i32 + ty as i32 - r;
+        for ox in 0..ow {
+            let o = (ty * ow + ox) * ch;
+            for (i, kv) in k.iter().enumerate() {
+                get(x0 as i32 + ox as i32 + i as i32 - r, y, &mut px);
+                for c in 0..ch {
+                    tmp[o + c] += px[c] * kv;
+                }
+            }
+        }
+    }
+    let mut out = vec![0.0f32; ow * oh * ch];
+    for oy in 0..oh {
+        for ox in 0..ow {
+            let o = (oy * ow + ox) * ch;
+            for (i, kv) in k.iter().enumerate() {
+                let s = ((oy + i) * ow + ox) * ch;
+                for c in 0..ch {
+                    out[o + c] += tmp[s + c] * kv;
+                }
+            }
+            if let Some(a) = alpha {
+                let av = out[o + a];
+                for c in 0..ch {
+                    if c != a {
+                        out[o + c] = if av > 1e-6 { out[o + c] / av } else { data[((y0 + oy) * w + x0 + ox) * ch + c] };
+                    }
+                }
+            }
+        }
+    }
+    out
+}

crates/engine/examples/bench_blur_tool.rs

@@ -0,0 +1,62 @@
+//! Reproducible soft Blur strokes on a 24 MP striped document. Run in release mode.
+//! Setup, Undo and the output fingerprint are excluded from the timing.
+use std::time::Instant;
+
+use photocraft_engine::{EngineError, Session};
+use photocraft_geom::Rect;
+use serde_json::json;
+
+fn run() -> Result<(), Box<dyn std::error::Error>> {
+    for depth in [8, 16, 32] {
+        let mut session = Session::new();
+        session.execute("file.new", json!({"width":6000,"height":4000,"depth":depth,"mode":"rgb"}))?;
+        session.execute("layer.new.layer", json!({}))?;
+        session.edit("stripes", |doc, active| {
+            let id = active.ok_or_else(|| EngineError::Other("no layer".into()))?;
+            let surface = doc.layer_mut(id).and_then(|l| l.surface_mut()).ok_or_else(|| EngineError::Other("no surface".into()))?;
+            for x in (0..6000).step_by(64) {
+                let color = if x % 128 == 0 { [0.9, 0.3, 0.1, 1.0] } else { [0.1, 0.4, 0.8, 1.0] };
+                surface.fill_rect(Rect::new(x, 0, (x + 64).min(6000), 4000), &color);
+            }
+            Ok(())
+        })?;
+        for size in [100, 500, 1000] {
+            let params = json!({"points":[[2500,2000],[3500,2000]],"size":size,"hardness":20,"strength":50,"spacing":25});
+            let mut times = Vec::new();
+            let mut fingerprint = 0u64;
+            for sample in 0..6 {
+                let start = Instant::now();
+                session.execute("paint.blur", params.clone())?;
+                let ms = start.elapsed().as_secs_f64() * 1000.0;
+                if sample > 0 {
+                    times.push(ms);
+                }
+                if sample == 5 {
+                    let st = session.active().ok_or("no document")?;
+                    let surface = st.doc.layer(st.active_layer.ok_or("no layer")?).and_then(|l| l.surface()).ok_or("no surface")?;
+                    for y in (1500..2500).step_by(17) {
+                        for x in (2000..4000).step_by(13) {
+                            for channel in surface.rgba(x, y) {
+                                fingerprint = fingerprint.wrapping_mul(16777619) ^ u64::from(channel.to_bits());
+                            }
+                        }
+                    }
+                }
+                session.execute("edit.undo", json!({}))?;
+            }
+            times.sort_by(f64::total_cmp);
+            println!(
+                "depth={depth} size={size} median_ms={:.3} min_ms={:.3} max_ms={:.3} sampled_fingerprint={fingerprint:016x}",
+                times[2], times[0], times[4]
+            );
+        }
+    }
+    Ok(())
+}
+
+fn main() {
+    if let Err(error) = run() {
+        eprintln!("{error}");
+        std::process::exit(1);
+    }
+}