Selection distance transforms: 7.6-10.6x faster and correct past 4K (#2093)

Co-authored-by: Brandon Thomas <[email protected]>
This commit is contained in:
AyanandBrandon Thomas authored and GitHub committed 2026-10-09 19:28:33 -04:00
1 parent 590c1496b4
commit 3953d78378
6 files changed
+469 -60

No files matched your search

+3
View File
@@ -63,3 +63,6 @@ sync_github_name = true
[people.skraus2]
sync_github_name = true
[people.excited-raccoon]
sync_github_name = true
+2 -60
View File
@@ -706,67 +706,9 @@ pub fn tone_range(px: &[[f32; 4]], lo: f32, hi: f32, falloff: f32) -> Vec<f32> {
.collect()
}
/// 1D squared distance transform (Felzenszwalb & Huttenlocher).
fn dt1(f: &[f32], out: &mut [f32], v: &mut [usize], z: &mut [f32]) {
let n = f.len();
let mut k = 0usize;
v[0] = 0;
z[0] = f32::NEG_INFINITY;
z[1] = f32::INFINITY;
for q in 1..n {
loop {
let p = v[k];
let s = ((f[q] + (q * q) as f32) - (f[p] + (p * p) as f32)) / (2.0 * (q as f32 - p as f32));
if s <= z[k] && k > 0 {
k -= 1;
continue;
}
if s <= z[k] {
v[0] = q;
z[0] = f32::NEG_INFINITY;
z[1] = f32::INFINITY;
break;
}
k += 1;
v[k] = q;
z[k] = s;
z[k + 1] = f32::INFINITY;
break;
}
}
k = 0;
for (q, o) in out.iter_mut().enumerate() {
while z[k + 1] < q as f32 {
k += 1;
}
let d = q as f32 - v[k] as f32;
*o = d * d + f[v[k]];
}
}
mod distance;
/// Euclidean distance to the nearest `true` pixel.
pub fn edt(inside: &[bool], w: usize, h: usize) -> Vec<f32> {
let mut g: Vec<f32> = inside.iter().map(|&b| if b { 0.0 } else { 1e20 }).collect();
let n = w.max(h).max(1);
let (mut f, mut o, mut v, mut z) = (vec![0.0; n], vec![0.0; n], vec![0usize; n], vec![0.0f32; n + 1]);
for x in 0..w {
for y in 0..h {
f[y] = g[y * w + x];
}
dt1(&f[..h], &mut o[..h], &mut v, &mut z);
for y in 0..h {
g[y * w + x] = o[y];
}
}
for y in 0..h {
f[..w].copy_from_slice(&g[y * w..(y + 1) * w]);
dt1(&f[..w], &mut o[..w], &mut v, &mut z);
for x in 0..w {
g[y * w + x] = o[x].sqrt();
}
}
g
}
pub use distance::edt;
/// Grows the selection by `r` pixels.
pub fn expand(m: &[f32], w: usize, h: usize, r: f32) -> Vec<f32> {
+132
View File
@@ -0,0 +1,132 @@
//! Binary Euclidean distance transform: nearest seeds along columns, then a lower envelope
//! of squared-distance parabolas along rows (Felzenszwalb and Huttenlocher, 2012).
const NO_SEED: u32 = u32::MAX;
// Keep the public API's finite result for a mask with no selected pixel.
const EMPTY_DISTANCE: f32 = 1e10;
const BAND_ROWS: usize = 32;
/// Euclidean distance to the nearest `true` pixel, in row-major order.
///
/// Empty or inconsistent dimensions return an empty buffer. A mask with no seeds returns
/// `1e10` at every pixel, as before. Native passes run in parallel; wasm uses the same scalar
/// kernel. Working storage is one u32 per pixel in addition to the returned f32 buffer, plus
/// a 32-row band and a row envelope per task.
pub fn edt(inside: &[bool], w: usize, h: usize) -> Vec<f32> {
let Some(n) = w.checked_mul(h) else { return Vec::new() };
if n == 0 || n > isize::MAX as usize / std::mem::size_of::<u32>() || n != inside.len() || u32::try_from(w).is_err() || u32::try_from(h).is_err() {
return Vec::new();
}
if !inside.iter().any(|seed| *seed) {
return vec![EMPTY_DISTANCE; n];
}
if inside.iter().all(|seed| *seed) {
return vec![0.0; n];
}
let Some(band_len) = w.checked_mul(BAND_ROWS) else { return Vec::new() };
// Column-major storage makes both distance sweeps contiguous. Only the boolean input
// is read with a stride. Binary costs need two nearest-seed sweeps, not an envelope.
let mut vertical = vec![NO_SEED; n];
for_chunks(&mut vertical, h, |x, column| {
let mut nearest = None;
for (y, d) in column.iter_mut().enumerate() {
if inside.get(y * w + x).copied().unwrap_or(false) {
nearest = Some(y);
}
if let Some(seed) = nearest {
*d = (y - seed) as u32;
}
}
nearest = None;
for (y, d) in column.iter_mut().enumerate().rev() {
if *d == 0 {
nearest = Some(y);
}
if let Some(seed) = nearest {
*d = (*d).min((seed - y) as u32);
}
}
});
let mut out = vec![0.0; n];
for_chunks(&mut out, band_len, |band, dst| {
let y0 = band * BAND_ROWS;
let rows = dst.len() / w;
let mut costs = vec![NO_SEED; dst.len()];
// Copy small column spans into a row band. Reading whole rows straight from the
// transposed plane would revisit one cache line per pixel on a large canvas.
for (x, column) in vertical.chunks_exact(h).enumerate() {
if let Some(span) = column.get(y0..y0 + rows) {
for (y, d) in span.iter().enumerate() {
if let Some(cost) = costs.get_mut(y * w + x) {
*cost = *d;
}
}
}
}
let mut sites = vec![0usize; w];
let mut starts = vec![0.0f64; w];
for (row, result) in costs.chunks_exact(w).zip(dst.chunks_exact_mut(w)) {
row_distance(row, result, &mut sites, &mut starts);
}
});
out
}
fn for_chunks<T: Send + Sync>(buf: &mut [T], stride: usize, f: impl Fn(usize, &mut [T]) + Sync + Send) {
// Avoid starting Rayon tasks for the tiny masks used by tools and thumbnails.
if buf.len() < 65_536 {
for (i, chunk) in buf.chunks_mut(stride).enumerate() {
f(i, chunk);
}
} else {
crate::photo_util::par_rows(buf, stride, 1, f);
}
}
/// Lower envelope of parabolas with integer vertical distances. Work is O(row length): each
/// site enters and leaves the envelope at most once. Empty columns never become candidates.
fn row_distance(costs: &[u32], out: &mut [f32], sites: &mut [usize], starts: &mut [f64]) {
let mut count = 0usize;
for (q, &distance) in costs.iter().enumerate().filter(|(_, d)| **d != NO_SEED) {
let fq = f64::from(distance).powi(2);
let mut start = f64::NEG_INFINITY;
while count > 0 {
let Some(&p) = sites.get(count - 1) else { return };
let Some(&dp) = costs.get(p) else { return };
// (q² - p²)/(2(q-p)) = (q+p)/2. This avoids subtracting large rounded
// coordinate squares. f64 boundaries retain the half-pixel precision on PSBs.
start = (fq - f64::from(dp).powi(2)) / (2.0 * (q - p) as f64) + (q as f64 + p as f64) * 0.5;
if starts.get(count - 1).is_some_and(|s| start > *s) {
break;
}
count -= 1;
}
if count == 0 {
start = f64::NEG_INFINITY;
}
if let (Some(site), Some(boundary)) = (sites.get_mut(count), starts.get_mut(count)) {
*site = q;
*boundary = start;
}
count += 1;
}
if count == 0 {
out.fill(EMPTY_DISTANCE);
return;
}
let mut k = 0usize;
for (x, d) in out.iter_mut().enumerate() {
while k + 1 < count && starts.get(k + 1).is_some_and(|s| *s < x as f64) {
k += 1;
}
if let Some(&site) = sites.get(k)
&& let Some(&vertical) = costs.get(site)
{
let dx = x as f64 - site as f64;
*d = (dx * dx + f64::from(vertical).powi(2)).sqrt() as f32;
}
}
}
#[cfg(test)]
mod tests;
+251
View File
@@ -0,0 +1,251 @@
use super::*;
fn brute(inside: &[bool], w: usize, h: usize) -> Vec<f32> {
let sites: Vec<_> = inside.iter().enumerate().filter(|(_, b)| **b).map(|(i, _)| (i % w, i / w)).collect();
(0..w * h)
.map(|i| {
sites
.iter()
.map(|&(x, y)| {
let dx = (i % w) as f64 - x as f64;
let dy = (i / w) as f64 - y as f64;
(dx * dx + dy * dy).sqrt() as f32
})
.fold(EMPTY_DISTANCE, f32::min)
})
.collect()
}
#[test]
fn matches_brute_force_for_every_small_mask() {
for (w, h) in [(1, 1), (1, 8), (8, 1), (3, 3), (4, 3)] {
for bits in 0u64..1 << (w * h) {
let inside: Vec<_> = (0..w * h).map(|i| bits & (1 << i) != 0).collect();
assert_eq!(edt(&inside, w, h), brute(&inside, w, h), "{w}x{h}: {bits}");
}
}
}
#[test]
fn rectangular_and_partial_bands_match_the_previous_transform() {
for (w, h) in [(1, 101), (101, 1), (17, 33), (73, 65), (9, 255), (513, 137)] {
for kind in 0..4 {
let inside: Vec<_> = (0..w * h)
.map(|i| match kind {
0 => false,
1 => true,
2 => i == 0 || i == w * h - 1,
_ => (i * 101 + i / w * 37) % 29 == 0,
})
.collect();
assert_eq!(edt(&inside, w, h), old_edt(&inside, w, h));
}
}
}
#[test]
fn large_coordinates_keep_nearby_parabolas_distinct() {
// Adjacent sites late in a PSB-sized row lost their half-pixel boundaries when
// the previous transform subtracted f32 coordinate squares.
for (w, h) in [(65_537, 2), (2, 65_537)] {
let mut inside = vec![false; w * h];
if w > h {
inside[w - 3] = true;
inside[w - 2] = true;
inside[2 * w - 1] = true;
} else {
inside[(h - 3) * w] = true;
inside[(h - 2) * w] = true;
inside[h * w - 1] = true;
}
assert_eq!(edt(&inside, w, h), brute(&inside, w, h));
}
}
#[test]
fn selected_pixels_stay_at_zero_past_coordinate_4096() {
let inside = vec![true; 6000];
assert_eq!(edt(&inside, 6000, 1), vec![0.0; 6000]);
assert!(old_edt(&inside, 6000, 1).iter().any(|d| *d != 0.0), "the previous transform reproduces the regression");
}
#[test]
fn inconsistent_and_empty_dimensions_do_not_panic() {
for (mask, w, h) in
[(&[][..], 0, 0), (&[][..], 1, 0), (&[][..], 0, 1), (&[][..], 1, 1), (&[true][..], 2, 1), (&[true, false][..], 1, 1), (&[][..], usize::MAX, 2)]
{
assert!(edt(mask, w, h).is_empty());
}
}
#[test]
fn selection_modifiers_keep_fractional_coverage() {
use crate::selection::{border, contract, expand};
let (w, h) = (47, 39);
let m: Vec<f32> = (0..w * h).map(|i| if (8..31).contains(&(i % w)) && (5..27).contains(&(i / w)) { 0.8 } else { 0.2 }).collect();
let reference_expand = |m: &[f32], r: f32| {
let inside: Vec<_> = m.iter().map(|v| *v >= 0.5).collect();
m.iter().zip(old_edt(&inside, w, h)).map(|(v, d)| v.max((r + 1.0 - d).clamp(0.0, 1.0))).collect::<Vec<_>>()
};
for r in [0.0, 1.0, 1.25, 3.75, 16.0] {
assert_eq!(expand(&m, w, h, r), reference_expand(&m, r));
let inverse: Vec<_> = m.iter().map(|v| 1.0 - v).collect();
let contracted: Vec<_> = reference_expand(&inverse, r).iter().map(|v| 1.0 - v).collect();
assert_eq!(contract(&m, w, h, r), contracted);
let outer = reference_expand(&m, r / 2.0);
let inner: Vec<_> = reference_expand(&inverse, r / 2.0).iter().map(|v| 1.0 - v).collect();
assert_eq!(border(&m, w, h, r), outer.iter().zip(inner).map(|(o, i)| (o - i).max(0.0)).collect::<Vec<_>>());
}
}
#[cfg(not(target_arch = "wasm32"))]
#[test]
fn results_do_not_depend_on_the_worker_count() {
let (w, h) = (521, 139);
let inside: Vec<_> = (0..w * h).map(|i| (i * 101 + i / w * 37) % 97 == 0).collect();
let expected = old_edt(&inside, w, h);
for workers in [1, 2, 4] {
let pool = rayon::ThreadPoolBuilder::new().num_threads(workers).build().unwrap();
assert_eq!(pool.install(|| edt(&inside, w, h)), expected, "{workers} workers");
}
}
#[cfg(not(target_arch = "wasm32"))]
#[test]
#[ignore = "release comparison on 24 and 36 MP masks"]
#[allow(clippy::assertions_on_constants)] // This opt-in benchmark needs an optimized build.
fn distance_transform_release_comparison() {
use std::{hint::black_box, time::Instant};
assert!(!cfg!(debug_assertions), "run with --release");
let workers = std::env::var("EDT_WORKERS").ok().and_then(|s| s.parse().ok()).unwrap_or(16);
let pool = rayon::ThreadPoolBuilder::new().num_threads(workers).build().unwrap();
for (w, h) in [(6000usize, 4000usize), (7200, 5000)] {
for kind in ["ellipse", "sparse", "dense"] {
let inside: Vec<_> = (0..w * h)
.map(|i| {
let (x, y) = (i % w, i / w);
match kind {
"ellipse" => {
let dx = x as f64 - w as f64 * 0.5;
let dy = y as f64 - h as f64 * 0.5;
(dx / (w as f64 * 0.35)).powi(2) + (dy / (h as f64 * 0.3)).powi(2) < 1.0
}
"sparse" => x % 251 == 17 && y % 241 == 23,
_ => (i.wrapping_mul(2_654_435_761) ^ (i >> 8)) % 5 < 2,
}
})
.collect();
let before = old_edt(&inside, w, h);
let after = pool.install(|| edt(&inside, w, h));
assert!(inside.iter().zip(&after).all(|(seed, distance)| !seed || *distance == 0.0));
let mut different = 0usize;
let mut max_error = 0.0f32;
for (a, b) in before.iter().zip(&after) {
if a != b {
different += 1;
max_error = max_error.max((a - b).abs());
}
}
// Validate the large transform independently at corners, seed sites and a grid.
let sites: Vec<_> = inside.iter().enumerate().filter(|(_, b)| **b).map(|(i, _)| (i % w, i / w)).collect();
for i in [0, w - 1, w * (h - 1), w * h - 1, w * h / 2 + w / 2] {
let expected = sites
.iter()
.map(|&(x, y)| {
let dx = (i % w) as f64 - x as f64;
let dy = (i / w) as f64 - y as f64;
(dx * dx + dy * dy).sqrt() as f32
})
.fold(EMPTY_DISTANCE, f32::min);
assert_eq!(after[i], expected);
}
drop(sites);
drop(before);
drop(after);
let mut old_ms = Vec::new();
let mut new_ms = Vec::new();
for run in 0..3 {
let mut measure = |fast: bool| {
let t = Instant::now();
let result = if fast { pool.install(|| edt(black_box(&inside), w, h)) } else { old_edt(black_box(&inside), w, h) };
let ms = t.elapsed().as_secs_f64() * 1000.0;
black_box(&result);
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!(
"EDT {w}x{h} {kind} workers={workers} old_ms={old_ms:?} new_ms={new_ms:?} speedup={:.2} different={different} max_error={max_error}",
old_ms[1] / new_ms[1]
);
}
}
}
/// 1D squared distance transform (Felzenszwalb & Huttenlocher).
fn old_dt1(f: &[f32], out: &mut [f32], v: &mut [usize], z: &mut [f32]) {
let n = f.len();
let mut k = 0usize;
v[0] = 0;
z[0] = f32::NEG_INFINITY;
z[1] = f32::INFINITY;
for q in 1..n {
loop {
let p = v[k];
let s = ((f[q] + (q * q) as f32) - (f[p] + (p * p) as f32)) / (2.0 * (q as f32 - p as f32));
if s <= z[k] && k > 0 {
k -= 1;
continue;
}
if s <= z[k] {
v[0] = q;
z[0] = f32::NEG_INFINITY;
z[1] = f32::INFINITY;
break;
}
k += 1;
v[k] = q;
z[k] = s;
z[k + 1] = f32::INFINITY;
break;
}
}
k = 0;
for (q, o) in out.iter_mut().enumerate() {
while z[k + 1] < q as f32 {
k += 1;
}
let d = q as f32 - v[k] as f32;
*o = d * d + f[v[k]];
}
}
/// Euclidean distance to the nearest `true` pixel.
fn old_edt(inside: &[bool], w: usize, h: usize) -> Vec<f32> {
let mut g: Vec<f32> = inside.iter().map(|&b| if b { 0.0 } else { 1e20 }).collect();
let n = w.max(h).max(1);
let (mut f, mut o, mut v, mut z) = (vec![0.0; n], vec![0.0; n], vec![0usize; n], vec![0.0f32; n + 1]);
for x in 0..w {
for y in 0..h {
f[y] = g[y * w + x];
}
old_dt1(&f[..h], &mut o[..h], &mut v, &mut z);
for y in 0..h {
g[y * w + x] = o[y];
}
}
for y in 0..h {
f[..w].copy_from_slice(&g[y * w..(y + 1) * w]);
old_dt1(&f[..w], &mut o[..w], &mut v, &mut z);
for x in 0..w {
g[y * w + x] = o[x].sqrt();
}
}
g
}
+6
View File
@@ -71,6 +71,12 @@ i7-9750H, 12 workers, 24 MP RGBA8, Good quality, amount 1: Spin 66.10 → 16.43
(4.02×), Zoom 19.19 → 5.35 s (3.59×), one paired run each, exact output equality.
These are local measurements at amount 1; see [method and limits](radial-blur-performance.md).
2026-10-10: native selection distance transforms measured 7.59–10.57x faster on six
24–36 MP synthetic masks on an AWS c7i.4xlarge (16 workers, three paired release runs).
The transform also corrects f32 envelope errors past coordinate 4096, including nonzero
distances at selected pixels. The parallel path needs an additional four bytes per pixel;
see [measurements, correctness and limits](selection-distance-performance.md).
**Bottom line.** Two days after 0.2.0 we had merged ~96 PRs and closed ~48 issues, but **real
Photoshop parity is still well below 50%**. The biggest gaps are AI, missing tools, professional
workflow depth and the plug-in ecosystem. Most fixes since 0.2.0 have passed our tests but have
+75
View File
@@ -0,0 +1,75 @@
# Selection distance transforms
The distance transform used by Select > Modify > Expand, Contract and Border now runs
as two native parallel passes. It also serves matting and mask feathering callers in the
engine. The compositor has a separate implementation and is outside this change.
The first pass finds the nearest selected pixel above and below each pixel. Since its
input is binary, two sweeps replace the general parabola envelope for that axis. Vertical
distances are stored as integers in column order. The second pass copies 32-row bands
into contiguous rows and evaluates the lower envelope of squared-distance parabolas.
Both passes are linear in the number of pixels. WebAssembly uses the same scalar kernels.
The row envelope follows [Felzenszwalb and Huttenlocher, 2012](https://toc.cs.uchicago.edu/articles/v008a019/).
Its intersections use `(fq - fp) / (2 * (q - p)) + (q + p) / 2` in f64. This avoids the old
subtraction of f32 coordinate squares, which lost half-pixel boundary precision on large
canvases. Distances are rounded to f32 only at the output.
## Correctness
The previous transform can return a positive distance at a selected pixel past coordinate
4096. A 6000 x 1 fully selected mask reproduces it. The regression tests also compare sparse,
competing sites in 65537 x 2 and 2 x 65537 masks against direct Euclidean distances.
Tests enumerate every mask up to 4 x 3 against a brute-force oracle, check rectangular and
partial row bands against the previous implementation, preserve fractional Expand/Contract/
Border coverage on smaller inputs, and compare native worker counts 1, 2 and 4. Empty or
inconsistent dimensions return an empty buffer. No-seed masks retain the finite 1e10 result.
Empty and fully selected masks skip the transform passes.
Large benchmark inputs check every selected pixel has zero distance and check independent
nearest-site distances at corners and the centre. The sparse-grid cases match the old
output exactly. Ellipse and dense cases intentionally correct old distance errors, so this
is not a claim of identical output everywhere or of measured Photoshop parity.
## Measurements
Recorded 2026-10-10 on an AWS c7i.4xlarge in ap-southeast-2: 16 vCPUs, 32 GiB RAM,
Rust 1.99.0, the workspace's normal release profile, 16 Rayon workers. Synthetic masks are
an ellipse, a sparse regular grid, and a deterministic dense pattern. Input generation,
reference validation and destruction of the output happen outside the timed interval.
Each implementation gets a warmup, then three measured runs with alternating order.
| Size | Mask | Before p50 | After p50 | Speedup | Changed distances | Largest correction |
|---|---|---:|---:|---:|---:|---:|
| 24 MP | ellipse | 630.74 ms | 83.11 ms | 7.59x | 827,500 | 1 px |
| 24 MP | sparse | 621.81 ms | 71.03 ms | 8.75x | 0 | 0 px |
| 24 MP | dense | 1048.21 ms | 99.73 ms | 10.51x | 2,034,343 | 1.4142135 px |
| 36 MP | ellipse | 994.14 ms | 124.48 ms | 7.99x | 2,431,817 | 1 px |
| 36 MP | sparse | 984.52 ms | 109.94 ms | 8.96x | 0 | 0 px |
| 36 MP | dense | 1588.23 ms | 150.28 ms | 10.57x | 4,771,985 | 1.4142135 px |
These are kernel timings. Mask conversion, selection history, document compositing and
canvas upload are excluded. They do not establish an end-to-end interaction budget.
## Memory and limits
The old transform used one f32 plane. The new transform needs one u32 plane plus the f32
result, or eight bytes per pixel, excluding the caller's input mask. It also uses up to
`32 * width * 4 + width * (size_of::<usize>() + 8)` scratch bytes per active row task.
At 24 MP the two planes total 192 MB; at 36 MP they total 288 MB. Fully empty and fully
selected masks need only the result plane. Small mixed masks below 65536 pixels run serially.
The extra plane buys contiguous writes and independent parallel bands. This remains a dense
transform, so it does not solve streaming selections on very large PSB documents. Dimension
multiplication, element-size limits and input length are checked before allocation.
## Reproduce
```sh
cargo test --release -p photocraft-algo distance_transform_release_comparison -- --ignored --nocapture
EDT_WORKERS=1 cargo test --release -p photocraft-algo distance_transform_release_comparison -- --ignored --nocapture
```
The previous scalar implementation is retained only in the test module as the reference
and benchmark baseline.