From e93b30e60d33b44ab7e0da080ba06bcb91396277 Mon Sep 17 00:00:00 2001 From: Ayan <108534604+excited-raccoon@users.noreply.github.com> Date: Sat, 10 Oct 2026 04:34:26 +0500 Subject: [PATCH] Proximity Match: 1.5-3.8x faster stroke searches, smaller mask tables (#2101) Co-authored-by: Brandon Thomas --- crates/algo/src/inpaint.rs | 27 +- crates/algo/src/inpaint/proximity_integral.rs | 46 ++++ .../src/inpaint/proximity_integral/tests.rs | 238 ++++++++++++++++++ docs/proximity-match-performance.md | 71 ++++++ docs/roadmap.md | 6 + 5 files changed, 381 insertions(+), 7 deletions(-) create mode 100644 crates/algo/src/inpaint/proximity_integral.rs create mode 100644 crates/algo/src/inpaint/proximity_integral/tests.rs create mode 100644 docs/proximity-match-performance.md diff --git a/crates/algo/src/inpaint.rs b/crates/algo/src/inpaint.rs index 6b85d015..b777bf4d 100644 --- a/crates/algo/src/inpaint.rs +++ b/crates/algo/src/inpaint.rs @@ -20,6 +20,9 @@ use crate::poisson::membrane_fill; +mod proximity_integral; +use proximity_integral::HoleIntegral; + /// Below this many grid cells a level runs single-threaded (thread hand-off costs more than it saves). #[cfg(not(target_arch = "wasm32"))] const PAR_MIN: usize = 128 * 128; @@ -529,7 +532,9 @@ pub fn best_offset(w: usize, h: usize, ch: usize, img: &[f32], hole: &[bool], ri return None; } let ring = ring.max(1).min(extent as usize) as i32; - let hs = integral(w, h, hole); + // Every set cell lies in the hole bounds. Counts outside that rectangle are zero, + // so source rejection needs a table for the hole, not for the whole photo. + let hs = HoleIntegral::new(w, hole, (bx0, by0, bx1, by1))?; // Sample points: the hole's ring (dilated minus hole), subsampled for speed. let mut pts = Vec::new(); let (x0, y0) = (bx0 as i32 - ring, by0 as i32 - ring); @@ -539,35 +544,43 @@ pub fn best_offset(w: usize, h: usize, ch: usize, img: &[f32], hole: &[bool], ri if hole[y as usize * w + x as usize] { continue; } - if window_count(&hs, w, h, x - ring, y - ring, x + ring + 1, y + ring + 1) > 0 { + if hs.count(x - ring, y - ring, x + ring + 1, y + ring + 1) > 0 { pts.push((x, y)); } } } let stride = (pts.len() / 2000).max(1); let pts: Vec<(i32, i32)> = pts.into_iter().step_by(stride).collect(); - let eval = |dx: i32, dy: i32| -> Option { + let eval = |dx: i32, dy: i32, cutoff: f32| -> Option { let (hx0, hy0, hx1, hy1) = (x0 + dx, y0 + dy, x1 + dx, y1 + dy); let inside = hx0 >= 0 && hy0 >= 0 && hx1 <= w as i32 && hy1 <= h as i32; - if (dx, dy) == (0, 0) || !inside || window_count(&hs, w, h, hx0, hy0, hx1, hy1) != 0 { + if (dx, dy) == (0, 0) || !inside || hs.count(hx0, hy0, hx1, hy1) != 0 { return None; } let mut e = 0.0f32; - for &(x, y) in &pts { + let n = pts.len().max(1) as f32; + let penalty = 1.0 + 1e-3 * ((dx as f32).hypot(dy as f32) / max_radius as f32); + for (i, &(x, y)) in pts.iter().enumerate() { let (a, b) = ((y as usize * w + x as usize) * ch, ((y + dy) as usize * w + (x + dx) as usize) * ch); for c in 0..ch { let d = img[a + c] - img[b + c]; e += d * d; } + // SSD terms are nonnegative. The partial score is a lower bound on the + // final score; use the same division and penalty to keep float rounding + // and strict tie ordering unchanged. Check in batches to keep this cheap. + if i % 32 == 31 && cutoff.is_finite() && e / n * penalty >= cutoff { + return None; + } } // Mild preference for nearer sources on ties. - Some(e / pts.len().max(1) as f32 * (1.0 + 1e-3 * ((dx as f32).hypot(dy as f32) / max_radius as f32))) + Some(e / n * penalty) }; // Coarse grid over the search window, then a full-resolution refinement around the winner. let step = (max_radius / 24).max(1); let mut best: Option<((i32, i32), f32)> = None; let consider = |dx: i32, dy: i32, best: &mut Option<((i32, i32), f32)>| { - if let Some(e) = eval(dx, dy) + if let Some(e) = eval(dx, dy, best.map_or(f32::INFINITY, |(_, error)| error)) && best.is_none_or(|(_, be)| e < be) { *best = Some(((dx, dy), e)); diff --git a/crates/algo/src/inpaint/proximity_integral.rs b/crates/algo/src/inpaint/proximity_integral.rs new file mode 100644 index 00000000..7271b321 --- /dev/null +++ b/crates/algo/src/inpaint/proximity_integral.rs @@ -0,0 +1,46 @@ +//! Summed-area table restricted to the nonzero bounds of a healing hole. + +pub(super) struct HoleIntegral { + x0: i32, + y0: i32, + w: usize, + h: usize, + data: Vec, +} + +impl HoleIntegral { + pub(super) fn new(stride: usize, hole: &[bool], (x0, y0, x1, y1): (usize, usize, usize, usize)) -> Option { + let (w, h) = (x1.checked_sub(x0)?, y1.checked_sub(y0)?); + let row_stride = w.checked_add(1)?; + let len = row_stride.checked_mul(h.checked_add(1)?)?; + let mut data = Vec::new(); + data.try_reserve_exact(len).ok()?; + data.resize(len, 0u32); + for (y, sy) in (y0..y1).enumerate() { + let base = sy.checked_mul(stride)?; + let row = hole.get(base.checked_add(x0)?..base.checked_add(x1)?)?; + let mut count = 0u32; + for (x, &set) in row.iter().enumerate() { + count = count.checked_add(u32::from(set))?; + let above = *data.get(y * row_stride + x + 1)?; + *data.get_mut((y + 1) * row_stride + x + 1)? = above.checked_add(count)?; + } + } + Some(Self { x0: i32::try_from(x0).ok()?, y0: i32::try_from(y0).ok()?, w, h, data }) + } + + pub(super) fn count(&self, x0: i32, y0: i32, x1: i32, y1: i32) -> u32 { + super::window_count( + &self.data, + self.w, + self.h, + x0.saturating_sub(self.x0), + y0.saturating_sub(self.y0), + x1.saturating_sub(self.x0), + y1.saturating_sub(self.y0), + ) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/algo/src/inpaint/proximity_integral/tests.rs b/crates/algo/src/inpaint/proximity_integral/tests.rs new file mode 100644 index 00000000..5fab71df --- /dev/null +++ b/crates/algo/src/inpaint/proximity_integral/tests.rs @@ -0,0 +1,238 @@ +use super::*; +use crate::inpaint::{best_offset, hole_bbox, integral, window_count}; + +#[test] +fn cropped_counts_match_full_image_for_clipped_and_disjoint_windows() { + let (w, h) = (37, 29); + for kind in 0..5 { + let hole: Vec<_> = (0..w * h) + .map(|i| { + let (x, y) = (i % w, i / w); + match kind { + 0 => (11..16).contains(&x) && (8..14).contains(&y), + 1 => x < 3 && y < 4, + 2 => x >= w - 3 && y >= h - 4, + 3 => (x == 3 && y == 8) || (x == 31 && y == 25), + _ => (x * 37 + y * 71) % 17 == 0, + } + }) + .collect(); + let cropped = HoleIntegral::new(w, &hole, hole_bbox(w, h, &hole).unwrap()).unwrap(); + let full = integral(w, h, &hole); + for x0 in [-50, -1, 0, 3, 11, 15, 36, 40] { + for x1 in [-2, 0, 3, 13, 16, 37, 50] { + for y0 in [-30, -1, 0, 8, 13, 28, 35] { + for y1 in [-1, 0, 10, 14, 29, 40] { + assert_eq!(cropped.count(x0, y0, x1, y1), window_count(&full, w, h, x0, y0, x1, y1), "kind {kind}: {x0},{y0}..{x1},{y1}"); + } + } + } + } + } +} + +#[test] +fn table_size_depends_on_the_hole_bounds() { + let (w, h) = (1000, 700); + let hole: Vec<_> = (0..w * h).map(|i| (503..514).contains(&(i % w)) && (307..320).contains(&(i / w))).collect(); + let table = HoleIntegral::new(w, &hole, hole_bbox(w, h, &hole).unwrap()).unwrap(); + assert_eq!((table.w, table.h, table.data.len()), (11, 13, 12 * 14)); + assert_eq!(table.count(0, 0, 1000, 700), 143); + assert_eq!(table.count(i32::MIN, i32::MIN, i32::MAX, i32::MAX), 143); + assert!(HoleIntegral::new(w, &hole, (10, 5, 9, 6)).is_none()); + assert!(HoleIntegral::new(w, &[], (3, 4, 5, 6)).is_none()); + assert!(HoleIntegral::new(w, &[], (0, 0, usize::MAX, 2)).is_none()); +} + +#[test] +fn proximity_offsets_match_full_tables_at_several_depths_and_channels() { + let (w, h) = (61, 49); + for ch in [1, 3, 4, 5] { + for depth in [8, 16, 32] { + let img: Vec<_> = (0..w * h * ch) + .map(|i| { + let value = ((i * 101 + i / (w * ch) * 37) % 251) as f32 / 250.0; + match depth { + 8 => (value * 255.0).round() / 255.0, + 16 => (value * 65535.0).round() / 65535.0, + _ => value * 2.0 - 0.5, + } + }) + .collect(); + for kind in 0..5 { + let hole: Vec<_> = (0..w * h) + .map(|i| { + let (x, y) = (i % w, i / w); + match kind { + 0 => (27..32).contains(&x) && (20..25).contains(&y), + 1 => x < 2 && y < 3, + 2 => (x == 11 && y == 12) || (x == 38 && y == 30), + 3 => x == 30 && y == 24, + _ => false, + } + }) + .collect(); + for (ring, radius) in [(1, 5), (3, 17), (5, 30), (1, 49)] { + assert_eq!( + best_offset(w, h, ch, &img, &hole, ring, radius), + old_best_offset(w, h, ch, &img, &hole, ring, radius), + "ch {ch} depth {depth} hole {kind} ring {ring} radius {radius}" + ); + } + } + } + } +} + +#[test] +fn ties_nonfinite_samples_and_full_holes_keep_the_previous_result() { + let (w, h, ch) = (40, 35, 3); + let hole: Vec<_> = (0..w * h).map(|i| (18..22).contains(&(i % w)) && (16..20).contains(&(i / w))).collect(); + for value in [0.5, f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + let img = vec![value; w * h * ch]; + assert_eq!(best_offset(w, h, ch, &img, &hole, 3, 20), old_best_offset(w, h, ch, &img, &hole, 3, 20)); + } + let full = vec![true; w * h]; + assert_eq!(best_offset(w, h, ch, &vec![0.5; w * h * ch], &full, 3, 20), None); +} + +#[test] +#[ignore = "release comparison on 24 and 36 MP images"] +#[allow(clippy::assertions_on_constants)] +fn proximity_integral_release_comparison() { + use std::{hint::black_box, time::Instant}; + assert!(!cfg!(debug_assertions), "run with --release"); + for (w, h, diameters, radius) in [ + (182usize, 182usize, vec![50usize], 66i32), + (332, 332, vec![100], 116), + (632, 632, vec![200], 216), + (932, 932, vec![300], 316), + (6000, 4000, vec![32, 256], 512), + (7200, 5000, vec![32, 256], 512), + ] { + for diameter in diameters { + let ch = 3; + let img: Vec<_> = (0..w * h * ch) + .map(|i| { + let (x, y, c) = ((i / ch) % w, (i / ch) / w, i % ch); + ((x % 53 * 7 + y % 47 * 11 + c * 19) % 251) as f32 / 250.0 + }) + .collect(); + let hole: Vec<_> = (0..w * h) + .map(|i| { + let dx = (i % w) as f64 - (w / 2) as f64; + let dy = (i / w) as f64 - (h / 2) as f64; + dx * dx + dy * dy < (diameter as f64 * 0.5).powi(2) + }) + .collect(); + let bbox = hole_bbox(w, h, &hole).unwrap(); + let cropped = HoleIntegral::new(w, &hole, bbox).unwrap(); + let old = old_best_offset(w, h, ch, &img, &hole, (diameter / 8).clamp(3, 16), radius); + assert_eq!(best_offset(w, h, ch, &img, &hole, (diameter / 8).clamp(3, 16), radius), old); + let mut old_ms = Vec::new(); + let mut new_ms = Vec::new(); + for run in 0..5 { + let mut measure = |fast: bool| { + let start = Instant::now(); + let result = if fast { + best_offset(w, h, ch, black_box(&img), &hole, (diameter / 8).clamp(3, 16), radius) + } else { + old_best_offset(w, h, ch, black_box(&img), &hole, (diameter / 8).clamp(3, 16), radius) + }; + let ms = start.elapsed().as_secs_f64() * 1000.0; + assert_eq!(result, old); + if fast { + new_ms.push(ms); + } else { + old_ms.push(ms); + } + }; + measure(run % 2 == 0); + measure(run % 2 != 0); + } + old_ms.sort_by(f64::total_cmp); + new_ms.sort_by(f64::total_cmp); + println!( + "PROXIMITY {w}x{h} hole_diameter={diameter} old_ms={old_ms:?} new_ms={new_ms:?} speedup={:.2} full_table_bytes={} cropped_table_bytes={} offset={old:?}", + old_ms[2] / new_ms[2], + (w + 1) * (h + 1) * 4, + cropped.data.len() * 4 + ); + } + } +} +fn old_best_offset(w: usize, h: usize, ch: usize, img: &[f32], hole: &[bool], ring: usize, max_radius: i32) -> Option<(i32, i32)> { + let (bx0, by0, bx1, by1) = hole_bbox(w, h, hole)?; + // No valid translation can move farther than the source image's extent. + // Bound arbitrary public API inputs before constructing i32 search ranges: + // usize -> i32 truncation and signed increments otherwise overflow or hang. + let extent = w.max(h).min((i32::MAX / 4) as usize) as i32; + let max_radius = max_radius.max(0).min(extent); + if max_radius == 0 { + return None; + } + let ring = ring.max(1).min(extent as usize) as i32; + let hs = integral(w, h, hole); + // Sample points: the hole's ring (dilated minus hole), subsampled for speed. + let mut pts = Vec::new(); + let (x0, y0) = (bx0 as i32 - ring, by0 as i32 - ring); + let (x1, y1) = (bx1 as i32 + ring, by1 as i32 + ring); + for y in y0.max(0)..y1.min(h as i32) { + for x in x0.max(0)..x1.min(w as i32) { + if hole[y as usize * w + x as usize] { + continue; + } + if window_count(&hs, w, h, x - ring, y - ring, x + ring + 1, y + ring + 1) > 0 { + pts.push((x, y)); + } + } + } + let stride = (pts.len() / 2000).max(1); + let pts: Vec<(i32, i32)> = pts.into_iter().step_by(stride).collect(); + let eval = |dx: i32, dy: i32| -> Option { + let (hx0, hy0, hx1, hy1) = (x0 + dx, y0 + dy, x1 + dx, y1 + dy); + let inside = hx0 >= 0 && hy0 >= 0 && hx1 <= w as i32 && hy1 <= h as i32; + if (dx, dy) == (0, 0) || !inside || window_count(&hs, w, h, hx0, hy0, hx1, hy1) != 0 { + return None; + } + let mut e = 0.0f32; + for &(x, y) in &pts { + let (a, b) = ((y as usize * w + x as usize) * ch, ((y + dy) as usize * w + (x + dx) as usize) * ch); + for c in 0..ch { + let d = img[a + c] - img[b + c]; + e += d * d; + } + } + // Mild preference for nearer sources on ties. + Some(e / pts.len().max(1) as f32 * (1.0 + 1e-3 * ((dx as f32).hypot(dy as f32) / max_radius as f32))) + }; + // Coarse grid over the search window, then a full-resolution refinement around the winner. + let step = (max_radius / 24).max(1); + let mut best: Option<((i32, i32), f32)> = None; + let consider = |dx: i32, dy: i32, best: &mut Option<((i32, i32), f32)>| { + if let Some(e) = eval(dx, dy) + && best.is_none_or(|(_, be)| e < be) + { + *best = Some(((dx, dy), e)); + } + }; + let mut dy = -max_radius; + while dy <= max_radius { + let mut dx = -max_radius; + while dx <= max_radius { + consider(dx, dy, &mut best); + dx += step; + } + dy += step; + } + if step > 1 + && let Some(((bx, by), _)) = best + { + for dy in (by - step)..=(by + step) { + for dx in (bx - step)..=(bx + step) { + consider(dx, dy, &mut best); + } + } + } + best.map(|b| b.0) +} diff --git a/docs/proximity-match-performance.md b/docs/proximity-match-performance.md new file mode 100644 index 00000000..5abd9aa1 --- /dev/null +++ b/docs/proximity-match-performance.md @@ -0,0 +1,71 @@ +# Proximity Match search and memory + +Spot Healing's Proximity Match chooses a source displacement by comparing the ring around +the healing hole. Two changes reduce that search's work: + +- Source-validity counts use a summed-area table over the hole's nonzero bounding box. + Every excluded cell outside that box is zero, so translating and clipping queries gives + the same counts as a full-image table. This also covers separated holes and canvas edges. +- Each candidate's SSD is compared with the best score after every 32 ring samples. All + SSD terms are nonnegative, so the partial sum divided by the final sample count and + multiplied by the same distance penalty is a lower bound on the final score. A candidate + that reaches the best score cannot win the existing strict comparison. + +The sample order, accumulation order, distance penalty, candidate visit order, tie handling, +search radius and source-pixel copying stay the same. Nonfinite best scores disable pruning. +The cropped table checks allocation, arithmetic and source slices and returns None on failure. +The other summed-area tables used by content-aware completion and texture synthesis are unchanged. + +## Measurements + +Recorded 2026-10-10 on an AWS c7i.4xlarge, Intel Xeon Platinum 8488C, 16 vCPUs, 32 GiB RAM, +Rust 1.99.0, normal workspace release profile. This kernel is serial. The input is a synthetic +periodic RGB float texture and a circular hole; values are generated before timing. Each path +gets a warmup and five measured runs, alternating old/new order. Every run checks identical +source displacement against the old implementation. + +The four small regions follow the existing retouch benchmark's stroke-domain sizing: +`width = height = diameter + 2 * (diameter + 16)`, search radius `diameter + 16`. +Ring width is `clamp(diameter / 8, 3, 16)`. Large-image cases use search radius 512. +These are kernel calls, excluding brush rasterization, surface reads, Poisson blending, +document history and compositing. + +| Input region | Hole diameter | Before p50 | After p50 | Speedup | Old table | New table | +|---|---:|---:|---:|---:|---:|---:| +| 182 x 182 | 50 px | 2.924 ms | 0.765 ms | 3.82x | 133,956 B | 10,000 B | +| 332 x 332 | 100 px | 0.828 ms | 0.409 ms | 2.02x | 443,556 B | 40,000 B | +| 632 x 632 | 200 px | 1.959 ms | 1.272 ms | 1.54x | 1,602,756 B | 160,000 B | +| 932 x 932 | 300 px | 4.816 ms | 1.372 ms | 3.51x | 3,481,956 B | 360,000 B | +| 6000 x 4000 | 32 px | 67.862 ms | 14.672 ms | 4.63x | 96,040,004 B | 4,096 B | +| 6000 x 4000 | 256 px | 96.509 ms | 15.469 ms | 6.24x | 96,040,004 B | 262,144 B | +| 7200 x 5000 | 32 px | 99.693 ms | 21.658 ms | 4.60x | 144,048,804 B | 4,096 B | +| 7200 x 5000 | 256 px | 132.300 ms | 22.397 ms | 5.91x | 144,048,804 B | 262,144 B | + +Table storage alone falls by roughly 90–93 percent in these stroke regions. On a full +24–36 MP input it falls from 96–144 MB to 4 KiB for the 32 px hole, or 256 KiB for the +256 px hole. Those are table sizes, not total process memory: the input image, hole mask +and sampled ring still exist. The initial hole-bounds scan remains linear in image size. + +The engine already crops a normal healing stroke to a local region. The large-image cases +show the public kernel's scaling, not an additional 4.6–6.2x UI speedup. Source textures, +hole shapes, candidate validity and the quality of the current best match affect pruning; +these measurements do not establish a worst-case speedup or Photoshop image-quality parity. +A hole whose bounds cover the image still needs a full-size table. + +## Correctness and reproduction + +Tests compare cropped counts with full-table counts for rectangles wholly inside, outside, +crossing or clipped against the hole's bounds. They include separated holes, canvas edges, +very large query coordinates and checked table allocation/bounds failures. + +Search tests compare exact source displacements with the previous implementation over +1, 3, 4 and 5 channels; 8-bit, 16-bit and float working values; several ring/search radii; +empty, one-pixel and separated holes; uniform ties and nonfinite samples. A completely +unknown image still has no source candidate. + +```sh +cargo test --release -p photocraft-algo proximity_integral_release_comparison -- --ignored --nocapture +``` + +The old search is retained only in the test module. The selected source displacement is +unchanged, so downstream source-patch copying and gradient-domain blending get the same input. diff --git a/docs/roadmap.md b/docs/roadmap.md index dbf31120..2a430743 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -125,6 +125,12 @@ a conservative 512 MiB working-set estimate, excluding stored Surface tiles and memory. The crossover is a heuristic; images dominated by near-cutoff alpha can need extra scalar work. These measurements do not establish Photoshop filter parity. +2026-10-10: Proximity Match search medians on an AWS c7i.4xlarge improved 1.54–3.82x +for four synthetic stroke regions (50–300 px holes), with identical source displacements. +Bounding the mask table to the hole and pruning losing SSD candidates reduced table storage +by about 90–93 percent there. Full 24–36 MP kernel inputs improved 4.60–6.24x, but normal +engine strokes already use a cropped region. See [method and limits](proximity-match-performance.md). + 2026-10-08: the tools update above is checked against the current toolbar groups in [`panels.rs`](../crates/ui-egui/src/panels.rs), the Pencil and tool-cycle tests in [`pencil_tests.rs`](../crates/ui-egui/src/pencil_tests.rs), Patch's live-preview test in