perf(analysis): viewshed rows on several cores

crates/analysis/src/parallel.rs

@@ -12,8 +12,9 @@ pub(crate) const BATCH: usize = 16;
 /// Each worker's stack, as large as the main thread's usually is.
 const STACK_BYTES: usize = 8 << 20;
 
-/// The work below which a raster tool's rows are made on one thread: cell-point pairs, say,
-/// about 30 ns each for IDW, so some 5 ms of it, against tens of microseconds to start a thread.
+/// The work below which a raster tool's rows are made on one thread, in the tool's own units:
+/// cell-point pairs for IDW (about 30 ns each, so some 5 ms of it), the same plus each cell's
+/// system for kriging, ray steps for a viewshed. Starting a thread costs tens of microseconds.
 pub(crate) const MIN_ROW_WORK: usize = 200_000;
 /// The scratch memory a raster tool's threads may take together, beyond one thread's.
 const SCRATCH_BUDGET: usize = 256 << 20;
@@ -46,12 +47,14 @@ fn cores() -> usize {
 
 /// `f` on each of a raster's `rows` (its output rows, `chunks_mut` of the bands it writes) with
 /// the row's number, on `threads` threads ([`threads_for_rows`]), this one among them. Each row
-/// is written where it lies, so the output is held once, as on one thread. A row is costly
-/// enough to be taken one at a time, so rows that cost more than others are shared out too.
+/// is written where it lies, so the output is held once, as on one thread. Rows are taken a few
+/// at a time, by the rows still left: one each when there are few per thread (and so at the
+/// end), so rows that cost more than others are shared out, up to [`BATCH`] when there are many,
+/// so cheap rows don't wait on the lock.
 /// Each thread makes its own `scratch()` once; one that panics stops the others after the rows
 /// they are making, and its panic is raised again here.
 pub(crate) fn for_each_row<T: Send, S>(
-    rows: impl Iterator<Item = T> + Send,
+    rows: impl ExactSizeIterator<Item = T> + Send,
     threads: usize,
     scratch: impl Fn() -> S + Sync,
     f: impl Fn(&mut S, usize, T) + Sync,
@@ -67,10 +70,19 @@ pub(crate) fn for_each_row<T: Send, S>(
     let work = || {
         let _stop = StopOnPanic(&failed);
         let mut s = scratch();
+        let mut taken = Vec::with_capacity(BATCH);
         while !failed.load(Ordering::Relaxed) {
-            // The lock is held only to take a row; a poisoned one means a thread panicked.
-            let Some((i, row)) = rows.lock().ok().and_then(|mut r| r.next()) else { break };
-            f(&mut s, i, row);
+            // The lock is held only to take rows; a poisoned one means a thread panicked.
+            if let Ok(mut r) = rows.lock() {
+                let take = (r.len() / threads.saturating_mul(BATCH)).clamp(1, BATCH);
+                taken.extend(r.by_ref().take(take));
+            }
+            if taken.is_empty() {
+                break;
+            }
+            for (i, row) in taken.drain(..) {
+                f(&mut s, i, row);
+            }
         }
     };
     std::thread::scope(|scope| {
@@ -109,33 +121,11 @@ pub(crate) fn try_map_on<T: Sync, S, R: Send, E: Send>(
     threads: usize,
     scratch: impl Fn() -> S + Sync,
     f: impl Fn(&mut S, usize, &T) -> Result<R, E> + Sync,
-) -> Result<Vec<R>, E> {
-    try_each_on(items.len(), threads, scratch, |s, i| f(s, i, &items[i]))
-}
-
-/// Sets the stop flag when the thread it lives on panics, so the others stop after the items
-/// they are making rather than make the rest for nothing.
-struct StopOnPanic<'a>(&'a std::sync::atomic::AtomicBool);
-
-impl Drop for StopOnPanic<'_> {
-    fn drop(&mut self) {
-        if std::thread::panicking() {
-            self.0.store(true, std::sync::atomic::Ordering::Relaxed);
-        }
-    }
-}
-
-/// [`try_map_on`] over the indices `0..len`.
-fn try_each_on<S, R: Send, E: Send>(
-    len: usize,
-    threads: usize,
-    scratch: impl Fn() -> S + Sync,
-    f: impl Fn(&mut S, usize) -> Result<R, E> + Sync,
 ) -> Result<Vec<R>, E> {
     use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
-    if threads <= 1 || len <= 1 {
+    if threads <= 1 || items.len() <= 1 {
         let mut s = scratch();
-        return (0..len).map(|i| f(&mut s, i)).collect();
+        return items.iter().enumerate().map(|(i, item)| f(&mut s, i, item)).collect();
     }
     let next = AtomicUsize::new(0);
     let failed = AtomicBool::new(false);
@@ -145,21 +135,20 @@ fn try_each_on<S, R: Send, E: Send>(
         let mut made = Vec::new();
         while !failed.load(Ordering::Relaxed) {
             let from = next.fetch_add(BATCH, Ordering::Relaxed);
-            if from >= len {
-                break;
-            }
-            for i in from..len.min(from.saturating_add(BATCH)) {
-                match f(&mut s, i) {
-                    Ok(r) => made.push((i, r)),
+            let Some(batch) = items.get(from..items.len().min(from.saturating_add(BATCH))).filter(|b| !b.is_empty()) else { break };
+            for (k, item) in batch.iter().enumerate() {
+                match f(&mut s, from + k, item) {
+                    Ok(r) => made.push((from + k, r)),
                     Err(e) => {
                         failed.store(true, Ordering::Relaxed);
-                        return Err((i, e));
+                        return Err((from + k, e));
                     }
                 }
             }
         }
         Ok(made)
     };
+
     let mut parts = Vec::with_capacity(threads);
     std::thread::scope(|scope| {
         let others: Vec<_> =
@@ -171,7 +160,7 @@ fn try_each_on<S, R: Send, E: Send>(
             parts.push(h.join().unwrap_or_else(|e| std::panic::resume_unwind(e)));
         }
     });
-    let mut made: Vec<(usize, R)> = Vec::with_capacity(len);
+    let mut made: Vec<(usize, R)> = Vec::with_capacity(items.len());
     let mut first: Option<(usize, E)> = None;
     for part in parts {
         match part {
@@ -190,6 +179,18 @@ fn try_each_on<S, R: Send, E: Send>(
     Ok(made.into_iter().map(|(_, r)| r).collect())
 }
 
+/// Sets the stop flag when the thread it lives on panics, so the others stop after the items
+/// they are making rather than make the rest for nothing.
+struct StopOnPanic<'a>(&'a std::sync::atomic::AtomicBool);
+
+impl Drop for StopOnPanic<'_> {
+    fn drop(&mut self) {
+        if std::thread::panicking() {
+            self.0.store(true, std::sync::atomic::Ordering::Relaxed);
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -321,6 +322,13 @@ mod tests {
             );
             assert_eq!(out, alone, "{threads} threads");
         }
+        // Many cheap rows per thread, taken several at a time.
+        let tall: Vec<usize> = (0..5_000).collect();
+        for threads in [2, 16] {
+            let mut out = vec![0usize; 5_000];
+            for_each_row(out.chunks_mut(1), threads, || (), |_, r, row: &mut [usize]| row[0] = r);
+            assert_eq!(out, tall, "{threads} threads");
+        }
         let mut none: Vec<u8> = Vec::new();
         for_each_row(none.chunks_mut(5), 4, || (), |_, _, _: &mut [u8]| panic!("no rows"));
         // Small work, and the browser: one thread; never more threads than rows.

crates/analysis/src/rasterops.rs

@@ -314,6 +314,11 @@ pub(crate) const MAX_IDW_WORK: usize = 400_000_000;
 
 /// Cells visible from an observer (1) or not (0) within `radius_m`; outside the radius no-data.
 pub fn viewshed(r: &Raster, b: usize, observer: Coord, height: f64, radius_m: f64) -> Result<Vec<f32>, String> {
+    viewshed_on(None, r, b, observer, height, radius_m)
+}
+
+/// [`viewshed`] on `threads` threads, or as many as its work calls for.
+fn viewshed_on(threads: Option<usize>, r: &Raster, b: usize, observer: Coord, height: f64, radius_m: f64) -> Result<Vec<f32>, String> {
     let (oc, or) = r.pixel_at(observer).ok_or("the observer is outside the raster")?;
     let z0 = r.value(b, oc, or).ok_or("the observer stands on no-data")? as f64 + height;
     let (cx, cy) = r.cell_size_m();
@@ -328,46 +333,58 @@ pub fn viewshed(r: &Raster, b: usize, observer: Coord, height: f64, radius_m: f6
     let rad_r = (radius_m / cy).ceil() as i64;
     // Each cell walks a ray back to the observer: bound cells × ray length, not just cells.
     let cells = (2 * rad_c + 1).saturating_mul(2 * rad_r + 1);
-    if cells > 25_000_000 || cells.saturating_mul(rad_c.max(rad_r) / 2 + 1) > MAX_VIEWSHED_WORK {
+    let ray_work = cells.saturating_mul(rad_c.max(rad_r) / 2 + 1);
+    if cells > 25_000_000 || ray_work > MAX_VIEWSHED_WORK {
         return Err("the radius covers too many cells; use a smaller radius".into());
     }
     let w = r.width as usize;
     let mut out = vec![NODATA; w * r.height as usize];
     const R: f64 = 6_371_000.0 * 4.0 / 3.0;
-    for row in (or - rad_r).max(0)..=(or + rad_r).min(r.height as i64 - 1) {
-        for col in (oc - rad_c).max(0)..=(oc + rad_c).min(r.width as i64 - 1) {
-            let (dx, dy) = ((col - oc) as f64 * cx, (row - or) as f64 * cy);
-            let dist = dx.hypot(dy);
-            if dist > radius_m {
-                continue;
-            }
-            let idx = row as usize * w + col as usize;
-            let Some(zt) = r.value(b, col, row) else { continue };
-            if dist == 0.0 {
-                if let Some(o) = out.get_mut(idx) {
-                    *o = 1.0;
+    // Each cell walks its own ray and writes only itself, so the rows within the radius are made
+    // on several threads when the work is large (#1252).
+    let (first, last) = ((or - rad_r).max(0), (or + rad_r).min(r.height as i64 - 1));
+    let rows = usize::try_from(last - first + 1).unwrap_or(0);
+    let threads = threads.unwrap_or_else(|| crate::parallel::threads_for_rows(rows, usize::try_from(ray_work).unwrap_or(usize::MAX), 0));
+    let window = out.chunks_mut(w).skip(usize::try_from(first).unwrap_or(0)).take(rows);
+    crate::parallel::for_each_row(
+        window,
+        threads,
+        || (),
+        |_, i, line: &mut [f32]| {
+            let row = first + i as i64;
+            for col in (oc - rad_c).max(0)..=(oc + rad_c).min(r.width as i64 - 1) {
+                let (dx, dy) = ((col - oc) as f64 * cx, (row - or) as f64 * cy);
+                let dist = dx.hypot(dy);
+                if dist > radius_m {
+                    continue;
                 }
-                continue;
-            }
-            let target = (zt as f64 - dist * dist / (2.0 * R) - z0) / dist;
-            // Walk the ray from the observer; any cell higher than the line hides the target.
-            let steps = ((col - oc).abs().max((row - or).abs())) as usize;
-            let mut visible = true;
-            for s in 1..steps {
-                let t = s as f64 / steps as f64;
-                let (c, rr) = ((oc as f64 + (col - oc) as f64 * t).round() as i64, (or as f64 + (row - or) as f64 * t).round() as i64);
-                let Some(zi) = r.value(b, c, rr) else { continue };
-                let di = dist * t;
-                if (zi as f64 - di * di / (2.0 * R) - z0) / di > target {
-                    visible = false;
-                    break;
+                let Some(zt) = r.value(b, col, row) else { continue };
+                if dist == 0.0 {
+                    if let Some(o) = line.get_mut(col as usize) {
+                        *o = 1.0;
+                    }
+                    continue;
+                }
+                let target = (zt as f64 - dist * dist / (2.0 * R) - z0) / dist;
+                // Walk the ray from the observer; any cell higher than the line hides the target.
+                let steps = ((col - oc).abs().max((row - or).abs())) as usize;
+                let mut visible = true;
+                for s in 1..steps {
+                    let t = s as f64 / steps as f64;
+                    let (c, rr) = ((oc as f64 + (col - oc) as f64 * t).round() as i64, (or as f64 + (row - or) as f64 * t).round() as i64);
+                    let Some(zi) = r.value(b, c, rr) else { continue };
+                    let di = dist * t;
+                    if (zi as f64 - di * di / (2.0 * R) - z0) / di > target {
+                        visible = false;
+                        break;
+                    }
+                }
+                if let Some(o) = line.get_mut(col as usize) {
+                    *o = if visible { 1.0 } else { 0.0 };
                 }
             }
-            if let Some(o) = out.get_mut(idx) {
-                *o = if visible { 1.0 } else { 0.0 };
-            }
-        }
-    }
+        },
+    );
     Ok(out)
 }
 
@@ -1106,6 +1123,35 @@ impl Expr {
 mod tests {
     use super::*;
 
+    /// Viewshed rows made on several threads see what one thread sees (#1252), with the radius
+    /// inside the raster and reaching past its edges.
+    #[test]
+    fn viewshed_on_several_threads_is_viewshed_on_one() {
+        let (w, h) = (90u32, 70u32);
+        let data = (0..w * h)
+            .map(|i| {
+                let (x, y) = ((i % w) as f32, (i / w) as f32);
+                if i % 97 == 5 { -9999.0 } else { 300.0 + 40.0 * (x * 0.21).sin() * (y * 0.17).cos() + 25.0 * ((x + y) * 0.4).sin() }
+            })
+            .collect();
+        let transform = GeoTransform { origin_x: 0.0, origin_y: 45.0, pixel_width: 0.0003, pixel_height: -0.0003 };
+        let dem = Raster::new("dem", w, h, vec![data], transform, Crs::LonLat, Some(-9999.0), ColorHint::Gray, "Float32")
+            .unwrap_or_else(|e| panic!("{e}"));
+        // Visible, hidden and no-data cells, as one thread gives them: a change to the
+        // ray walk shows here, not only a difference between thread counts.
+        for (at, radius, expected) in
+            [(Coord::new(0.012, 44.990), 500.0, [207, 775, 5318]), (Coord::new(0.003, 44.985), 5_000.0, [185, 6050, 65])]
+        {
+            let one = viewshed_on(Some(1), &dem, 0, at, 2.0, radius).unwrap_or_else(|e| panic!("{e}"));
+            let counts = [1.0, 0.0, NODATA].map(|k| one.iter().filter(|v| v.to_bits() == k.to_bits()).count());
+            assert_eq!(counts, expected, "radius {radius}");
+            for threads in [2, 7, 16] {
+                let many = viewshed_on(Some(threads), &dem, 0, at, 2.0, radius).unwrap_or_else(|e| panic!("{e}"));
+                assert!(one.iter().map(|v| v.to_bits()).eq(many.iter().map(|v| v.to_bits())), "{threads} threads, radius {radius}");
+            }
+        }
+    }
+
     /// Rows made on several threads give the same surface, bit for bit, as one thread (#1252),
     /// from every point and from the nearest few.
     #[test]